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.

Hand-derive the bubble: p stages, m microbatches → $\text{bubble}=\frac{p-1}{m+p-1}$; give numeric examples (p=4,m=4 → 3/7≈43%; p=8,m=32 → 7/39≈18%) and how to reduce it (larger m, interleaved chunks).

Describe TP column/row parallel with dual AllReduce: why column-parallel outputs need AllReduce summation and why row-parallel backward does too; how many per layer per iteration (2 forward + 2 backward = 4).

Compare greedy vs 1F1B scheduling: why 1F1B has lower peak memory (in-flight microbatches ≤ p) and earlier backprop start; why the bubble is not eliminated (inter-stage pipeline stalls).

Derive Ring-Attention: why per-rank KV memory is O(1) after splitting the sequence into N chunks; per-layer communication ≈ 2× the full KV; bandwidth/latency profile (constant bandwidth, N−1 serial hops) and its relation to FlashAttention.

3D integration: explain the DP/TP/PP division by communication level (TP intra-node frequent small messages, PP inter-node infrequent large messages, DP gradient AllReduce), and why 3D parallelism grows parallelism while cutting per-GPU memory and communication cost.