app.py

import os
import re
import webbrowser
import pandas as pd
import gradio as gr
from huggingface_hub import HfApi
from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError
from accelerate.commands.estimate import create_empty_model, check_has_model
from accelerate.utils import convert_bytes, calculate_maximum_sizes

# We need to store them as globals because gradio doesn't have a way for us to pass them in to the button
HAS_DISCUSSION = True
MODEL_NAME = None
LIBRARY = None
USER_TOKEN = None
TOKEN = os.environ.get("HUGGINGFACE_API_LOGIN", None)

def check_for_discussion(model_name:str):
 "Checks if an automated discussion has been opened on the model by `model-sizer-bot`"
 global TOKEN
 api = HfApi(token=TOKEN)
 discussions = list(api.get_repo_discussions(model_name))
 return any(discussion.title == "[AUTOMATED] Model Memory Requirements" and discussion.author == "model-sizer-bot" for discussion in discussions)

def report_results():
 "Reports the results of a memory calculation to the model's discussion page, and opens a new tab to it afterwards"
 global MODEL_NAME, LIBRARY, TOKEN, USER_TOKEN
 api = HfApi(token=TOKEN)
 results, data = calculate_memory(MODEL_NAME, LIBRARY, ["fp32", "fp16", "int8", "int4"], access_token=USER_TOKEN, raw=True)
 minimum = data[0]

 USER_TOKEN = None
 post = f"""# Model Memory Requirements\n
You will need about {minimum[1]} VRAM to load this model for inference, and {minimum[3]} VRAM to train it using Adam.
 
These calculations were measured from the [Model Memory Utility Space](https://hf.co/spaces/hf-accelerate/model-memory-utility) on the Hub.
 
The minimum recommended vRAM needed for this model assumes using [Accelerate or `device_map="auto"`](https://huggingface.co/docs/accelerate/usage_guides/big_modeling) and is denoted by the size of the "largest layer". 
When performing inference, expect to add up to an additional 20% to this, as found by [EleutherAI](https://blog.eleuther.ai/transformer-math/). More tests will be performed in the future to get a more accurate benchmark for each model.
When training with `Adam`, you can expect roughly 4x the reported results to be used. (1x for the model, 1x for the gradients, and 2x for the optimizer).
## Results:
{results}
"""
 discussion = api.create_discussion(MODEL_NAME, "[AUTOMATED] Model Memory Requirements", description=post)
 webbrowser.open_new_tab(discussion.url)

def convert_url_to_name(url:str):
 "Converts a model URL to its name on the Hub"
 results = re.findall(r"huggingface.co\/(.*?)#", url)
 if len(results) < 1:
 raise ValueError(f"URL {url} is not a valid model URL to the Hugging Face Hub")
 return results[0]

