Part II — Systems

Data parallelism — DDP → ZeRO → FSDP

Lesson 07 made one GPU as fast as it can be: kernels fused, FlashAttention killing the O(T²) attention activations, the chip running near its roofline. But a fast GPU is still one GPU, and lesson 5 already proved the verdict — a 7B model's training state is ≈112 GB, and an H100 holds 80. No amount of kernel cleverness fits 112 into 80. The only move left is more devices. This lesson takes the obvious first step (replicate the model and split the batch), shows it scales speed but not memory, and then shows the idea that actually breaks the memory wall: stop replicating the redundant state and shard it.

The previous step left this broken

The memory ledger from lesson 5 is unmoved by anything in lesson 6 or 07. Fusion and FlashAttention shrink the activation term; they do nothing to the 112 GB of training state — 2 bytes of bf16 weights, 4 of fp32 master, 4+4 of Adam m,v, 2 of grads, ≈16 bytes/param. That state still overflows one 80 GB H100 before a single activation is stored. We can now make a GPU efficient; we still cannot make the model fit on it. The work must cross device boundaries — and the cheapest way to cross a boundary, the one every framework reaches for first, is to put a whole copy of the model on each GPU and feed each copy a different slice of the batch.

Linear position

Forced by: a model's training state (≈16N bytes) exceeds one GPU's memory, and lesson 7 has already spent every single-device trick.
New idea (two steps): DDP — replicate the model on G GPUs, split the global batch G ways, run local forward/backward, then all-reduce the gradients so every replica steps identically (this scales throughput, not memory); then ZeRO / FSDP — notice params, grads, and optimizer state are identical across replicas, so shard each of the three across the G ranks and gather just-in-time (all-gather in forward, reduce-scatter in backward), cutting per-GPU state ≈G×.
Forces next: even fully sharded, a single layer too big to compute on one device, or a sequence too long to hold its activations, still won't fit — and DP can never shrink the per-step compute below one whole model; you must split the model itself (lesson 9).

The plan

Five moves. (1) Build DDP: replicate, scatter the batch, local fwd/bwd, all-reduce gradients — and price the all-reduce at ≈2(G−1)/G·N bytes, overlapped with backward. (2) Name the limitation: every rank still holds the full ≈16N state, so memory is unchanged. (3) The ZeRO insight: that state is redundant across replicas → shard it. ZeRO-1 (optimizer) / ZeRO-2 (+grads) / ZeRO-3 = FSDP (+params, gathered per layer then freed). (4) Plot per-GPU memory ≈16N/G and price ZeRO-3's extra comm (~1.5× DDP). (5) Drive the meter: set ZeRO-3 and watch a model DDP can't hold drop under 80 GB.

1 · DDP: replicate the model, split the batch

The first multi-GPU idea is the one that needs the least new machinery, because it leaves the model untouched. Put a full, identical copy of the model on each of G GPUs. Take the global batch B and hand each GPU a disjoint slice of size B/G (the "data" the parallelism is named for). Every GPU runs its own forward and backward on its own slice, producing its own gradient. The four steps of one training step:

replicateat startup, broadcast one copy of the weights to all G ranks so every GPU holds identical parameters.

scatter batchsplit the global batch B into G micro-batches of B/G; rank i sees only its slice.

local fwd / bwdeach rank computes loss and gradients on its slice — no communication yet, pure local compute.

all-reduce gradsaverage the gradients across all ranks so every rank gets the same global-batch gradient, then every rank applies the identical AdamW step.

The all-reduce is the whole trick and the whole cost. Mathematically, a sum over the global batch equals the sum of the per-rank sums, so averaging each rank's local gradient gives exactly the gradient you would have computed on the full batch on one (impossibly large) GPU. Because every rank ends the step with the same averaged gradient and started with the same weights, they apply the same update and stay bit-identical replicas forever — no parameter sync is ever needed, only the gradient sync.

What the all-reduce costs

A naive all-reduce — every rank sends its full gradient to one rank to be summed and broadcast back — costs O(G·N) and bottlenecks on one link. The algorithm everyone actually uses is the ring all-reduce: arrange the G ranks in a ring, split each gradient into G chunks, and run G−1 reduce-scatter steps (each rank accumulates one chunk) followed by G−1 all-gather steps (each rank receives the final value of every chunk). Each GPU sends and receives

all-reduce bytes / GPU ≈ 2 · (G−1)/G · N · (bytes per grad element)

