Parallelism II

The new unit of compute is the data center

Percy covered the mechanics of parallelism last lecture; this one is Tatsu’s tour of “knowledge, details, and trivia about how modern parallelism for language model training works.” Two bottlenecks force us onto many machines: compute (the fastest supercomputers have exaflops; one GPU is nowhere close) and memory (a single GPU can’t fit most large models, so they must be sharded). The whole lecture lives at the level of collective communication primitives — all-reduce, all-gather, reduce-scatter — never packets, and the one identity to keep in your pocket is that an all-reduce equals a reduce-scatter followed by an all-gather, at the same cost in the bandwidth-limited regime. That equivalence is what makes the first big algorithm of the day free.

The hardware substrate shapes which strategies win. TPUs network as a toroidal mesh — every chip talks to its neighbors, wrap-around at the edges — so the topology scales indefinitely and neighbor-to-neighbor traffic is cheap and beefy. GPUs use a fat tree: GPUs tightly connected in a box, boxes in pods, pods joined by spine switches — more flexible for “random,” unpredictable communication. Bill Dally and Jeff Dean framed it as workloads: GPUs suit mixture-of-experts models whose tokens route unpredictably; TPUs suit dense models with very regular partitions. And the morning of the lecture, Google announced TPU8i/t with much more all-to-all, tree-style connectivity — “a little bit of a convergent evolution,” because modern models are MoEs and the workload is defining the network.

What we want from multi-machine scaling is simple to state: linear memory scaling, linear compute scaling, and losslessness — use everything you paid for.

Data parallelism and the ZeRO ladder

Naive data parallelism splits a batch of $B$ examples across $M$ machines, each computes gradients on $B/M$ examples, and one all-reduce synchronizes them — communication of $2 \times$ the parameter count per step. Compute scaling is perfect as long as there are enough examples per GPU. Memory scaling is zero: every GPU replicates everything. And the memory situation “is actually just terrible,” because training state is far more than parameters. Rule of thumb: five copies of the weights, 16 bytes per parameter — 2 (bf16 parameters) + 2 (bf16 gradients) + 4 (fp32 master weights) + 4 + 4 (fp32 Adam first and second moments). The optimizer state dominates.

ZeRO’s core idea: shard the expensive parts across the data-parallel workers, and use the reduce-scatter equivalence to do it for free. Three stages, one per thing sharded:

• ZeRO stage 1 shards optimizer state only. Each worker owns the update for one slice of the parameters: everyone computes a full gradient on their data, reduce-scatters gradients so each worker gets the summed gradient for its slice, updates its slice, then all-gathers the updated parameters. Reduce-scatter + all-gather = all-reduce, so this costs exactly what naive DDP costs — “free memory savings.”
• ZeRO stage 2 shards gradients too. You can never materialize the full gradient vector, so it becomes a systems trick: sweep backward through the compute graph, reduce each layer’s gradient to its owner as soon as it’s computed, and free it immediately. Incremental or all-at-once, same total communication.
• ZeRO stage 3 (FSDP) shards the parameters as well — no GPU ever holds the whole model. Parameters are all-gathered on demand layer by layer: gather, forward, free; then in the backward pass gather again, compute, reduce-scatter the gradients out, free. Two all-gathers plus one reduce-scatter, so $3\times$ parameters — 1.5x the communication.

Scheme	Primitive	Comm cost	Memory per GPU
Naive DDP	one all-reduce	2 × params	$(2+2+K)\Psi$
ZeRO 1	reduce-scatter + all-gather	2 × params	$2\Psi + 2\Psi + K\Psi/N_d$
ZeRO 2	incremental reduce + all-gather	2 × params	$2\Psi + (2+K)\Psi/N_d$
ZeRO 3 (FSDP)	2 all-gathers + reduce-scatter	3 × params	$(2+2+K)\Psi/N_d$

With $K=12$, $\Psi = 7.5$B and 64 GPUs, that’s 120 GB per GPU down to 31.4, 16.6, and finally 1.9 GB.

Why data parallel isn’t enough — and pipeline parallelism

Data parallelism consumes a finite resource: batch size. With batch 8 you can never use more than eight workers, and you can’t just grow the batch forever — past the critical batch size, an extra batch element helps less than an extra SGD step on it; “an infinitely large batch size is not infinitely better than infinite steps.” And ZeRO stages 1 and 2 don’t scale memory across model size at all, while stage 3 never touches activation memory. So we turn to model parallelism: like FSDP it splits parameters across GPUs, but the conceptual difference is that it communicates activations, not parameters.

Pipeline parallelism cuts along depth: layers 0–7 on GPU 0, 8–15 on GPU 1, and so on, passing activations forward and partial gradients backward. Done naively, utilization is terrible — with $n$ GPUs each is active $1/n$ of the time, waiting in the bubble. The fix is micro-batching: send off micro-batch 1 and immediately start micro-batch 2, so stages overlap.

