The redundancy hiding in data parallelism

Standard data parallelism is the simplest way to train on many GPUs: replicate the entire model on every device, give each a different slice of the batch, run forward and backward independently, then all-reduce the gradients so every replica applies the same update. It is easy to reason about and scales throughput almost linearly. But it is brutally memory-inefficient. Each of the N GPUs holds a full copy of the parameters, a full copy of the gradients, and a full copy of the optimizer states — and after the synchronized update those copies are bit-for-bit identical.

That is pure redundancy. The model-state memory does not shrink as you add GPUs; it is duplicated N times over. So data parallelism gives you more compute but not more capacity: a model that does not fit on one GPU still does not fit on a hundred. ZeRO attacks exactly this. It keeps the data-parallel execution model — each GPU still processes its own micro-batch through the whole network — but stops storing N identical copies of the training state. Instead it splits that state across the GPUs and reconstructs pieces on demand.

Advertisement

The memory model: 2 + 2 + K bytes per parameter

To reason about the savings you need a memory model. Modern large-model training uses mixed precision with Adam, and for a model of Ψ parameters the per-parameter memory splits into three buckets:

fp16 parameters      : 2Ψ bytes
fp16 gradients       : 2Ψ bytes
optimizer states (K) : KΨ bytes
-------------------------------------
total                : (2 + 2 + K) · Ψ bytes

The first 2Ψ is the working copy of the weights in fp16; the second 2Ψ is the fp16 gradients. The interesting term is K, the optimizer state per parameter. Mixed-precision Adam keeps an fp32 master copy of the weights (4 bytes) plus the fp32 first moment / momentum (4 bytes) and second moment / variance (4 bytes) — so K = 4 + 4 + 4 = 12. Those fp32 states, not the fp16 weights, are the memory hogs: they are 12Ψ of the 16Ψ total. That imbalance is precisely why ZeRO partitions the optimizer states first.

Advertisement

The vanilla baseline: 16Ψ on every GPU

Plug K = 12 into the model and vanilla data parallelism costs (2 + 2 + 12)Ψ = 16Ψ bytes of model-state memory on each GPU, independent of how many GPUs you have. For a 1.5-billion-parameter model that is 16 × 1.5e9 = 24 GB — already tight on a 32 GB card once activations are added. Scale to 7.5 billion parameters and it is 16 × 7.5e9 = 120 GB, which fits on no single accelerator made.

Adding GPUs does not help, because every one independently needs the full 120 GB. This is the wall ZeRO is built to break: the 16Ψ is fixed, dominated by the 12Ψ optimizer term, and completely redundant across the data-parallel group. Divide that 16Ψ by the number of GPUs instead of replicating it, and the hardware that could not hold the model at all holds it comfortably. That division is what the three ZeRO stages deliver, one bucket at a time.

Stage 1 — P_os: partition the optimizer states

ZeRO Stage 1, written P_os (partition optimizer states), targets the fattest bucket. The 12Ψ bytes of fp32 master weights, momentum, and variance are split into N equal shards; GPU i owns only shard i. Parameters and gradients are still replicated in full, but each GPU is now responsible for updating only its 1/N slice of the weights.

The per-GPU memory becomes 2Ψ + 2Ψ + 12Ψ/N. As N grows the optimizer term collapses toward zero and the cost approaches 4Ψ — a 4× reduction from the 16Ψ baseline. Crucially, Stage 1 costs nothing in extra communication: a standard data-parallel step already moves 2Ψ of gradient/parameter traffic, and by reorganizing that same traffic as a reduce-scatter of gradients plus an all-gather of the updated weights, P_os keeps total communication identical. It is the cheapest, safest first move — explored fully in the Stage 1 deep-dive (tm_zero_1_math).

Stage 2 — P_g: partition the gradients too

Stage 2, P_os+g, adds gradient partitioning (P_g) on top of Stage 1. The observation is that once GPU i only updates its 1/N shard of parameters, it only ever needs the gradients for that shard. So there is no reason to keep the full 2Ψ gradient buffer on every device. During the backward pass, as each layer’s gradients are produced they are reduce-scattered to the owning GPU and the rest are discarded.

Now per-GPU memory is 2Ψ + (2 + 12)Ψ/N = 2Ψ + 14Ψ/N, which approaches 2Ψ for large N — an 8× reduction. Like Stage 1, P_g preserves the baseline communication volume of 2Ψ, because the reduce-scatter it uses moves no more data than the all-reduce it replaces. This is the reason Stage 2 is the popular default: it roughly doubles the savings of Stage 1 for free in bandwidth terms. The full treatment lives in the Stage 2 deep-dive (tm_zero_2_math).

Stage 3 — P_p: partition the parameters too

Stage 3, P_os+g+p, takes the final step and partitions the parameters themselves (P_p). Now nothing is fully replicated: each GPU permanently stores only its 1/N slice of the fp16 weights, its slice of the gradients, and its slice of the optimizer states. When a layer is needed for the forward or backward pass, its parameters are all-gathered from the owners just in time, used, and then freed again.

Per-GPU model-state memory drops to the clean (2 + 2 + 12)Ψ/N = 16Ψ/N — a reduction that is linear in the number of GPUs. With enough devices the model state per GPU becomes arbitrarily small, which is what makes trillion-parameter training feasible. The catch is communication: because parameters must be gathered twice (once in forward, once in backward) on top of the gradient reduce-scatter, Stage 3 raises total volume to about 3Ψ, roughly 1.5× the baseline. This stage is functionally equivalent to PyTorch FSDP; see the deep-dive (tm_zero_3_math).

The per-stage memory formula and a worked example

Collecting the three stages into one table, with N the data-parallel degree and K = 12:

baseline (DP) : 2Ψ + 2Ψ + KΨ            = 16Ψ
Stage 1 (P_os)  : 2Ψ + 2Ψ + KΨ/N          → 4Ψ   (~4x)
Stage 2 (P_g)   : 2Ψ + (2 + K)Ψ/N          → 2Ψ   (~8x)
Stage 3 (P_p)   : (2 + 2 + K)Ψ/N = 16Ψ/N   → 0     (Nx)

Now make it concrete with a 7.5B-parameter model on N = 64 GPUs, so Ψ = 7.5e9:

baseline : 16 × 7.5        = 120  GB / GPU   (fits nowhere)
Stage 1  : 4×7.5 + 12×7.5/64 = 31.4 GB / GPU
Stage 2  : 2×7.5 + 14×7.5/64 = 16.6 GB / GPU
Stage 3  : 16 × 7.5 / 64      = 1.9  GB / GPU

The progression is stark: a model that fit on no GPU at 120 GB drops to 1.9 GB per device under Stage 3 — a 64× cut matching the GPU count exactly. Even Stage 1 alone turns the impossible into a comfortable fit.