import matplotlib.pyplot as plt 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 = [] block_costs = [] # Embedding layer (treated as zero cost in the original implementation) pipeline_blocks.append("embedding") block_costs.append(0) # Decoder layers decoder_cost = (4 * num_attention_heads * (hidden_size//num_attention_heads) * hidden_size + 3 * intermediate_size * hidden_size) for _ in range(num_layers): pipeline_blocks.append("decoder") block_costs.append(decoder_cost) # LM head pipeline_blocks.append("lm_head") block_costs.append(vocab_size * hidden_size) # Now follow the same logic as the original code total_cost = sum(block_costs) target_cost_per_rank = total_cost / pp_size blocks_in_rank0 = 0 current_cost = 0 for block_idx, block_cost in enumerate(block_costs): current_cost += block_cost blocks_in_rank0 += 1 # Check if we should move to next rank remaining_ranks = pp_size - 1 # -1 because we're calculating for rank 0 remaining_nonzero_blocks = sum(1 for c in block_costs[block_idx+1:] if c > 0) if (remaining_ranks > 0 and remaining_nonzero_blocks <= remaining_ranks) or (current_cost >= target_cost_per_rank): break num_hidden_layers_in_pp = blocks_in_rank0 - 1 # We exclude first rank as it's the embedding layer print("num_hidden_layers_in_pp", num_hidden_layers_in_pp) return num_hidden_layers_in_pp def calculate_memory_components( hidden_size, num_attention_heads, num_key_value_heads, num_layers, vocab_size, intermediate_size, seq_len, mbs, batch_accum, tp, pp, dp, zero_stage, tie_word_embeddings ): # Calculate base components first 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 vocab_embeddings = vocab_size * hidden_size * (2 if (not tie_word_embeddings and pp==1) else 1) layer_params = ( (hidden_size * hidden_size * (1 + 2*num_key_value_heads/num_attention_heads)) # qkv_proj + (hidden_size * hidden_size) # out_proj + (hidden_size * 2 * intermediate_size) # gate_up_proj + (intermediate_size * hidden_size) # down_proj ) model_bf16 = (vocab_embeddings + num_hidden_layers_in_pp * layer_params) * (2 / 1024 / 1024) / tp # 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 # Activations # decoder_layer_mib = (seq_len * mbs * hidden_size/tp) * (2/1024/1024) * (4*intermediate_size/hidden_size + 10) 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 + 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 else: cast_to_fp32 = sharded_cross_entropy = seq_len * mbs * vocab_size * (2 / 1024 / 1024) * 2 / tp activs = num_layers * decoder_layer_mib + cast_to_fp32 + sharded_cross_entropy # Calculate aggregate metrics memory_usage_after_optimstates = ( model_bf16 + fp32_params/dp_if_zero + fp32_grads + optimstates/dp_if_zero + ddp_grads_buffers + overhead ) memory_usage_before_optimstates = ( model_bf16 + fp32_params/dp_if_zero + fp32_grads + ddp_grads_buffers ) memory_usage_peak_tbi = ( model_bf16 + fp32_params/dp_if_zero + fp32_grads + optimstates/dp_if_zero + ddp_grads_buffers + overhead + activs ) return { "Components": { "Model BF16": model_bf16, "FP32 Parameters": fp32_params/dp_if_zero, "FP32 Gradients": fp32_grads, "Optimizer States": optimstates/dp_if_zero, "DDP Gradient Buffers": ddp_grads_buffers, "Overhead": overhead, "Activations": activs }, "Aggregates": { "Memory Before Optimizer States": memory_usage_before_optimstates, "Memory After Optimizer States": memory_usage_after_optimstates, "Peak Memory (TBI)": memory_usage_peak_tbi } } def plot_memory_breakdown( hidden_size, num_attention_heads, num_key_value_heads, num_layers, vocab_size, intermediate_size, seq_len, mbs, batch_accum, tp, pp, dp, zero_stage, tie_word_embeddings ): results = calculate_memory_components( hidden_size, num_attention_heads, num_key_value_heads, num_layers, vocab_size, intermediate_size, seq_len, mbs, batch_accum, tp, pp, dp, zero_stage, tie_word_embeddings ) # Create figure for components plot plt.close('all') fig1 = plt.figure(figsize=(10, 6)) ax1 = fig1.add_subplot(1, 1, 1) # Plot components components = results["Components"] names = list(components.keys()) values = list(components.values()) bars1 = ax1.bar(range(len(components)), values) # Add value labels with better positioning for bar in bars1: height = bar.get_height() ax1.text(bar.get_x() + bar.get_width()/2., height, f'{height:.1f} MiB', ha='center', va='bottom', rotation=0) # Remove rotation for better readability # Customize the first plot ax1.set_xticks(range(len(components))) ax1.set_xticklabels(names, rotation=45, ha='right') ax1.set_ylabel('Memory (MiB)') ax1.set_title('Memory Component Breakdown', pad=20) plt.tight_layout() # Create figure for aggregates plot fig2 = plt.figure(figsize=(10, 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') # Customize the second plot ax2.set_xticks(range(len(aggregates))) ax2.set_xticklabels(names, 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)) # Adjust layout to prevent text overlap plt.tight_layout() return fig1, fig2