| | import gradio as gr |
| | from typing import Dict, Union |
| | from huggingface_hub import get_safetensors_metadata, hf_hub_download, login |
| | import json |
| |
|
| | |
| | bytes_per_dtype: Dict[str, float] = { |
| | "int4": 0.5, |
| | "int8": 1, |
| | "float8": 1, |
| | "float16": 2, |
| | "float32": 4, |
| | } |
| |
|
| | def extract_keys(json_obj, keys_to_extract): |
| | extracted_values = {} |
| | def recursive_search(obj): |
| | if isinstance(obj, dict): |
| | for key, value in obj.items(): |
| | if key in keys_to_extract: |
| | extracted_values[key] = value |
| | recursive_search(value) |
| | elif isinstance(obj, list): |
| | for item in obj: |
| | recursive_search(item) |
| | recursive_search(json_obj) |
| | return extracted_values |
| |
|
| | def calculate_kv_cache_memory(context_size: int, model_id: str, dtype: str, token: str = None): |
| | try: |
| | file_path = hf_hub_download(repo_id=model_id, filename="config.json", token=token) |
| | with open(file_path, 'r') as f: |
| | config = json.load(f) |
| | |
| | keys_to_find = {"num_hidden_layers", "num_key_value_heads", "hidden_size", "num_attention_heads"} |
| | config = extract_keys(config, keys_to_find) |
| |
|
| | num_layers = config["num_hidden_layers"] |
| | num_att_heads = config.get("num_key_value_heads", config["num_attention_heads"]) |
| | dim_att_head = config["hidden_size"] // config["num_attention_heads"] |
| | dtype_bytes = bytes_per_dtype[dtype] |
| | |
| | memory_per_token = num_layers * num_att_heads * dim_att_head * dtype_bytes * 2 |
| | context_size_memory_footprint_gb = (context_size * memory_per_token) / 1_000_000_000 |
| | |
| | return context_size_memory_footprint_gb |
| | except Exception as e: |
| | return f"Error: {str(e)}" |
| |
|
| | def calculate_model_memory(parameters: float, dtype: str) -> float: |
| | bytes = bytes_per_dtype[dtype] |
| | return round((parameters * 4) / (32 / (bytes * 8)) * 1.18, 2) |
| |
|
| | def get_model_size(model_id: str, dtype: str, token: str = None) -> Union[float, str]: |
| | try: |
| | metadata = get_safetensors_metadata(model_id, token=token) |
| | if not metadata or not metadata.parameter_count: |
| | return "Error: Could not fetch metadata." |
| | model_parameters = int(list(metadata.parameter_count.values())[0]) / 1_000_000_000 |
| | return calculate_model_memory(model_parameters, dtype) |
| | except Exception as e: |
| | return f"Error: {str(e)}" |
| |
|
| | def estimate_vram(model_id, dtype, context_size, hf_token): |
| | if hf_token: |
| | login(token=hf_token) |
| | |
| | if dtype not in bytes_per_dtype: |
| | return "Error: Unsupported dtype" |
| | |
| | model_memory = get_model_size(model_id, dtype, hf_token) |
| | context_memory = calculate_kv_cache_memory(context_size, model_id, dtype, hf_token) |
| | |
| | if isinstance(model_memory, str) or isinstance(context_memory, str): |
| | return model_memory if isinstance(model_memory, str) else context_memory |
| | |
| | total_memory = model_memory + context_memory |
| | return f"Model VRAM: {model_memory:.2f} GB\nContext VRAM: {context_memory:.2f} GB\nTotal VRAM: {total_memory:.2f} GB" |
| |
|
| | iface = gr.Interface( |
| | fn=estimate_vram, |
| | inputs=[ |
| | gr.Textbox(label="Hugging Face Model ID", value="google/gemma-3-27b-it"), |
| | gr.Dropdown(choices=list(bytes_per_dtype.keys()), label="Data Type", value="float16"), |
| | gr.Number(label="Context Size", value=128000), |
| | gr.Textbox(label="Hugging Face Access Token", type="password", placeholder="Optional - Needed for gated models") |
| | ], |
| | outputs=gr.Textbox(label="Estimated VRAM Usage"), |
| | title="LLM GPU VRAM Calculator", |
| | description="Estimate the VRAM requirements of a model and context size. Optionally provide a Hugging Face token for gated models." |
| | ) |
| |
|
| | iface.launch() |
| |
|
