ZeRO Stage 1 Math

The 16-bytes-per-parameter budget

Start with the honest accounting for one parameter under mixed-precision Adam. Let Ψ be the number of parameters in the model. During training each GPU in a standard DDP (DistributedDataParallel) setup holds five things per parameter:

fp16 weights   (forward/backward)  : 2 bytes
fp16 gradients                     : 2 bytes
fp32 master weights  (optimizer)   : 4 bytes
fp32 momentum  m     (optimizer)   : 4 bytes
fp32 variance  v     (optimizer)   : 4 bytes
-----------------------------------------------
total per parameter                : 16 bytes

So the per-GPU footprint of the model+optimizer is 16·Ψ bytes. It is conventional to split this into the model/gradient part (2 + 2 = 4 bytes) and the optimizer part, written K·Ψ with K = 12. That factor of 12 is not a detail — it is three-quarters of the whole budget. ZeRO Stage 1 is, in one sentence, the decision to stop replicating those 12 bytes on every rank.

Where the 12 bytes come from

Why does the optimizer need 12·Ψ bytes when the model itself lives happily in 2·Ψ bytes of fp16? Because Adam is a stateful optimizer that demands full precision to remain stable. Its update, for each parameter, is:

m_t = β1·m_(t-1) + (1-β1)·g_t          (momentum)
v_t = β2·v_(t-1) + (1-β2)·g_t^2        (variance)
θ_t = θ_(t-1) - α·m̂_t / (√v̂_t + ε)     (step)

The running estimates m and v must persist across steps and accumulate tiny contributions, so they are kept in fp32 (4 bytes each). The step itself would lose precision if applied to the fp16 weights directly — subtracting a small α·m̂/√v̂ from an fp16 number often rounds to a no-op — so training also keeps an fp32 master copy of the weights (another 4 bytes) and rounds it down to fp16 only for the forward pass. Master weights + momentum + variance = 4 + 4 + 4 = 12 bytes. That is the block Stage 1 shards.

The core idea: shard the states, not the compute

Here is the observation ZeRO is built on. In DDP, every rank runs the identical optimizer step on the identical full set of states, producing the identical updated weights. That is redundant work on redundant memory: N ranks each store all 12·Ψ bytes and each recompute the whole update. Nothing forces this.

ZeRO Stage 1 — sometimes written P_os, for ‘partition optimizer states’ — slices the parameter space into N disjoint shards and assigns shard i to rank i. Rank i stores the fp32 master weights, momentum, and variance only for its shard, and runs the Adam update only for those parameters. No rank owns the full optimizer state anymore; collectively the N ranks hold exactly one copy, spread out. Crucially, the forward and backward passes are untouched — every rank still holds the complete fp16 model and computes complete gradients, exactly as in DDP. Only the update is partitioned. That is what keeps Stage 1 simple and cheap: it changes who does the optimizer step, not how the network runs.

The memory math

Write out the per-GPU footprint after sharding. The fp16 weights and fp16 gradients are still full on every rank (Stage 1 leaves them alone), but the 12-byte optimizer block is now divided by N:

DDP    per-GPU = 2Ψ + 2Ψ + 12Ψ      = 16Ψ
ZeRO-1 per-GPU = 2Ψ + 2Ψ + 12Ψ/N   = 4Ψ + 12Ψ/N

Two limits make the behaviour clear. At N = 1 (one GPU) the formula returns 16Ψ — no partition, no saving, as expected. As N → ∞ the optimizer term vanishes and the footprint approaches 4Ψ, a hard floor set by the fp16 weights and gradients that Stage 1 never touches. So the best Stage 1 can ever do is a 4× reduction of the model+optimizer memory (from 16Ψ to 4Ψ), and most of that win arrives quickly: the 12Ψ/N term is already an order of magnitude smaller by N = 16. The saving is real precisely because the term being divided — the 12-byte block — was the largest one to begin with.

A worked example

Take a 7.5-billion-parameter model, Ψ = 7.5×10^9, and count bytes per GPU (using 1 GB = 10^9 bytes). Baseline DDP:

16 × 7.5e9 = 120e9 bytes = 120 GB per GPU

That alone will not fit on an 80 GB accelerator — the optimizer states are why. Now shard across N = 64 ranks:

(4 + 12/64) × 7.5e9 = 4.1875 × 7.5e9 ≈ 31.4 GB per GPU