# Based on the following doc:
#
# - https://huggingface.co/docs/transformers/v4.31.0/perf_train_gpu_one#anatomy-of-models-memory
# - https://blog.eleuther.ai/transformer-math/
# - https://kipp.ly/transformer-inference-arithmetic/
# - https://github.com/ray-project/llm-numbers
#
def calc_vram_f32(model, optimizer, sequence_len, micro_batch_size, device_count, gradient_checkpointing):
 # is_16bit = cfg.bf16 or cfg.bfloat16 or cfg.load_in_8bit or cfg.fp16 or cfg.float16

 # if torch.cuda.device_count() > 1 or cfg.fsdp or os.environ.get("ACCELERATE_USE_DEEPSPEED") == "true" or cfg.adapter:
 # return { 'supported': False }

 # Model Weights
 #
 # Hf doc counts:
 #
 # - 4 bytes * number of parameters for fp32 training
 # - 6 bytes * number of parameters for mixed precision training (maintains a model in fp32 and one in fp16 in memory)
 #
 # But we follow https://blog.eleuther.ai/transformer-math/#model-weights to count 2 bytes here for mixed precision training,
 # leave the rest to optimizor state.
 #
 # Here we calculate only for fp32, will adjust for each dtype outside.
 #
 # for param in model.parameters():
 # print(f'{type(param)} {param.shape} {param.element_size()}')
 #
 # print(f'total parameters = {sum([param.nelement() for param in model.parameters()])}')

 param_element_size = 4
 vram_model = sum([param.nelement() * param_element_size for param in model.parameters()])

 # Buffers
 #
 # Buffers are tensors that do not require gradients and not registered as parameters.
 # e.g. mean and std in batch norm layers.
 # - https://github.com/huggingface/transformers/blob/d4bd33cc9f11ca48635e54983d75249c78d72e2a/src/transformers/modeling_utils.py#L1897
 # - https://discuss.pytorch.org/t/gpu-memory-that-model-uses/56822/2
 #
 # for buf in model.buffers():
 # print(f'buf.element_size() = {buf.element_size()}')
 vram_buffer = sum([buf.nelement() * buf.element_size() for buf in model.buffers()])

 # Optimizer States:
 # - 8 bytes * number of parameters for normal AdamW (maintains 2 states)
 # - 2 bytes * number of parameters for 8-bit AdamW optimizers like bitsandbytes
 # - 4 bytes * number of parameters for optimizers like SGD with momentum (maintains only 1 state)
 #
 # For now we use AdamW/SGD as the baseline for the estimation, even for other more memory-efficient optimizers
 # ADAMW_HF = "adamw_hf"
 # ADAMW_TORCH = "adamw_torch"
 # ADAMW_TORCH_FUSED = "adamw_torch_fused"
 # ADAMW_TORCH_XLA = "adamw_torch_xla"
 # ADAMW_APEX_FUSED = "adamw_apex_fused"
 # ADAFACTOR = "adafactor"
 # ADAMW_ANYPRECISION = "adamw_anyprecision"
 # SGD = "sgd"
 # ADAGRAD = "adagrad"
 # ADAMW_BNB = "adamw_bnb_8bit"
 # ADAMW_8BIT = "adamw_8bit" # just an alias for adamw_bnb_8bit
 # LION_8BIT = "lion_8bit"
 # LION = "lion_32bit"
 # PAGED_ADAMW = "paged_adamw_32bit"
 # PAGED_ADAMW_8BIT = "paged_adamw_8bit"
 # PAGED_LION = "paged_lion_32bit"
 # PAGED_LION_8BIT = "paged_lion_8bit"
 # optimizer = cfg.optimizer
 optimizer_state_size_per_param = 4 if 'sgd' in optimizer else (2 if '8bit' in optimizer else 8)
 vram_optimizer = sum([param.nelement() * optimizer_state_size_per_param for param in model.parameters()])

 # Gradients
 #
 # 4 bytes * number of parameters for either fp32 or mixed precision training (gradients are always kept in fp32)
 # but we will follow transformer-math to treat it conditionally outside
 # for now we ignores whether is mixed precision training
 #
 gradient_element_size = 4 # 2 if is_16bit else 4
 vram_gradient = sum([param.nelement() * gradient_element_size for param in model.parameters()])

 # Forward Activations
 # size depends on many factors, the key ones being sequence length, hidden size and batch size.
 s = sequence_len # cfg.sequence_len
 b = micro_batch_size # cfg.micro_batch_size
 h = model.config.hidden_size
 L = model.config.num_hidden_layers
 t = device_count # max(1, torch.cuda.device_count()) # len(DataParallel(model).device_ids) #torch.cuda.device_count()
 a = model.config.num_attention_heads
 print(f's={s} b={b} h={h} L={L} t={t} a={a}')

 sbHL = s * b * h * L
 print(f'sbHL = {sbHL / 1e9} GB')

 print(f'10 + {24 / t} + {5 * a * s / (h * t)}')

 vram_activation = sbHL * (10 + 24 / t) if gradient_checkpointing else sbHL * (10 + 24 / t + 5 * a * s / (h * t))

 return {
 # 'supported': True,
 'param_element_size': param_element_size,
 'total': vram_model + vram_buffer + vram_optimizer + vram_activation,
 'model': vram_model,
 'buffer': vram_buffer,
 'optimizer': vram_optimizer,
 'activation': vram_activation,
 }
