EXPERIMENTAL: add zero2 zero3

app.py

@@ -107,7 +107,7 @@ with gr.Blocks() as demo:
                     tp = gr.Number(1, label="Tensor Parallelism")
                     pp = gr.Number(1, label="Pipeline Parallelism")
                     dp = gr.Number(1, label="Data Parallelism")
-                    zero_stage = gr.Radio([0, 1], value=0, label="ZeRO Stage")
+                    zero_stage = gr.Radio([0, 1, 2, 3], value=0, label="ZeRO Stage")
 
                 manual_submit = gr.Button("Calculate Memory (Manual Input)")
         with gr.Column(scale=2):

utils.py

@@ -57,7 +57,6 @@ def calculate_memory_components(
     if pp == 1:
         num_hidden_layers_in_pp = num_layers
     else:
-        # num_hidden_layers_in_pp = num_layers // pp
         num_hidden_layers_in_pp = get_num_hidden_layers_in_pp(hidden_size, num_layers, vocab_size, intermediate_size, num_attention_heads, pp)
     
     # Model BF16 calculation
@@ -70,13 +69,26 @@ def calculate_memory_components(
         + (intermediate_size * hidden_size)      # down_proj
     )
     
-    model_bf16 = (vocab_embeddings + num_hidden_layers_in_pp * layer_params) * (2 / 1024 / 1024) / tp
+    model_bf16_full = (vocab_embeddings + num_hidden_layers_in_pp * layer_params) * (2 / 1024 / 1024) / tp
+
+    # Adjust model components based on ZeRO stage
+    if zero_stage == 3:
+        # In ZeRO-3, model parameters are sharded across dp ranks
+        model_bf16 = model_bf16_full / dp
+        fp32_params = 2 * model_bf16
+        fp32_grads = 2 * model_bf16
+        optimstates = 4 * model_bf16
+        # Additional communication buffers for ZeRO-3
+        zero3_buffers = 2 * model_bf16  # For parameter gathering during forward/backward
+    else:
+        # For ZeRO-0/1/2
+        dp_if_zero = 1 if zero_stage == 0 else dp
+        model_bf16 = model_bf16_full
+        fp32_params = 2 * model_bf16 / dp_if_zero
+        fp32_grads = 2 * model_bf16
+        optimstates = 4 * model_bf16 / dp_if_zero
+        zero3_buffers = 0
 
-    # Other components
-    dp_if_zero = 1 if zero_stage == 0 else dp
-    fp32_params = 2 * model_bf16
-    fp32_grads = 2 * model_bf16
-    optimstates = 4 * model_bf16
     use_ddp = zero_stage == 0 and dp > 1
     ddp_grads_buffers = model_bf16 if use_ddp else 0
     overhead = 72 + 32 * mbs
@@ -84,7 +96,6 @@ def calculate_memory_components(
     # Activations
     is_mha = num_key_value_heads == num_attention_heads
     decoder_layer_mib = (seq_len * mbs * hidden_size/tp) * (2/1024/1024) * (4*intermediate_size/hidden_size + 6 + 2*num_key_value_heads/num_attention_heads + 2)
-    # decoder_layer_mib = (seq_len * mbs * hidden_size/tp) * (2/1024/1024) * (4*intermediate_size/hidden_size + 12 + 2*num_key_value_heads/num_attention_heads + (2 if is_mha else 0))
     
     if pp > 1:
         activs = min(pp, batch_accum) * num_hidden_layers_in_pp * decoder_layer_mib
@@ -95,26 +106,29 @@ def calculate_memory_components(
     # Calculate aggregate metrics
     memory_usage_after_optimstates = (
         model_bf16 + 
-        fp32_params/dp_if_zero + 
+        fp32_params + 
         fp32_grads + 
-        optimstates/dp_if_zero + 
+        optimstates + 
         ddp_grads_buffers + 
+        zero3_buffers +
         overhead
     )
 
     memory_usage_before_optimstates = (
         model_bf16 + 
-        fp32_params/dp_if_zero + 
+        fp32_params + 
         fp32_grads + 
-        ddp_grads_buffers
+        ddp_grads_buffers +
+        zero3_buffers
     )
 
     memory_usage_peak_tbi = (
         model_bf16 + 
-        fp32_params/dp_if_zero + 
+        fp32_params + 
         fp32_grads + 
-        optimstates/dp_if_zero + 
+        optimstates + 
         ddp_grads_buffers + 
+        zero3_buffers +
         overhead + 
         activs
     )
@@ -122,10 +136,11 @@ def calculate_memory_components(
     return {
         "Components": {
             "Model BF16": model_bf16,
-            "FP32 Parameters": fp32_params/dp_if_zero,
+            "FP32 Parameters": fp32_params,
             "FP32 Gradients": fp32_grads,
-            "Optimizer States": optimstates/dp_if_zero,
+            "Optimizer States": optimstates,
             "DDP Gradient Buffers": ddp_grads_buffers,
+            "ZeRO-3 Buffers": zero3_buffers,
             "Overhead": overhead,
             "Activations": activs
         },
@@ -189,28 +204,33 @@ def plot_memory_breakdown(
         "Model Init": [
             ("Model BF16", c["Model BF16"]),
             ("DDP Gradient Buffers", c["DDP Gradient Buffers"]),
+            ("ZeRO-3 Buffers", c["ZeRO-3 Buffers"]),
         ],
         "Gradient Accumulator Init": [
             ("Model BF16", c["Model BF16"]),
             ("DDP Gradient Buffers", c["DDP Gradient Buffers"]),
+            ("ZeRO-3 Buffers", c["ZeRO-3 Buffers"]),
             ("FP32 Parameters", c["FP32 Parameters"]),
             ("FP32 Gradients", c["FP32 Gradients"])
         ],
         "Fwd-Bwd Peak": [
             ("Model BF16", c["Model BF16"]),
             ("DDP Gradient Buffers", c["DDP Gradient Buffers"]),
+            ("ZeRO-3 Buffers", c["ZeRO-3 Buffers"]),
             ("FP32 Parameters", c["FP32 Parameters"]),
             ("FP32 Gradients", c["FP32 Gradients"]),
             ("Activations", c["Activations"])
         ],
         "Optimizer Step": [
             ("Model BF16", c["Model BF16"]),
+            ("ZeRO-3 Buffers", c["ZeRO-3 Buffers"]),
             ("FP32 Parameters", c["FP32 Parameters"]),
             ("FP32 Gradients", c["FP32 Gradients"]),
             ("Optimizer States", c["Optimizer States"])
         ],
         "2nd Fwd-Bwd Peak": [
             ("Model BF16", c["Model BF16"]),
+            ("ZeRO-3 Buffers", c["ZeRO-3 Buffers"]),
             ("FP32 Parameters", c["FP32 Parameters"]),
             ("FP32 Gradients", c["FP32 Gradients"]),
             ("Optimizer States", c["Optimizer States"]),
@@ -219,6 +239,7 @@ def plot_memory_breakdown(
         ],
         "2nd Optimizer Step": [
             ("Model BF16", c["Model BF16"]),
+            ("ZeRO-3 Buffers", c["ZeRO-3 Buffers"]),
             ("FP32 Parameters", c["FP32 Parameters"]),
             ("FP32 Gradients", c["FP32 Gradients"]),
             ("Optimizer States", c["Optimizer States"]),