Why architecture matters here
The architecture matters because the alternatives partition along different axes and each has a different cost profile. Tensor parallelism splits individual matrices across devices, which works but demands an all-reduce inside every layer's forward and backward pass — extremely bandwidth-hungry, viable essentially only within a single NVLink node. Pipeline parallelism splits layers across devices, which communicates little but introduces bubbles and forces you to restructure the model into stages. ZeRO keeps the data-parallel programming model completely intact — you still write a normal model and a normal training loop — and buys memory by removing redundancy rather than by splitting math. That preservation of the programming model is why it is usually the first thing to reach for.
It matters because the three stages are genuinely different bargains, not merely increasing intensities of one knob. Stage 1 removes the optimizer state — twelve of the sixteen bytes, the largest single chunk — for essentially no additional communication over standard data parallelism, because the gradient all-reduce that data parallelism already performs can be restructured as a reduce-scatter plus all-gather with the same total volume. Stage 1 is close to free, and the fact that many jobs still run without it is a straightforward waste of most of their memory.
Stage 2 adds gradient sharding, which is still nearly free in communication terms because the reduce-scatter is already happening — you simply stop all-gathering the full gradient afterward, since each rank only updates its own shard. Stage 3 is where the bargain changes qualitatively: sharding the parameters themselves means a rank does not hold the weights it needs to compute a layer, so it must all-gather them before every forward and again before every backward. That is roughly a fifty percent increase in communication volume, and it is the point at which the interconnect becomes the deciding factor.
The interconnect dependency is the architectural fact that most determines whether ZeRO works for you. Stage 3 across NVLink inside a node is comfortable; the same configuration across a slow Ethernet fabric turns a compute-bound job into a communication-bound one, and GPUs sit idle waiting for weights. This is why the meaningful question is never 'which stage saves the most memory' — stage 3 always does — but 'which is the least aggressive stage that fits', because every stage beyond what you need is throughput you paid for and did not have to.
Finally, the architecture matters because overlap is what makes it viable at all, and overlap is a property of the implementation rather than of the algorithm. The all-gather for layer k+1 can be issued while layer k is still computing, hiding the transfer under compute. When prefetching works, stage 3's extra communication largely disappears into the shadow of the math. When it does not — because the bucket size is wrong, the layers are too small to cover the transfer, or a synchronization point forces a stall — the same configuration exposes every byte of it. The difference between a well-tuned and badly-tuned ZeRO-3 job is routinely a factor of two in throughput on identical hardware.
The architecture: every piece explained
Start with the state taxonomy, because ZeRO is organized around it precisely. Persistent training state divides into three categories with very different sizes: parameters (2 bytes each in fp16), gradients (2 bytes each), and optimizer state (12 bytes each for Adam — the fp32 master weight, first moment, second moment). Activations are a fourth category and are explicitly not what ZeRO addresses; that is activation checkpointing's job, and conflating the two leads people to expect memory relief that ZeRO structurally does not provide.
ZeRO-1 partitions only the optimizer state. Each rank owns 1/N of the parameters for update purposes and holds Adam's moments only for its slice. Every rank still holds all parameters and computes all gradients. After the backward pass, a reduce-scatter delivers each rank the fully reduced gradients for exactly its slice; it updates its slice; then an all-gather redistributes the updated fp16 parameters. Total bytes on the wire equal a standard all-reduce, which is why stage 1 costs nothing and removes twelve of sixteen bytes.
ZeRO-2 additionally partitions the gradients. Rather than materializing a full gradient buffer and reducing it, gradients are reduce-scattered as they are produced during the backward pass, bucket by bucket, and each rank retains only its shard. This removes two more bytes per parameter and — subtly but importantly — removes a full-size gradient buffer that was a peak-memory contributor. Communication volume is unchanged from stage 1 for the same reason: the reduce-scatter was always going to happen.
ZeRO-3 partitions the parameters themselves, and the execution model changes. A rank holds only 1/N of each layer's weights at rest. Before computing a layer's forward, the ranks all-gather that layer's full parameters into a temporary buffer, compute, and then immediately free the buffer. The same all-gather happens again during the backward pass, because the weights are needed to compute gradients and were discarded. Two extra all-gathers of the parameters per step is the cost; near-linear memory scaling with N is the reward. Only one layer's worth of full parameters is live at any moment, which is what allows models far larger than a device to train.
Offload extends the hierarchy beyond GPU memory. ZeRO-Offload moves optimizer state and the parameter update itself to CPU RAM, which works because the update is memory-bound rather than compute-bound and a CPU can do it acceptably while the GPU handles the math. ZeRO-Infinity extends further to NVMe, making the total state limited by disk rather than by aggregate HBM. Both are escape hatches with a real cost: the PCIe link becomes the bottleneck, and a job that offloads to keep from OOMing may train several times slower than one that fits. Offload buys feasibility, not speed, and that distinction should drive the decision.
End-to-end flow
Walk a ZeRO-3 training step. At rest, each of eight ranks holds one-eighth of every layer's parameters, one-eighth of the gradients, and one-eighth of the Adam state. No rank can compute anything yet — it does not have a complete layer. The step begins when rank-local data loading hands each rank its own distinct micro-batch, exactly as in ordinary data parallelism.
The forward pass starts at layer one. The framework issues an all-gather for layer one's parameters across all eight ranks, and every rank now holds that layer's complete weights in a temporary buffer. Each rank computes layer one's forward on its own micro-batch. The moment the layer is done, the temporary buffer is freed — the rank goes back to holding only its shard. Critically, while layer one is computing, the framework has already issued the all-gather for layer two, so the transfer proceeds concurrently with the math. This prefetch is the whole reason stage 3 is usable.
This repeats layer by layer to the loss. Peak parameter memory is not the whole model — it is one rank's shard of everything plus the full weights of however many layers are gathered at once (the current one and the prefetched next). That is the structural insight: the transient buffer, not the model size, sets the parameter memory floor, which is why a 70B model can train on devices that could never hold it.
The backward pass reverses, and pays the parameter cost again. Layer N's weights were freed after the forward, so they must be all-gathered a second time to compute gradients. As each layer's gradients are produced, they are immediately reduce-scattered: the gradients are summed across ranks and each rank keeps only its own slice. The full gradient for a layer never exists on any single device. The freed weights and the discarded gradient remainder both release memory as the pass walks backward.
The optimizer step is then embarrassingly local, and this is the elegant part. Each rank holds its parameter shard, its gradient shard, and its Adam state shard — all aligned to the same slice. It runs the Adam update on its slice with no communication at all. There is no optimizer all-reduce, because no rank needs to know anything about any other rank's slice. The updated fp16 parameters are simply the new resting shard, ready for the next step's all-gathers.
The synthesis: ZeRO turns the persistent memory cost into a transient one. Instead of holding all state all the time, each rank holds 1/N all the time and materializes full tensors only for the microseconds it needs them. Every performance question reduces to whether those materializations hide under compute. When prefetch depth, bucket size, and layer granularity line up, the GPUs stay fed and stage 3 costs little; when they do not, the same job is a sequence of stalls, and the profiler — not the memory counter — is where you find out which one you have.