So why endure it? Pipelines save memory versus DDP, and above all they have the best communication profile available: only activations ($b \times s \times h$, almost always smaller than a parameter matrix) and strictly point to point. In practice, pipelines live on your slowest links — across nodes, pods, even data centers. Cleverer schedules shrink the bubble further (the DeepSeek interleaving), and zero-bubble pipelining exploits the structure of backward itself: each node must (1) propagate partials down the graph — urgent, the next stage is blocked on it — and (2) compute the weight gradient — a leaf, doable “whenever.” Compute the B’s as fast as possible, defer the W’s into the gaps, and the pipeline fills almost completely. “Such clever things that you can do with systems.”

Tensor parallelism and the activation-memory problem

Pipeline cuts depth; tensor parallelism cuts width, using the fact that a matmul decomposes into submatrices whose partial sums add back up — the same primitive as tiling. In the Megatron MLP, $A$ splits column-wise into $[A_1, A_2]$ and $B$ row-wise, each GPU computes $\text{GeLU}(XA_i)B_i$, and the results are summed. There is a duality to remember when implementing it: in the forward pass the input operator $f$ is the identity (copy $X$) and the output operator $g$ is an all-reduce; in the backward pass they flip. Across a transformer block: column-wise cuts on QKV and the MLP up-projection, row-wise cuts on the attention output and down-projection, and small layers — norms, non-linearities, MoE routers — fully replicated.

Tensor parallel is extremely communication hungry: an activation-sized all-reduce at every matmul, roughly $8bsh \cdot (n-1)/n$ per layer versus the pipeline’s $bsh$ point to point. So on GPUs you keep it inside a node — up to 8, on NVLink — and performance drops sharply cross-node. (TPUs, with their uniform mesh, can push tensor parallelism much wider; that’s the Google claim.) The pros: no bubble, no large-batch requirement, low complexity.

But memory isn’t just parameters. Profile a real training step and there’s a huge dynamic hump of activations, peaking just after the backward sweep begins; at scale, activations dwarf parameter memory. Storing everything costs, per transformer layer,

$\text{activation bytes} = sbh\left(34 + 5\,\frac{as}{h}\right),$

where the $sbh$ dependence is fundamental (something per sequence position, batch element, and hidden unit) and the $5as/h$ term is the quadratic attention piece, droppable via FlashAttention-style recomputation. Tensor parallelism divides the matmul parts — $sbh(10 + 24/t + 5as/(ht))$ — but the remaining $10\,sbh$ of layer norms, dropouts, and residual inputs is not split, no matter how many GPUs you throw at $t$.

Expert parallelism (and context parallelism)

Now that most big models are MoEs, expert parallelism — sharding whole experts across devices and routing token activations to them — joins the toolbox. Systems-wise it behaves like tensor parallelism: high-bandwidth, activation-reducing, best kept on fast interconnects. But when a layer is an MoE, prefer EP over TP: Megatron’s guidelines note that EP keeps local matmuls big (TP shreds them and utilization suffers), routes sparse token activations rather than dense sliced ones, and when EP equals the expert count, token permutation is free. It is still seriously hard: routing is all-to-all and latency-critical, which is why DeepSeek’s DeepEP went as far as undocumented PTX instructions to accelerate dispatch, and NVIDIA built Hybrid-EP for the same reason.

Two composition wrinkles. First, EP conventionally shards experts across the data-parallel replicas, so EP is bounded by DP. Second, MoEs change only the MLPs — so you’d like high TP for attention but low TP for the experts. Modern systems resolve the conflict by decoupling: attention layers get one parallelism (TP × CP × DP), MoE layers another (ETP × EP × EDP). Finally, context parallelism (ring attention) splits very long sequences across accelerators, passing KV blocks around a ring — the tool for long-context extension and serving; Tatsu skips the details as conceptually overlapping what’s already been covered.

4D parallelism: the recipe and the wild

No strategy dominates — every row of the recap table has something in red. FSDP is wonderful but doesn’t touch activations and consumes global batch; tensor parallel fixes activations but demands NVLink bandwidth; pipelines tolerate slow links but need micro-batches. The TPU-book analysis makes it quantitative: for a big enough per-chip batch, FSDP alone stays compute-bound; as batch shrinks it goes communication-bound, and mixing in model parallelism pushes the efficient frontier out to smaller batches. Hence 3D/4D parallelism. The prescription is genuinely simple:

The Narayanan 2021 Megatron sweeps bear it out: TP rises to 8 and caps; PP then grows with model size; DP shrinks last (down to 6 for the 1T model) — and utilization stays flat near 50% of peak from 32 to 3,072 GPUs. One counterintuitive finding: aggressive activation recomputation can raise throughput, because the memory it frees becomes batch size, and batch size buys utilization.

