Nearest B or M

app.py

@@ -8,7 +8,7 @@ def format_params(params):
     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"] else 4
+    bytes_per_param = 2 if precision in ["FP16/BF16", "BF16"] else 4
     
     # Model Weights
     model_memory = params * bytes_per_param
@@ -28,7 +28,7 @@ def calculate_training_memory(params, precision, batch_size, seq_length, num_hea
     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"] else 4
+    bytes_per_param = 2 if precision in ["FP16/BF16", "BF16"] else 4
     
     # Model Weights
     model_memory = params * bytes_per_param
@@ -42,7 +42,7 @@ def calculate_inference_memory(params, precision, batch_size, seq_length, num_he
     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"] else 4
+    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
@@ -56,7 +56,7 @@ with gr.Blocks() as app:
         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")
+                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)
@@ -71,7 +71,7 @@ with gr.Blocks() as app:
         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")
+                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)
@@ -91,7 +91,7 @@ with gr.Blocks() as app:
                 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")
+                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)