update

app.py

@@ -70,10 +70,23 @@ def format_config_display(config):
     if not config:
         return "No configuration loaded"
 
+    # Calculate number of parameters
+    vocab_embeddings = config['vocab_size'] * config['hidden_size'] * (1 if config['tie_word_embeddings'] else 2)
+    
+    layer_params = (
+        (config['hidden_size'] * config['hidden_size'] * (1 + 2*config['num_key_value_heads']/config['num_attention_heads']))  # qkv_proj
+        + (config['hidden_size'] * config['hidden_size'])     # out_proj
+        + (config['hidden_size'] * 2 * config['intermediate_size'])  # gate_up_proj
+        + (config['intermediate_size'] * config['hidden_size'])      # down_proj
+    )
+    total_params = (vocab_embeddings + config['num_layers'] * layer_params) / config['tp']
+    params_billions = total_params / 1_000_000_000
+
     sections = {
         "Model Architecture": [
-            "hidden_size", "num_layers", "vocab_size",
-            "intermediate_size", "tie_word_embeddings", "num_attention_heads", "num_key_value_heads"
+            "hidden_size", "num_layers", "vocab_size", 
+            "intermediate_size", "tie_word_embeddings", "num_attention_heads", "num_key_value_heads",
+            ("num_params", f"{params_billions:.2f}B")  # Show params in billions
         ],
         "Training Configuration": [
             "seq_len", "mbs", "batch_accum"
@@ -87,8 +100,13 @@ def format_config_display(config):
     for section_name, params in sections.items():
         output += f"<div style='flex: 1; padding-right: 20px;'><h3>{section_name}</h3>"
         for param in params:
-            value = config.get(param, 'N/A')
-            output += f"<b>{param}</b>: {value}<br>"
+            if isinstance(param, tuple):
+                # Handle custom parameter display
+                param_name, value = param
+                output += f"<b>{param_name}</b>: {value}<br>"
+            else:
+                value = config.get(param, 'N/A')
+                output += f"<b>{param}</b>: {value}<br>"
         output += "</div>"
     output += "</div>"
     return output

utils.py

@@ -71,6 +71,9 @@ def calculate_memory_components(
     
     model_bf16_full = (vocab_embeddings + num_hidden_layers_in_pp * layer_params) * (2 / 1024 / 1024) / tp
 
+    # Calculate number of parameters in billions
+    num_params_in_B = (vocab_embeddings + num_layers*layer_params) / 1e9
+
     # Adjust model components based on ZeRO stage
     if zero_stage == 3:
         # In ZeRO-3, model parameters are sharded across dp ranks
@@ -148,7 +151,8 @@ def calculate_memory_components(
             "DDP Gradient Buffers": ddp_grads_buffers,
             "ZeRO-3 Buffers": zero3_buffers,
             "Overhead": overhead,
-            "Activations": activs
+            "Activations": activs,
+            "num_params_in_B": num_params_in_B
         },
         "Aggregates": {
             "Memory Before Optimizer States": memory_usage_before_optimstates,