ZeRO stages 1, 2 and 3
Look again at what data parallelism stores. Every GPU holds the full 16 bytes for every weight, and the copies are identical. Eight GPUs therefore hold eight identical copies of 447 GiB of state, which is 3,576 GiB of memory to store 447 GiB of information.
Most of that is never needed in eight places. During the update, each GPU applies the same average to the same weights and produces the same answer, so seven of the eight are repeating work already being done. The idea is to give each GPU responsibility for one slice of the weights, let it keep only the state for that slice, and have it fetch anything else it needs when it needs it. Every GPU receives a different slice, so all eight hold the same amount of state and none of them sits idle.
That idea is called ZeRO, which stands for Zero Redundancy Optimizer. It was published by Microsoft in 2019. It describes an arrangement of storage rather than a change to the mathematics — the model learns exactly what it would have learned under plain data parallelism.
Redundancy is not the same as replication
Data parallelism replicates the weights for a reason: every card runs a full forward pass and needs them. What it also replicates, without a reason, is the optimizer state, which is touched once per step by an update that produces the same answer on every card. ZeRO removes the part that had no reason.
Three stages, each removing one more class of duplicate
Stage 1 splits the optimizer state, which is the 32-bit weight copy and the two running averages, twelve of the sixteen bytes. Each GPU keeps one slice and updates only that slice. The updated slices are then shared out so every GPU has the current weights again.
Stage 2 splits the gradients as well. A GPU only ever needs the gradients for the slice of weights it is responsible for updating, so the others are discarded as soon as they have been sent where they are needed.
Stage 3 splits the weights themselves. Each GPU stores one slice of the model. When the forward pass reaches a layer, the GPUs collect that layer's weights from each other, use them, and discard them again immediately.
That gather-use-discard cycle is what makes the memory saving real rather than notional. A card never holds the whole model, only its slice plus whichever single layer is currently in use. It is also where the extra communication comes from: the weights have to be fetched before every layer of the forward pass and again during the backward pass, which is the difference between ZeRO-3's 3P and the 2P of the earlier stages.
The limit this exposes
The peak is one assembled layer, so ZeRO-3 needs the largest single layer of the model to fit on one card. If a layer on its own exceeds the card, no amount of sharding state helps — the model has to be cut across cards rather than merely its state, which is tensor parallelism and a different session. This is the boundary where ZeRO stops being enough.
What each stage leaves on a card
Eight GPUs, eight weights, four bands. Switch the arrangement and watch which cells stay solid and which become dashed gaps — a dashed cell is state some other card is holding on this weight's behalf.
The bands from the first page, now with a divisor applied to whichever of them the stage has sharded. N is the number of GPUs; the figures are for eight.
| Bytes per weight, 8 GPUs | 30B model per GPU | Communication | |
|---|---|---|---|
| data parallelism | 16.00 | 447.0 GiB | 2P |
| ZeRO-1 | 5.50 | 153.7 GiB | 2P |
| ZeRO-2 | 3.75 | 104.8 GiB | 2P |
| ZeRO-3 | 2.00 | 55.9 GiB | 3P |
Stages 1 and 2 return ten of the sixteen bytes at the communication volume data parallelism was already paying. The reason is the equivalence from the collectives page. Data parallelism performs an all-reduce, which is internally a reduce-scatter and an all-gather. Stages 1 and 2 perform those same two phases and simply keep the intermediate slice instead of discarding it. Stage 3 adds a further all-gather of the weights in the forward pass and again in the backward pass, which raises the total from 2P to 3P.