Add application file

app.py

@@ -0,0 +1,93 @@
+import gradio as gr
+
+def calculate_training_memory(params, precision, batch_size, seq_length, num_heads, head_dim, num_layers):
+    bytes_per_param = 2 if precision == "FP16" 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 == "FP16" 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 == "FP16" 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", "FP32"], label="Precision", value="FP16")
+            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", "FP32"], label="Precision", value="FP16")
+            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", "FP32"], label="Precision", value="FP16")
+            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()
+