Megatron 3D & Pipeline Bubble

Megatron-LM's 3D parallelism (TP × PP × DP) implementation. TP column/row parallel with dual AllReduce: inside each transformer layer, the attention QKV projection is column-parallel (output split along hidden), the output projection row-parallel (input split along hidden), and the MLP the same — a column-parallel branch's output is a per-rank partial sum that must be AllReduced to recover the full result; the forward has 2 AllReduces per layer (attention output projection + MLP's second linear), the backward has 2 more (column-parallel weight gradients), 4 per layer per iteration, each message about $b \times s \times h$. PP 1F1B schedule: the naive greedy schedule (all microbatches' forwards, then all backwards) keeps m microbatches' activations alive at once, peaking memory and starting backprop late; 1F1B (one-forward-one-backward) alternates one forward and one backward per stage, capping in-flight microbatches at p and lowering peak memory. Bubble formula: with p stages and m microbatches per batch, the steady-state bubble fraction is $\text{bubble} = \frac{p-1}{m+p-1}$ — the larger m, the smaller the bubble for fixed p; numeric examples: p=4, m=4 → $\frac{3}{7} \approx 43\%$; p=8, m=32 → $\frac{7}{39} \approx 18\%$; p=8, m=64 → $\frac{7}{71} \approx 10\%$; approaching 0 as m grows. Interleaved (V-shaped/chunked) scheduling splits each stage's layers into c chunks, cutting the bubble by roughly c× at extra communication and memory cost. Ring-Attention sequence parallelism: the sequence is split into N chunks across N ranks; each rank first computes partial attention on its local chunk, then KV blocks circulate around the ring, and after N−1 rounds every rank has seen all KV — per-rank KV memory is O(1) (independent of seq length), per-layer communication ≈ 2× the full KV size (each rank reads all K and V once), bandwidth is constant while latency grows with N−1 serial hops; it shares the online-softmax technique with FlashAttention.

the hot zone of training-infra interviews: derive the bubble formula, why TP runs 2 AllReduces per layer, how 1F1B beats greedy, how Ring-Attention's memory and communication work out, and how a 175B model is configured in 3D.

TP spreads per-layer compute across GPUs at the cost of frequent small messages (only intra-node NVLink can carry them); PP splits layers so memory is amortized, but idle bubbles eat utilization — 1F1B squeezes memory, large m squeezes the bubble, interleaved trades further; Ring-Attention parallelizes the longest dimension (the sequence) so ultra-long-context training no longer hits the per-GPU KV memory ceiling.

Test Your Knowledge: Practice Questions for "Megatron 3D & Pipeline Bubble"

Single choice pitfall questions with instant feedback and mistake tracking.