The full sweep shows how fast the optimizer term melts, while the 4Ψ = 30 GB floor stays put:

A 120 GB-per-GPU model becomes a comfortable 31 GB, and going past N = 64 barely helps — you are already scraping the 4Ψ floor. To break through it you need Stage 2 or Stage 3, which shard the other terms too.

The per-step mechanics

How does a step actually run when no rank owns the full optimizer state? The dance has four beats:

1. forward + backward   : each rank computes FULL gradients (2Ψ) on its microbatch
2. reduce-scatter grads : gradients are averaged AND scattered so rank i receives
                          only the reduced gradient for its own shard (1/N of Ψ)
3. local Adam update    : rank i updates m, v, and the fp32 master weights for its
                          shard, then casts its shard to fp16
4. all-gather params    : every rank broadcasts its fresh fp16 shard; all ranks end
                          with the identical, fully-updated fp16 model

The key move is step 2. In plain DDP you would all-reduce the gradients so every rank gets the full averaged gradient — but rank i only needs the slice it will actually update, so ZeRO uses a reduce-scatter instead, delivering exactly that slice. After the local update in step 3, the parameters are out of sync (each rank only refreshed its own shard), so step 4’s all-gather puts them back in lockstep before the next forward pass. Notice that Stage 1 still materializes full gradients in step 1 — that 2Ψ is not saved here; saving it is Stage 2’s job.

Communication: identical to DDP

The result that makes Stage 1 an easy ‘yes’ is that it moves the same number of bytes as the DDP it replaces. Recall how a ring all-reduce works: it is internally a reduce-scatter followed by an all-gather, each moving about Ψ elements per GPU, for a total volume of 2Ψ. DDP pays that every step to synchronize gradients.

DDP      : all-reduce grads          = reduce-scatter (Ψ) + all-gather (Ψ) = 2Ψ
ZeRO-1   : reduce-scatter grads (Ψ) + all-gather params (Ψ)         = 2Ψ

ZeRO Stage 1 simply splits that all-reduce across the step: it does the reduce-scatter half on the gradients before the update, and the all-gather half on the parameters after it. The total, 2Ψ, is unchanged. This is the crux: you get up to a 4× cut in optimizer memory for zero extra communication. That is unusual — memory savings in distributed training normally cost bandwidth. Stage 1 is the one tier where the ledger balances perfectly, which is exactly why it is the default first step and why Stage 3, which does pay more bandwidth, is reserved for when you truly need it.

Stage 1 vs Stage 2 vs Stage 3

ZeRO is a ladder, and Stage 1 is only the first rung. Each stage partitions one more of the three memory components across the same N ranks:

Stage 2 additionally shards the 2-byte gradients, so a rank keeps gradients only for its own shard; it still costs the same 2Ψ communication as Stage 1 and DDP, making it nearly as free. Stage 3 goes all the way and shards the fp16 parameters too — no rank holds the whole model at any instant. That drives memory to 16Ψ/N (linear in the device count, no floor), but it must all-gather parameter shards during both the forward and backward passes, raising communication to roughly 3Ψ, about 1.5× DDP. The lesson: Stage 1 is the free lunch — the biggest single memory term removed at no bandwidth cost. You climb to Stage 2 and 3 only when the remaining 4Ψ (or 2Ψ) floor is still too big to fit.

Practical notes and pitfalls

A few things bite in practice. First, Stage 1’s 4Ψ floor is hard: it never touches activations, which for long sequences or large batches can dwarf the weights — pair ZeRO with activation checkpointing rather than expecting it to solve activation memory. Second, the saving assumes N data-parallel ranks with a genuinely replicated model; combined with tensor or pipeline parallelism, N is only the data-parallel width, not your total GPU count, so do the arithmetic on the right axis. Third, the partition must be reasonably balanced — a lone enormous embedding or tied output matrix assigned to one rank creates a straggler that blows the per-GPU estimate; good implementations shard at a finer granularity than whole tensors. Finally, remember what Stage 1 does not buy: gradients are still full (2Ψ), so if gradient memory is your wall, Stage 1 alone will not clear it. For small-model or CPU-SLM training the takeaway is the ratio, not the GPU cluster: three-quarters of Adam’s footprint is optimizer state, so if memory is tight and you cannot shard, a stateless or lighter-state optimizer (SGD, or 8-bit Adam) attacks the same 12Ψ term from a different direction.