Real runs follow the recipe. OLMo 7B (the Dolma dataset): pure FSDP — small models scale surprisingly far on FSDP alone. DeepSeek V1: ZeRO-1 + tensor + sequence + pipeline parallel; V3, an MoE, swaps tensor for 64-way expert parallel (eight nodes grouped) with 16-way pipeline. Yi: the classic ZeRO-1 + TP + PP combo. Llama 3 405B, a giant dense model, publishes the full breakdown: TP 8 / CP 1 / PP 16 / DP 128 for main pretraining, then context parallel cranked to 16 (DP down to 8) for long-context extension — and 148 GPU failures during training, so redundancy is its own distributed-systems problem. Gemma 2 on TPUs: ZeRO-3 plus TP+SP over the big mesh, no pipeline — the Google philosophy realized. Mixtral 8x22B (per Megatron Bridge): EP 8, PP 4, TP 4 for the attention. Qwen 3: the DeepSeek-style recipe, EP 32 / PP 8 / TP 2. The common threads: maximize data parallel, keep TP at or below 8, and EP can now be big.

Check yourself

Napkin exercises in the lecture’s accounting style — bytes, bubbles, and bandwidth.

Problem 8.1 ZeRO memory accounting

A 30B-parameter model trains with the lecture’s 16-bytes-per-parameter recipe (bf16 params and gradients, fp32 master weights and both Adam moments) on 16 GPUs. Compute the per-GPU memory for parameters/gradients/optimizer state under naive DDP, ZeRO 1, ZeRO 2, and ZeRO 3. Which stages fit in 80 GB (ignoring activations)?

Show solution

Total state: $16 \times 30\times10^9 = 480$ GB. Per parameter: 2 (params) + 2 (grads) + 12 (master weights + two moments = optimizer state). DDP: all 480 GB replicated — no. ZeRO 1 shards the 12-byte optimizer state: $(2+2)\cdot30 + (12\cdot30)/16 = 120 + 22.5 = 142.5$ GB — no. ZeRO 2 also shards gradients: $2\cdot30 + (14\cdot30)/16 = 60 + 26.25 = 86.25$ GB — just misses. ZeRO 3 shards everything: $480/16 = 30$ GB — fits with room for activations. Only full FSDP works here.

Problem 8.2 pipeline bubble

You pipeline a model over 12 stages. How many micro-batches keep the bubble overhead at or below 5%? If each micro-batch must contain at least 4 sequences, what global batch size does that imply, and why does the lecture call batch size a “resource”?

Show solution

Bubble ratio $= (n_{\text{stages}}-1)/n_{\text{micro}} = 11/n_{\text{micro}} \le 0.05$, so $n_{\text{micro}} \ge 220$. At 4 sequences each, the global batch must be at least 880 sequences per pipeline — and data parallelism multiplies that further. Batch size is a resource because both DP (one example per worker minimum, diminishing returns past the critical batch size) and PP (micro-batches to hide the bubble) spend it, and the critical batch size caps how much you have to spend.

Problem 8.3 TP vs PP communication

Per transformer layer, tensor parallelism over $n$ devices moves about $8bsh \cdot (n-1)/n$ bytes in all-reduce traffic, while a pipeline stage boundary moves $bsh$ point to point. For $b=4$, $s=8192$, $h=8192$, bf16 activations, and $n=8$, compute both. Which link — NVLink inside a node or InfiniBand across nodes — should carry each, and why?

Show solution

$bsh = 4 \cdot 8192 \cdot 8192 \cdot 2$ bytes $\approx 0.54$ GB per boundary, point to point. TP: $8 \cdot 0.54 \cdot 7/8 \approx 3.76$ GB per layer, and it’s an all-reduce involving every device, repeated every layer. TP traffic is roughly 7x larger and collective and blocking (computation waits on it), so it belongs on the fastest interconnect — NVLink within one box, hence the TP-at-most-8 rule. The pipeline’s single smaller point-to-point transfer tolerates the slower cross-node links.

Problem 8.4 activation-memory lower bound

A model has $h = 12288$, $a = 96$ heads, $L = 96$ layers, trained at $s = 4096$ with micro-batch $b = 1$ and tensor+sequence parallelism at $t = 8$ with selective activation recomputation. Estimate total activation memory per GPU. What would dropping sequence parallelism cost?

Show solution

Lower bound per layer: $34\,sbh/t = 34 \cdot 4096 \cdot 1 \cdot 12288 / 8 \approx 214$ MB. Times $L = 96$: about 20.5 GB per GPU. Without sequence parallelism the formula is $sbh(10 + 24/t)$ (recomputation already dropped the attention term): $sbh \cdot 13 = 654$ MB per layer, about 62.8 GB total — the unsharded $10\,sbh$ of norms, dropouts, and residual inputs triples the bill. That’s exactly why sequence parallelism exists: the pointwise ops don’t shrink with $t$ until you split them along the sequence axis.