timeline view

utils.py

@@ -1,5 +1,6 @@
-
 import matplotlib.pyplot as plt
+import numpy as np
+
 def get_num_hidden_layers_in_pp(hidden_size, num_layers, vocab_size, intermediate_size, num_attention_heads, pp_size):
     # Get list of pipeline blocks and their costs
     pipeline_blocks = []
@@ -169,33 +170,92 @@ def plot_memory_breakdown(
     
     plt.tight_layout()
     
-    # Create figure for aggregates plot
-    fig2 = plt.figure(figsize=(10, 6))
+    # Create figure for timeline plot
+    fig2 = plt.figure(figsize=(12, 6))
     ax2 = fig2.add_subplot(1, 1, 1)
     
-    # Plot aggregate metrics
-    aggregates = results["Aggregates"]
-    names = list(aggregates.keys())
-    values = list(aggregates.values())
-    
-    bars2 = ax2.bar(range(len(aggregates)), values, color='orange')
-    
-    # Add value labels
-    for bar in bars2:
-        height = bar.get_height()
-        ax2.text(bar.get_x() + bar.get_width()/2., height,
-                f'{height:.1f} MiB',
-                ha='center', va='bottom')
+    # Define timeline steps and their components
+    c = results["Components"]
+    timeline_steps = {
+        "Model Init": [
+            ("Model BF16", c["Model BF16"]),
+            ("DDP Gradient Buffers", c["DDP Gradient Buffers"]),
+        ],
+        "Gradient Accumulator Init": [
+            ("Model BF16", c["Model BF16"]),
+            ("DDP Gradient Buffers", c["DDP Gradient 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"]),
+            ("FP32 Parameters", c["FP32 Parameters"]),
+            ("FP32 Gradients", c["FP32 Gradients"]),
+            ("Activations", c["Activations"])
+        ],
+        "After Fwd-Bwd": [
+            ("Model BF16", c["Model BF16"]),
+            ("DDP Gradient Buffers", c["DDP Gradient Buffers"]),
+            ("FP32 Parameters", c["FP32 Parameters"]),
+            ("FP32 Gradients", c["FP32 Gradients"])
+        ],
+        "Optimizer Step": [
+            ("Model BF16", c["Model BF16"]),
+            ("FP32 Parameters", c["FP32 Parameters"]),
+            ("FP32 Gradients", c["FP32 Gradients"]),
+            ("Optimizer States", c["Optimizer States"])
+        ],
+        "2nd Fwd-Bwd Peak": [
+            ("Model BF16", c["Model BF16"]),
+            ("FP32 Parameters", c["FP32 Parameters"]),
+            ("FP32 Gradients", c["FP32 Gradients"]),
+            ("Optimizer States", c["Optimizer States"]),
+            ("DDP Gradient Buffers", c["DDP Gradient Buffers"]),
+            ("Activations", c["Activations"])
+        ],
+        "2nd Optimizer Step": [
+            ("Model BF16", c["Model BF16"]),
+            ("FP32 Parameters", c["FP32 Parameters"]),
+            ("FP32 Gradients", c["FP32 Gradients"]),
+            ("Optimizer States", c["Optimizer States"]),
+            ("DDP Gradient Buffers", c["DDP Gradient Buffers"])
+        ]
+    }
+    # Plot timeline
+    x = range(len(timeline_steps))
+    bottom = np.zeros(len(timeline_steps))
+    colors = plt.cm.Set3(np.linspace(0, 1, len(c)))
+    color_map = dict(zip(c.keys(), colors))
     
-    # Customize the second plot
-    ax2.set_xticks(range(len(aggregates)))
-    ax2.set_xticklabels(names, rotation=45, ha='right')
+    for component in c.keys():
+        heights = []
+        for step_components in timeline_steps.values():
+            height = 0
+            for comp_name, comp_value in step_components:
+                if comp_name == component:
+                    height = comp_value
+            heights.append(height)
+        
+        ax2.bar(x, heights, bottom=bottom, label=component, color=color_map[component])
+        bottom += heights
+
+    # Customize the timeline plot
+    ax2.set_xticks(x)
+    ax2.set_xticklabels(timeline_steps.keys(), rotation=45, ha='right')
     ax2.set_ylabel('Memory (MiB)')
-    ax2.set_title('Aggregate Memory Metrics', pad=20)
-    max_y_value = max(values)
-    ax2.set_ylim(0, max(80000, max_y_value))
+    ax2.set_title('Memory Timeline', pad=20)
+    ax2.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
     
-    # Adjust layout to prevent text overlap
+    # Add total memory labels on top of each bar
+    for i, total in enumerate(bottom):
+        ax2.text(i, total, f'{total:.1f} MiB', ha='center', va='bottom')
+    
+    # Adjust layout
     plt.tight_layout()
     
+    # Set y-axis limit
+    max_y_value = max(bottom)
+    ax2.set_ylim(0, max(80000, max_y_value))
+    
     return fig1, fig2