def bytes_by_dtype(bytes, dtype):
 if dtype in ("fp16", "bf16", "float16/bfloat16"):
 return bytes / 2
 elif dtype == "int8":
 return bytes / 4
 elif dtype == "int4":
 return bytes / 8
 else:
 return bytes

def calculate_memory(model_name:str, library:str, dtypes:list, optimizer:str, access_token:str, raw=False):
 "Calculates the memory usage for a model"
 if library == "auto":
 library = None
 if "http" in model_name and "//" in model_name:
 try:
 model_name = convert_url_to_name(model_name)
 except ValueError:
 raise gr.Error(f"URL `{model_name}` is not a valid model URL to the Hugging Face Hub")
 try:
 model = create_empty_model(model_name, library_name=library, trust_remote_code=True, access_token=access_token)
 except GatedRepoError:
 raise gr.Error(f"Model `{model_name}` is a gated model, please ensure to pass in your access token and try again if you have access. You can find your access token here : https://huggingface.co/settings/tokens. ")
 except RepositoryNotFoundError:
 raise gr.Error(f"Model `{model_name}` was not found on the Hub, please try another model name.")
 except ValueError as e:
 raise gr.Error(f"Model `{model_name}` does not have any library metadata on the Hub, please manually select a library_name to use (such as `transformers`)")
 except (RuntimeError, OSError) as e:
 library = check_has_model(e)
 if library != "unknown":
 raise gr.Error(f"Tried to load `{model_name}` with `{library}` but a possible model to load was not found inside the repo.")
 
 total_size, largest_layer = calculate_maximum_sizes(model)

 data = []

 title = f"Memory Usage for '{model_name}'"

 vram_f32 = calc_vram_f32(model, optimizer=optimizer, sequence_len=2048, micro_batch_size=1, device_count=1, gradient_checkpointing=True)

 for dtype in dtypes:
 param_element_size = bytes_by_dtype(vram_f32['param_element_size'], dtype)
 vram_model = bytes_by_dtype(vram_f32['model'], dtype)
 vram_buffer = vram_f32['buffer']
 vram_optimizer = vram_f32['optimizer']
 vram_activation = vram_f32['activation']
 row = {
 "dtype": dtype,
 'inference_total': convert_bytes(vram_model),
 'training_total': convert_bytes(vram_model + vram_buffer + vram_optimizer + vram_activation),
 'model': convert_bytes(vram_model),
 'buffer': convert_bytes(vram_buffer),
 'optimizer': convert_bytes(vram_optimizer),
 'activation': convert_bytes(vram_activation),
 }

 data.append(row)
 # dtype_total_size = total_size
 # dtype_largest_layer = largest_layer[0]
 # if dtype in ("fp16", "bf16", "float16/bfloat16"):
 # dtype_total_size /= 2
 # dtype_largest_layer /= 2
 # elif dtype == "int8":
 # dtype_total_size /= 4
 # dtype_largest_layer /= 4
 # elif dtype == "int4":
 # dtype_total_size /= 8
 # dtype_largest_layer /= 8
 # dtype_training_size = convert_bytes(dtype_total_size * 4)
 # dtype_total_size = convert_bytes(dtype_total_size)
 # dtype_largest_layer = convert_bytes(dtype_largest_layer)
 # data.append({
 # "dtype": dtype,
 # "Largest Layer or Residual Group": dtype_largest_layer,
 # "Total Size": dtype_total_size,
 # "Training using Adam": dtype_training_size,
 # "Test": 12345
 # })
 # data.append({
 # "dtype": dtype,
 # "Largest Layer or Residual Group": dtype_largest_layer,
 # "Total Size": dtype_total_size,
 # "Training using Adam": dtype_training_size,
 # "Test": 12345
 # })
 global HAS_DISCUSSION, MODEL_NAME, LIBRARY
 HAS_DISCUSSION = check_for_discussion(model_name)
 MODEL_NAME = model_name
 LIBRARY = library

 if raw:
 return pd.DataFrame(data).to_markdown(index=False), data
 
 results = [
 f'## {title}', 
 gr.update(visible=True, value=pd.DataFrame(data)), 
 # gr.update(visible=not HAS_DISCUSSION)
 ]
 return results

