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.