ZeRO (Zero Redundancy Optimizer)

Definition

ZeRO (Rajbhandari et al., DeepSpeed, arXiv:1910.02054) eliminates the memory redundancy of data parallelism: instead of every rank holding an identical copy of all model states, ZeRO partitions optimizer states, gradients, and parameters across the $N_d$ data-parallel ranks, materializing a full tensor only transiently when a computation needs it. It comes in three cumulative stages, each sharding one more category. FSDP FULL_SHARD is ZeRO-3’s PyTorch-native equivalent.

Key math / How it works

Model states are $16\Psi$ bytes for $\Psi$ params under mixed-precision Adam, factored as $(2+2+K)\Psi$ with $K=12$ (the fp32 master + $m$ + $v$). Per-GPU model-state memory by stage:

$$\begin{array}{rl}\text{DDP (baseline):}& 16\Psi \\ \text{ZeRO-1 }(P_{os}):& 4\Psi +\frac{12\Psi }{N_d}\text{(shard optimizer states)}\\ \text{ZeRO-2 }(P_{os+g}):& 2\Psi +\frac{14\Psi }{N_d}\text{(+ shard gradients)}\\ \text{ZeRO-3 }(P_{os+g+p}):& \frac{16\Psi }{N_d}\text{(+ shard parameters)}\end{array}$$

Whatever a stage shards contributes $(\cdot )/N_d$; whatever it still replicates is an irreducible floor. Only ZeRO-3 → 0 as $N_d\to \infty$; ZeRO-1 leaves a $4\Psi$ floor, ZeRO-2 a $2\Psi$ floor that no number of GPUs removes. ZeRO-1 alone is ~4× for free (12 of 16 bytes are optimizer state).

Communication cost — the memory↔comms tradeoff:

• ZeRO-1 & ZeRO-2: 2Ψ, same as DDP. The gradient all-reduce is decomposed into reduce-scatter (grads land on the rank owning the matching optimizer shard) + all-gather (of updated params). 4×–8× memory savings for zero extra bytes on the wire — why ZeRO-2 is often the sweet spot.
• ZeRO-3: 3Ψ = 1.5× DDP. Params are sharded, so each layer’s weights must be all-gathered just-in-time in forward (Ψ) and again in backward (Ψ), plus reduce-scatter grads (Ψ). A 50% comms increase for up to $N_d\times$ memory reduction.

Worked (70B, $N_d=64$): DDP = 1120 GB/GPU (impossible); ZeRO-1 = 293 GB (the 280 GB $4\Psi$ floor kills it); ZeRO-2 = 155 GB (140 GB $2\Psi$ floor kills it); ZeRO-3 = 17.5 GB, fits. The dividing line for “must shard params” is exactly whether bf16 params ($2\Psi$) fit on one GPU.

Offloading trades bandwidth for capacity: ZeRO-Offload (2101.06840) pushes optimizer state + fp32 step to CPU (on ZeRO-2, PCIe ~12 GB/s bottleneck, needs ~2K batch/GPU); ZeRO-Infinity (2104.07857) extends to CPU+NVMe (on ZeRO-3), fitting ~1T params/DGX-2 node via bandwidth-centric partitioning — a capacity escape hatch, not a throughput win.

Why it matters

ZeRO is the reason a model whose training state exceeds one GPU can be trained at all without splitting the compute. It is the “DP that also shards the model,” implemented in DeepSpeed and mirrored by FSDP; ZeRO-3/FSDP is the default for large models and the “DP leg” of N-D parallelism. DeepSeek-V3 used ZeRO-1 DP as its outermost axis.

Taught in

• 09-zero-and-fsdp — §3 the 16 B/param count, §4 stages + comms, §6 offloading.

See also

• fsdp — ZeRO-3 the PyTorch way
• optimizer-state-memory — the 16 B/param being sharded
• data-parallelism — the baseline ZeRO de-duplicates
• collective-communication — reduce-scatter + all-gather
• 3d-parallelism — ZeRO as the DP axis
• gpu-systems-for-llms