The key feature: the (G−1)/G factor tends to 1 as G grows, so the per-GPU communication is ≈2N bytes regardless of how many GPUs you add. Doubling the cluster does not double each GPU's comm — that bandwidth-independence of G is why ring all-reduce scales. For a 7B model in bf16 (2 bytes/grad), per-GPU traffic is ≈ 2 · 7×10⁹ · 2 ≈ 28 GB per step, moved over NVLink/InfiniBand.

Overlap: hide the comm behind the backward

You do not wait for the backward to finish and then all-reduce — that would serialize comm after compute and tank MFU. Gradients become ready layer-by-layer, from the top of the network down, as the backward sweeps. DDP fires the all-reduce for a bucket of gradients the instant that bucket is complete, while the backward keeps computing the lower layers. Done right, almost all of the ≈28 GB of comm overlaps the backward compute and is nearly free in wall-clock. This is the same recurring move from lesson 5 — spend the resource you have spare (here, the backward's compute time) to hide the one you're short on (interconnect latency).

What DDP actually buys

DDP multiplies throughput: G GPUs chew through G× the tokens per second (minus a few % for un-overlapped comm), so a fixed D tokens finishes ≈G× faster — exactly the time ≈ 6ND/(G·peak·MFU) speedup from lesson 5. What it does not buy is memory. Read on.

2 · The limitation: every rank still holds the whole model

Look back at the four DDP steps and ask what each GPU has resident. A full copy of the parameters. A full set of gradients. A full optimizer state (fp32 master, Adam m, Adam v). That is the entire ≈16N ledger from lesson 5, sitting on every single rank. DDP replicated the model; it replicated the memory problem along with it.

DDP per-GPU state = 16 · N (unchanged — independent of G)

So a 7B model is still ≈112 GB on each of your 64 GPUs, and 112 > 80 on every one of them. DDP can make a model that already fits train G× faster, but it cannot make a model that doesn't fit fit at all. Splitting the batch was the wrong axis to split on: the batch was never the thing overflowing the device — the state was, and DDP keeps a whole copy of it per GPU. We need to split the state.

3 · The ZeRO insight: that state is redundant — shard it

Here is the observation that breaks the wall, from the ZeRO paper (Rajbhandari et al., 2019). Across all G DDP replicas, the parameters are identical. The optimizer states are identical. After the all-reduce, even the gradients are identical. We are storing the same ≈16N bytes G times over. That redundancy is pure waste — and the fix is to keep only 1/G of each on each rank, and fetch the missing pieces over the (fast, otherwise-idle) interconnect exactly when a given step needs them. ZeRO ("Zero Redundancy Optimizer") shards the three pieces in three increasingly aggressive stages:

ZeRO-1 — shard optimizer state

Each rank keeps only 1/G of the fp32 master + Adam m,v (the 12 fp32 bytes). Params and grads still replicated. The optimizer is the biggest piece (12 of 16 bytes), so this alone is the largest single drop. Comm ≈ unchanged from DDP.

ZeRO-2 — also shard gradients

On top of ZeRO-1, each rank keeps only the gradient slice for the params it optimizes. The all-reduce splits into a reduce-scatter of the grads (each rank ends up with just its averaged slice) plus an all-gather of the updated params — the same 2N total as DDP's all-reduce, no extra comm. Params still replicated.

ZeRO-3 = FSDP — also shard parameters

The full step: params too are sharded 1/G. A layer's full weights exist on no single rank, so before you can compute a layer you must all-gather its shards, use them, then free them. This is what PyTorch ships as Fully Sharded Data Parallel. It costs extra comm (§4) but cuts per-GPU state to ≈16N/G.

How FSDP runs a forward and backward

The mechanics of ZeRO-3 are the part worth holding precisely, because the comm pattern follows directly. Each rank permanently owns only its 1/G shard of every layer's parameters. To execute the model:

forward, per layerall-gather this layer's parameter shards from all ranks → every rank briefly holds the layer's full weights → compute the layer's output → immediately free the gathered weights, keeping only your shard. Peak param memory is one layer's worth, not the whole model's.

backward, per layerall-gather the layer's weights again (they were freed) → compute gradients → reduce-scatter the gradients so each rank receives only the averaged slice for the params it owns → free the gathered weights.

optimizer stepeach rank updates only its 1/G of params using its 1/G of optimizer state and its reduce-scattered gradient slice. No rank ever materializes the full model state.

So the model is sharded at rest and gathered just-in-time, one layer at a time, and the gathered copy is thrown away the moment that layer is done. The all-gather in forward and the reduce-scatter in backward are the new communications ZeRO-3 adds over DDP. ZeRO-1/2 keep params replicated, so they never gather params: each rank reduce-scatters the grads to its slice, updates that slice, then all-gathers the updated params back — reduce-scatter + all-gather = the same 2N total as DDP's all-reduce. That is why they add memory savings at essentially the same comm as DDP.

4 · Per-GPU memory ≈ 16N/G — and the cost of ZeRO-3

The payoff is a clean curve. Sharding the parts each stage owns divides that part's bytes by G; the unsharded parts stay full. Per parameter (the 2/2/12 bf16 ledger from lesson 5), per GPU:

stage	params (2)	grads (2)	optim (12)	bytes/param @ G	per-GPU at large G
DDP (ZeRO-0)	2	2	12	16 (constant)	16N — never shrinks
ZeRO-1	2	2	12/G	4 + 12/G	→ 4N
ZeRO-2	2	2/G	12/G	2 + 14/G	→ 2N
ZeRO-3 (FSDP)	2/G	2/G	12/G	16/G	→ 0 (≈16N/G)

Read the limits. ZeRO-1 floors at 4N bytes/GPU (the still-replicated params+grads); ZeRO-2 floors at 2N (the still-replicated params); only ZeRO-3 keeps falling as ≈16N/G with no floor, because all three pieces shard. Worked on the 7B: DDP is 112 GB/GPU on any cluster size; ZeRO-3 on 16 GPUs is ≈ 112/16 = 7 GB/GPU of state — trivially under 80 GB, with the rest of the budget free for activations. That is the model DDP can't fit, fitting.

Nothing is free. ZeRO-3 trades that memory for communication volume:

scheme	comm pattern / step	volume vs DDP
DDP / ZeRO-1 / ZeRO-2	gradient all-reduce (or reduce-scatter) ≈ 2N	1.0× (baseline)
ZeRO-3 / FSDP	all-gather params (fwd) + all-gather params (bwd) + reduce-scatter grads	≈ 1.5× DDP

The intuition: DDP moves the gradient once (≈2N). FSDP additionally moves the parameters twice (once to gather for forward, once to re-gather for backward), and a reduce-scatter is half an all-reduce — adding up to roughly 1.5× the DDP communication volume. As long as that extra comm overlaps compute (gather layer i+1 while computing layer i, the same prefetch trick as DDP's bucketed overlap), the wall-clock hit is small — but it is real, and it is why you do not jump to ZeRO-3 when a lower stage already fits.

5 · Practical: which stage, and the overlap that makes it cheap

The decision rule is a memory-vs-comm trade, read straight off the table above. Pick the least aggressive stage that fits, because each stage up adds communication:

params replicated (DDP / ZeRO-1/2)

Use when the model's params already fit per GPU and you only need to claw back optimizer/grad memory. Lowest comm; the params never move. Common for ≤~few-B dense models on 80 GB cards.

params sharded (ZeRO-3 / FSDP)

Required when the params themselves don't fit per GPU (7B+ on one node, or you want huge per-GPU activation/batch headroom). Pays ~1.5× comm; the default for training large dense models on a DP-only mesh.

+ mixed precision (always)

FSDP keeps a sharded fp32 master and runs bf16 compute (lesson 4). Gather/compute in bf16 → half the all-gather bytes vs fp32; the fp32 master is the 4 of the 12 sharded optimizer bytes. The two ideas compose directly.

Two practical levers decide whether ZeRO-3's extra comm actually hurts. First, overlap: prefetch the next layer's all-gather during the current layer's compute, so communication hides behind the GEMMs instead of stalling them — exactly DDP's bucketed-overlap idea, applied per layer. Second, topology: the all-gathers are bandwidth-hungry, so on a multi-node cluster you keep the heavy collectives on fast intra-node NVLink where you can and let the slower inter-node InfiniBand carry the less frequent traffic — which is precisely the composition with tensor/pipeline parallelism that lesson 9 builds, and the cluster-topology reasoning post-training also relies on (see the reinforcement_learning track's topology lessons, and the orchestration view in ray). The fused optimizer step that consumes the sharded gradient and master copy — and its peak-memory accounting — is the subject of gpu_kernels · 30 (the training step & peak memory).

6 · Drive the meter

Set the model size and the data-parallel degree G, then sweep the ZeRO stage 0→3. The stacked meter shows per-GPU bytes split into sharded params / grads / optimizer against the 80 GB H100 line; when the bar would cross it, the bar and the "fits?" KPI turn red. Start at 7B, ZeRO-0 (DDP), G=8 — and watch it overflow no matter how high you push G, because DDP's 16 bytes/param is constant in G. Then step to ZeRO-3 and watch the same 7B drop to ≈16N/G and slide under the line. Prove three things: (a) ZeRO-1 only shrinks the optimizer segment; (b) ZeRO-0's total never moves when you change G, while ZeRO-3's falls as 1/G; (c) the comm-per-step KPI ticks up ~1.5× the moment you reach ZeRO-3.

ZeRO stage memory — what each stage shards, and what it costs

Per-GPU bytes vs the 80 GB H100 wall. The hinge: ZeRO-0 (DDP) holds the full ≈16N state on every GPU regardless of G, so 7B is red on a cluster of any size. ZeRO-1 shards only the optimizer (12→12/G), ZeRO-2 also grads, ZeRO-3/FSDP also params → per-GPU state ≈16N/G with no floor. Watch ZeRO-3 fit a model DDP can't, and watch comm/step jump to ~1.5× DDP when params start being all-gathered.

N (params): #GPUs (G): 8 ZeRO stage:

params (2 B) grads (2 B) optimizer · master+m+v (12 B)

80 GB · 1× H100

State GB / GPU

—

Fits 80 GB?

—

Comm / GPU / step

—

bytes / param

—

The widget is the lesson in one screen: DDP's bar is pinned at 16 bytes/param and goes red on 7B for every G; only when params themselves shard (ZeRO-3) does the bar fall as 1/G and clear the wall — at the price of the comm-per-step KPI rising to ~1.5× DDP.

Failure modes & checklist

Failure modes

Expecting DDP to save memory. Adding GPUs and still OOMing at the same per-GPU number. Signal: the OOM byte count is independent of G — DDP replicates the full 16N state; you needed to shard (ZeRO), not replicate.
Jumping straight to ZeRO-3. Using FSDP when the params already fit per GPU. Signal: MFU dips and step time rises vs a lower stage for no memory reason — you paid ~1.5× comm to shard params you didn't need to.
Not overlapping the collectives. All-reduce/all-gather running after compute instead of behind it. Signal: a comm-shaped gap in every step's timeline and MFU well under ~40% — buckets/prefetch aren't firing.
Forgetting comm scales with N, not batch. Assuming bigger batch makes the all-reduce relatively cheaper. Signal: comm-bound at small batch because per-GPU gradient bytes are ≈2N regardless of batch size.
Sharding across a slow link. Running FSDP's all-gathers over inter-node InfiniBand when they should stay on NVLink. Signal: step time dominated by parameter gather; collectives are crossing node boundaries (a lesson-09 mesh problem).

Checklist

Diagnose the wall first. If params fit but optimizer/grads don't, use ZeRO-1/2; if params don't fit, use ZeRO-3/FSDP.
Pick the lowest stage that fits — each stage up adds communication; don't pay for sharding you don't need.
Verify the comm overlaps compute (bucketed all-reduce, prefetched all-gather) before blaming the network.
Price the all-reduce as ≈2(G−1)/G·N bytes per GPU, and ZeRO-3 at ~1.5× that — both ≈ independent of G.
Keep heavy collectives intra-node (NVLink) and compose mixed precision with FSDP to halve gather bytes.

Checkpoint

Try it

Open the meter at 7B, ZeRO-0 (DDP), G=8. Confirm the bar is red at ≈112 GB and "Fits?" says no. Now (a) drag G from 8 to 64 and confirm the total does not move — explain in one sentence why DDP's per-GPU memory is independent of G. (b) Switch to ZeRO-1: which single segment shrank, and why is it the biggest first win? (c) Switch to ZeRO-3 and confirm 7B now fits; read off the state GB/GPU and check it against 112/8 ≈ 14 GB. (d) Note the comm/step KPI before and after step (c): by what factor did it rise, and what new collective caused it? Finally, by hand: a 70B model in bf16 AdamW — how many GPUs of ZeRO-3 sharding (state only) before per-GPU state drops under 80 GB?

Where this points next

FSDP just dissolved the memory wall: shard params, grads, and optimizer across G ranks and per-GPU state falls as ≈16N/G with no floor, so a 7B — or, with enough ranks, a 70B — fits. But notice what is still true on every rank. Each GPU still executes the whole model on its batch shard; it just gathers the weights one layer at a time to do so. That leaves two things data parallelism structurally cannot fix. First, a single layer whose weights or activations are too big to compute on one device even momentarily — gather it and you OOM on the gathered copy. Second, a sequence so long that one layer's activations won't fit. And more fundamentally, DP can never reduce the per-step compute below one full model forward+backward — it only splits the batch. The forced next move is to stop sharding only the storage and start splitting the computation: cut each matmul across GPUs (tensor parallelism), cut the layer stack into stages (pipeline parallelism), cut the sequence (context parallelism), and compose all of it with DP into a 3D mesh. Next: 09 · Model parallelism — tensor, pipeline, the 3D mesh.

Takeaway

Lesson 05 left a 7B model's ≈112 GB of training state overflowing an 80 GB GPU; lesson 7 made the GPU fast but not bigger. DDP is the first repair: replicate the model on G GPUs, split the batch, run local fwd/bwd, and all-reduce the gradients (ring all-reduce, ≈2(G−1)/G·N bytes/GPU — nearly independent of G, overlapped behind the backward). That multiplies throughput by G but leaves per-GPU memory at the full 16N, because every rank still holds a complete copy of params, grads, and optimizer state — DDP replicated the memory problem. The ZeRO insight is that those three are identical across replicas, hence redundant: shard them. ZeRO-1 shards the optimizer (the biggest piece, 12 of 16 bytes), ZeRO-2 also gradients, ZeRO-3 = FSDP also parameters — all-gathered per layer in forward, freed, re-gathered and reduce-scattered in backward. Per-GPU state then falls to ≈16N/G, the only stage with no floor, at the cost of ≈1.5× DDP communication (params moved twice plus a reduce-scatter), which overlap and intra-node topology keep cheap. That is how a model DDP cannot hold fits. But every rank still runs the whole model on its batch shard, so a single oversized layer, a too-long sequence, or the irreducible one-model-per-step compute still breaks DP — forcing model parallelism in lesson 9.

Interview prompts

What does DDP scale, and what does it leave untouched? (§1–2 — it scales throughput ≈G× by splitting the batch and all-reducing gradients; it leaves per-GPU memory at the full 16N because every rank holds a complete copy of params, grads, and optimizer state.)
Derive the per-GPU cost of a ring all-reduce. (§1 — split the gradient into G chunks; G−1 reduce-scatter + G−1 all-gather steps, each GPU sends/receives ≈2(G−1)/G·N elements → ≈2N bytes, nearly independent of G, which is why it scales.)
How is comm overlapped with the backward in DDP? (§1 — gradients become ready top-down; DDP fires the all-reduce per bucket the moment that bucket completes, while the backward keeps computing lower layers, so most comm hides behind compute.)
What is the ZeRO insight, and what do stages 1/2/3 each shard? (§3 — params/grads/optimizer are identical across replicas, hence redundant; ZeRO-1 shards optimizer state, ZeRO-2 also gradients, ZeRO-3/FSDP also parameters (all-gathered per layer, then freed).)
Walk an FSDP forward and backward. (§3 — forward: all-gather a layer's param shards, compute, free; backward: re-all-gather, compute grads, reduce-scatter grads to owners, free; optimizer updates only the local 1/G shard.)
Why does only ZeRO-3 keep shrinking with G? (§4 — ZeRO-1 floors at 4N (replicated params+grads), ZeRO-2 at 2N (replicated params); only ZeRO-3 shards all three, so per-GPU state is ≈16N/G with no floor.)
What extra communication does ZeRO-3 add over DDP, and why ~1.5×? (§4 — DDP moves the gradient once (≈2N); FSDP also all-gathers params twice (fwd + bwd) and reduce-scatters grads (half an all-reduce), summing to ≈1.5× DDP volume; overlap keeps the wall-clock hit small.)
When would you stay on ZeRO-1/2 instead of ZeRO-3? (§5 — when the parameters already fit per GPU; you only need to reclaim optimizer/grad memory, so you avoid ZeRO-3's ~1.5× comm by keeping params replicated.)