ZeRO Stage 2 takes the idea behind Stage 1 — stop making every GPU store a full copy of everything — and pushes it one step further. Stage 1 partitioned the fat optimizer states across data-parallel ranks while every rank still held a complete gradient buffer. Stage 2 partitions the gradients too. The trick that makes this free is a change of collective: instead of all-reducing gradients so everyone ends up with the full averaged gradient, each rank keeps only the gradient shard it will actually use, produced by a reduce-scatter. Gradient memory falls from a full 2 bytes per parameter to 2/N, and — the elegant part — the total bytes moved across the network do not go up at all. This piece works through the memory arithmetic, the reduce-scatter reasoning, a concrete example, and exactly where Stage 2 sits between Stage 1 and Stage 3.
The memory model we are optimizing
Fix the accounting first, because ZeRO is entirely a story about bytes per parameter. Train a model of Ψ parameters with mixed-precision Adam and a standard optimizer keeps, per parameter: an fp16 weight (2 bytes), an fp16 gradient (2 bytes), and a bundle of fp32 optimizer states — a master fp32 copy of the weights (4), the Adam momentum (4), and the variance (4), so K = 12 bytes.
per-param bytes = 2 (fp16 weight)
+ 2 (fp16 gradient)
+ K (fp32 optimizer states), K = 12
= 16 bytes/param → total = 16ΨPlain distributed data parallel (DDP) replicates all 16Ψ bytes on every one of the N GPUs. That is the waste ZeRO attacks: the replicated state is identical across ranks, so storing N copies buys nothing. The question each ZeRO stage answers is simply which of those three buckets we are allowed to slice into 1/N-sized shards.
Where Stage 1 stopped
Stage 1 (P_os, partition optimizer states) shards only the K = 12 bytes. Each rank owns 1/N of the parameters’ optimizer states and is responsible for updating exactly that slice. But every rank still stores a full fp16 gradient buffer of 2Ψ bytes, because a normal all-reduce hands every rank the complete averaged gradient before the optimizer step.
Stage 1 memory/GPU = 2Ψ (weights)
+ 2Ψ (full gradients)
+ 12Ψ/N (sharded optimizer states)
= 4Ψ + 12Ψ/NNotice the tension: a rank keeps the whole gradient but only ever uses 1/N of it — the slice matching its optimizer-state shard. The other (N-1)/N of the gradient is dead weight the instant the update runs. Stage 2 is the observation that we never needed to materialize it.
What Stage 2 adds: partition the gradients
Stage 2 (P_os+g) keeps everything Stage 1 does and additionally partitions the gradients. After the backward pass, rank i ends up holding only the reduced gradient for the parameters it owns — a 2Ψ/N-byte shard — not the full 2Ψ buffer.
Stage 2 memory/GPU = 2Ψ (weights, still replicated)
+ 2Ψ/N (sharded gradients)
+ 12Ψ/N (sharded optimizer states)
= 2Ψ + 14Ψ/NGradient memory has dropped from 2 bytes/param to 2/N bytes/param. As the world size grows the two /N terms vanish and each GPU asymptotes to 2Ψ — just the fp16 weights, which Stage 2 deliberately leaves replicated. That last replicated 2Ψ is the floor Stage 2 cannot cross, and it is precisely the wall Stage 3 knocks down.
Reduce-scatter, not all-reduce
The enabling move is swapping the collective. In DDP every micro-batch computes a local gradient and an all-reduce sums them and broadcasts the result, so all N ranks come away with the identical, full averaged gradient. That is exactly what forces each rank to own 2Ψ bytes.
A reduce-scatter does the summation but not the broadcast. Conceptually, cut the gradient into N equal chunks; reduce-scatter sums chunk i across all ranks and delivers the result only to rank i. Every rank contributes to reducing all chunks but walks away with just one fully-reduced shard — the shard whose optimizer states it also owns. Rank i then does its fp32 Adam update on that slice locally, with no other rank’s gradients needed. In practice this is overlapped with the backward pass: as each layer’s gradients finish, they are reduce-scattered immediately, and the off-shard portions can be freed instead of accumulated, which is what actually realizes the 2Ψ/N memory rather than merely renaming a full buffer.
Why the communication volume is unchanged
The counter-intuitive result is that Stage 2 sends no more bytes than DDP. The key identity: an all-reduce is a reduce-scatter followed by an all-gather. A bandwidth-optimal ring all-reduce of a Ψ-element buffer moves about 2Ψ of data per GPU — Ψ in the reduce-scatter half and Ψ in the all-gather half.
DDP: all-reduce(grads) ≈ 2Ψ moved/GPU
Stage 2: reduce-scatter(grads) ≈ Ψ (get my gradient shard)
+ all-gather(weights) ≈ Ψ (rebuild full fp16 weights)
≈ 2Ψ moved/GPUStage 2 uses only the first half of an all-reduce to collect gradients, then — because each rank updated a different slice of the weights — needs a matching all-gather after the optimizer step to reassemble the full fp16 weight tensor for the next forward pass. Reduce-scatter plus all-gather equals one all-reduce’s worth of traffic: 1x the DDP volume. Stage 2 saves gradient memory for free in communication terms.
A worked example: 7.5B parameters, 64 GPUs
Take Ψ = 7.5 billion parameters and N = 64 data-parallel GPUs, counting 1 GB = 10^9 bytes for round numbers.
Baseline DDP : 16Ψ = 120 GB per GPU (will not fit 32 GB)
Stage 1 : 4Ψ + 12Ψ/64 = 30 + 1.41 = 31.4 GB
Stage 2 : 2Ψ + 14Ψ/64 = 15 + 1.64 = 16.6 GB
Stage 3 : 16Ψ/64 = 1.9 GBRead the columns. DDP is hopeless — 120 GB on a 32 GB card. Stage 1 already fits, dominated by the 4Ψ = 30 GB of replicated fp16 weights and the still-full fp16 gradients. Stage 2 halves that replicated block to 2Ψ = 15 GB by sharding the gradients, landing at 16.6 GB — comfortable headroom for activations. The whole difference between the Stage 1 and Stage 2 columns is one term: 2Ψ of gradient shrinking to 2Ψ/64, worth about 14.8 GB per GPU here, bought without moving a single extra byte over the network.
Stage 1 vs Stage 2 vs Stage 3 at a glance
The three ZeRO stages are a strict ladder — each shards one more of the three memory buckets, trading a little more communication and complexity for lower per-GPU memory.
| Stage | Sharded | Replicated | Memory/GPU | Comm. |
|---|---|---|---|---|
| DDP | nothing | weights, grads, opt | 16Ψ | 1x |
| 1 (P_os) | optimizer states | weights, grads | 4Ψ + 12Ψ/N | 1x |
| 2 (P_os+g) | opt + gradients | weights | 2Ψ + 14Ψ/N | 1x |
| 3 (P_os+g+p) | opt + grads + weights | nothing | 16Ψ/N | 1.5x |
The jump that matters: Stages 1 and 2 both cost the same 1x communication as ordinary DDP, so Stage 2 is nearly always worth enabling over Stage 1 — it is pure memory savings. Stage 3 additionally shards the fp16 weights, so every layer must be all-gathered on the fly during both forward and backward passes, pushing communication to 1.5x. That is the real dividing line: Stage 2 is free bandwidth-wise; Stage 3 is not.
Why Stage 2 leaves weights alone
Stage 2 stops one bucket short of Stage 3 on purpose. The fp16 weights stay replicated because every rank needs the whole weight tensor to run its forward and backward pass on its own micro-batch — that is the essence of data parallelism. Sharding the weights (Stage 3) means a rank no longer has the layers it is about to compute, so it must all-gather them just-in-time, use them, and discard them, layer by layer. That is powerful but adds the extra 0.5x of traffic and real scheduling complexity.
Stage 2 therefore hits a sweet spot: it captures the two cheapest-to-shard buckets — optimizer states and gradients, together 14 of the 16 bytes/param — while keeping the communication pattern and the mental model almost identical to DDP. For many training runs the replicated 2Ψ of weights is a perfectly affordable residue, and reaching for Stage 3 is only justified once that last block genuinely does not fit.
Practical notes and pitfalls
A few things bite in practice. First, the 2Ψ/N gradient figure assumes the reduce-scatter output actually replaces the full buffer; a naive implementation that all-reduces and then slices saves no memory at all, so the gradient bucketing and immediate freeing during backward is load-bearing, not an optimization detail. Second, gradient accumulation interacts with sharding: across accumulation micro-steps you must accumulate into the shard, reduce-scattering only on the boundary, or you lose the memory win. Third, gradient clipping by global norm now needs an extra small all-reduce of the per-shard squared norms, since no single rank sees the whole gradient — cheap, but easy to forget.
Finally, keep the model honest against the arithmetic: if measured memory does not track 2Ψ + 14Ψ/N, the usual culprits are activations (orthogonal to ZeRO — use checkpointing), fragmentation, or a framework silently falling back to a full gradient buffer. The formula is the yardstick.
2 to 2/N bytes per parameter and per-GPU memory drops from 4Ψ + 12Ψ/N to 2Ψ + 14Ψ/N. The beautiful part is that it is free on the wire — reduce-scatter plus the post-update weight all-gather together move exactly as many bytes as one DDP all-reduce, keeping communication at 1x. That is why Stage 2 is almost always preferable to Stage 1. It stops short of sharding the still-replicated fp16 weights, the 2Ψ floor that only Stage 3 removes — at the cost of an extra half of communication volume.