with gr.Blocks() as demo:
 with gr.Column():
 gr.Markdown(
 """<img src="https://huggingface.co/spaces/hf-accelerate/model-memory-usage/resolve/main/measure_model_size.png" style="float: left;" width="250" height="250"><h1>🤗 Model Memory Calculator</h1>
 This tool is modified from https://huggingface.co/spaces/hf-accelerate/model-memory-usage with the following changes:
 
 - Focus on transformers and gives more detailed estimation based on more configs
 - Will auto-calculate the proper batch size given a VRAM constraint later
 - LoRA/QLoRA etc. will be supported later

 Note:

 - inference_total = model
 - training_total = model + buffer + optimizer + activation

 """
 )
 out_text = gr.Markdown()
 out = gr.DataFrame(headers=[
 "dtype",
 'inference_total',
 'training_total',
 'model',
 'buffer',
 'optimizer',
 'activation',
 ],
 interactive=False,
 visible=False,
 )
 with gr.Row():
 inp = gr.Textbox(label="Model Name or URL", value="bert-base-cased")
 with gr.Row():
 library = gr.Radio(["transformers"], label="Library", value="transformers")
 dtypes = gr.CheckboxGroup(
 ["float32", "float16/bfloat16", "int8", "int4"],
 value=["float32", "float16/bfloat16", "int8", "int4"],
 label="Model Precision",
 )
 # ADAMW_HF = "adamw_hf"
 # ADAMW_TORCH = "adamw_torch"
 # ADAMW_TORCH_FUSED = "adamw_torch_fused"
 # ADAMW_TORCH_XLA = "adamw_torch_xla"
 # ADAMW_APEX_FUSED = "adamw_apex_fused"
 # ADAFACTOR = "adafactor"
 # ADAMW_ANYPRECISION = "adamw_anyprecision"
 # SGD = "sgd"
 # ADAGRAD = "adagrad"
 # ADAMW_BNB = "adamw_bnb_8bit"
 # ADAMW_8BIT = "adamw_8bit" # just an alias for adamw_bnb_8bit
 # LION_8BIT = "lion_8bit"
 # LION = "lion_32bit"
 # PAGED_ADAMW = "paged_adamw_32bit"
 # PAGED_ADAMW_8BIT = "paged_adamw_8bit"
 # PAGED_LION = "paged_lion_32bit"
 # PAGED_LION_8BIT = "paged_lion_8bit"
 optimizer = gr.Dropdown(choices=["adamw_hf", "adamw_torch", "sgd", "lion_32bit", "adamw_8bit", "lion_8bit", "paged_adamw_8bit", "paged_lion_8bit"],
 value="adamw_hf", label="Optimizer", allow_custom_value=True)
 access_token = gr.Textbox(label="API Token", placeholder="Optional (for gated models)")
 with gr.Row():
 btn = gr.Button("Calculate Memory Usage")
 # post_to_hub = gr.Button(value = "Report results in this model repo's discussions!\n(Will open in a new tab)", visible=False)
 USER_TOKEN = access_token

 btn.click(
 calculate_memory, inputs=[inp, library, dtypes, optimizer, access_token], outputs=[out_text, out],
 )
 
 # post_to_hub.click(report_results).then(lambda: gr.Button.update(visible=False), outputs=post_to_hub)


demo.launch() # (share=True, inline=False, debug=True)