File size: 5,133 Bytes
3c83b5b 2a4dcce 3c83b5b 872d649 3c83b5b 872d649 3c83b5b 872d649 3c83b5b 872d649 3c83b5b 872d649 3c83b5b 872d649 3c83b5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import gradio as gr
def format_params(params):
if params >= 1e9:
return f"{params / 1e9:.2f}B"
elif params >= 1e6:
return f"{params / 1e6:.2f}M"
return str(params)
def calculate_training_memory(params, precision, batch_size, seq_length, num_heads, head_dim, num_layers):
bytes_per_param = 2 if precision in ["FP16/BF16", "BF16"] else 4
# Model Weights
model_memory = params * bytes_per_param
# Optimizer States (Adam)
optimizer_memory = model_memory * 2
# Gradients
gradient_memory = model_memory
# Activation Memory (approximate formula)
activation_memory = batch_size * seq_length * num_heads * head_dim * num_layers * bytes_per_param
# Total Training Memory
total_memory = model_memory + optimizer_memory + gradient_memory + activation_memory
return f"Model Weights: {model_memory / 1e9:.2f} GB\nOptimizer: {optimizer_memory / 1e9:.2f} GB\nGradients: {gradient_memory / 1e9:.2f} GB\nActivation Memory: {activation_memory / 1e9:.2f} GB\nTotal Training Memory: {total_memory / 1e9:.2f} GB"
def calculate_inference_memory(params, precision, batch_size, seq_length, num_heads, head_dim, num_layers):
bytes_per_param = 2 if precision in ["FP16/BF16", "BF16"] else 4
# Model Weights
model_memory = params * bytes_per_param
# KV Cache
kv_cache_memory = batch_size * seq_length * num_heads * head_dim * 2 * num_layers * bytes_per_param
# Total Inference Memory
total_memory = model_memory + kv_cache_memory
return f"Model Weights: {model_memory / 1e9:.2f} GB\nKV Cache: {kv_cache_memory / 1e9:.2f} GB\nTotal Inference Memory: {total_memory / 1e9:.2f} GB"
def calculate_kv_cache(batch_size, seq_length, num_heads, head_dim, num_layers, precision):
bytes_per_param = 2 if precision in ["FP16/BF16", "BF16"] else 4
# KV Cache Calculation
kv_cache_memory = batch_size * seq_length * num_heads * head_dim * 2 * num_layers * bytes_per_param
return f"KV Cache Memory: {kv_cache_memory / 1e9:.2f} GB"
with gr.Blocks() as app:
gr.Markdown("# GPU Memory Calculator for Transformer Models")
with gr.Tabs():
with gr.Tab("Training Memory Calculation"):
with gr.Row():
params = gr.Number(label="Number of Parameters (e.g., 175B = 175e9)", value=175e9)
precision = gr.Radio(["FP16/BF16", "FP32"], label="Precision", value="FP16/BF16")
with gr.Row():
batch_size = gr.Number(label="Batch Size", value=1)
seq_length = gr.Number(label="Sequence Length", value=2048)
with gr.Row():
num_heads = gr.Number(label="Number of Attention Heads", value=96)
head_dim = gr.Number(label="Head Dimension", value=128)
num_layers = gr.Number(label="Number of Layers", value=96)
train_button = gr.Button("Calculate Training Memory")
train_output = gr.Textbox(label="Training Memory Usage")
train_button.click(calculate_training_memory, [params, precision, batch_size, seq_length, num_heads, head_dim, num_layers], train_output)
with gr.Tab("Inference Memory Calculation"):
with gr.Row():
params_inf = gr.Number(label="Number of Parameters (e.g., 175B = 175e9)", value=175e9)
precision_inf = gr.Radio(["FP16/BF16", "FP32"], label="Precision", value="FP16/BF16")
with gr.Row():
batch_size_inf = gr.Number(label="Batch Size", value=1)
seq_length_inf = gr.Number(label="Sequence Length", value=2048)
with gr.Row():
num_heads_inf = gr.Number(label="Number of Attention Heads", value=96)
head_dim_inf = gr.Number(label="Head Dimension", value=128)
num_layers_inf = gr.Number(label="Number of Layers", value=96)
infer_button = gr.Button("Calculate Inference Memory")
infer_output = gr.Textbox(label="Inference Memory Usage")
infer_button.click(calculate_inference_memory, [params_inf, precision_inf, batch_size_inf, seq_length_inf, num_heads_inf, head_dim_inf, num_layers_inf], infer_output)
with gr.Tab("KV Cache Calculation"):
with gr.Row():
batch_size_kv = gr.Number(label="Batch Size", value=1)
seq_length_kv = gr.Number(label="Sequence Length", value=2048)
with gr.Row():
num_heads_kv = gr.Number(label="Number of Attention Heads", value=96)
head_dim_kv = gr.Number(label="Head Dimension", value=128)
num_layers_kv = gr.Number(label="Number of Layers", value=96)
precision_kv = gr.Radio(["FP16/BF16", "FP32"], label="Precision", value="FP16/BF16")
kv_button = gr.Button("Calculate KV Cache Memory")
kv_output = gr.Textbox(label="KV Cache Memory Usage")
kv_button.click(calculate_kv_cache, [batch_size_kv, seq_length_kv, num_heads_kv, head_dim_kv, num_layers_kv, precision_kv], kv_output)
app.launch()
|