ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Large language models hit a wall that looks like “not enough GPUs,” but is often “each GPU still keeps a full copy of the training state.” Classic data parallelism (DP) is easy and efficient — until Adam’s momentum, variance, master weights, gradients, and parameters all replicate on every rank. Then a 1.5B GPT-2-style model that needs only about 3GB for fp16 weights still refuses to train on a 32GB GPU once optimizer states join the party.

ZeRO (Zero Redundancy Optimizer) keeps DP’s compute grain and communication pattern, but partitions model states across ranks instead of replicating them. Three cumulative stages — optimizer states ($P_{\mathrm{os}}$), gradients ($P_{\mathrm{g}}$), then parameters ($P_{\mathrm{p}}$) — turn cluster memory into one shared pool. The paper’s analysis says a trillion-parameter Adam mixed-precision model (~16TB of model states) fits in principle on 1024 GPUs at ~16GB per device.

What this paper explains

Training memory splits into two buckets:

1. Model states — parameters, gradients, and optimizer states (for Adam: momentum + variance, plus an fp32 master copy under mixed precision).
2. Residual states — activations, temporary buffers, and allocator fragmentation.

ZeRO attacks both:

• ZeRO-DP removes DP redundancy with staged partitions $P_{\mathrm{os}}$, $P_{\mathrm{g}}$, $P_{\mathrm{p}}$.
• ZeRO-R shrinks residual pressure (activation partition/offload, sized buffers, defragmentation).

The evaluated system in the paper, ZeRO-100B, ships $P_{\mathrm{os}}+P_{\mathrm{g}}$ plus ZeRO-R (parameter partitioning $P_{\mathrm{p}}$ was planned as the path to full trillion-scale support).

Prior limits

• DP keeps high compute granularity and modest communication, but replicates every model state → OOM above ~1.4B parameters on 32GB GPUs in their baseline stack.
• Model parallelism (MP) shards layers and saves memory, but communication between layers hurts once you leave a high-bandwidth node (they report ~5 TFLOPS/V100 on a 40B Megatron run across two DGX-2 nodes).
• Pipeline parallelism can partition parameters/activations but brings micro-batch, tied-weight, and batch-norm friction; large micro-batch counts to hide bubbles also grow activation memory.
• CPU offload of full model states can burn large fractions of step time on PCIe copies.

The core tension: DP is usable and fast but memory-wasteful; MP/PP are memory-aware but harder and often less efficient across nodes.

The mechanism

Where the bytes go (mixed-precision Adam)

For $\Psi$ parameters, mixed-precision Adam needs roughly:

Optimizer-state multiplier $K = 12$, so total model-state memory is $2\Psi + 2\Psi + K\Psi = 16\Psi$ bytes. GPT-2 1.5B → about 24GB of model states, not 3GB of fp16 weights.

ZeRO-DP stages (cumulative)

$P_{\mathrm{os}}$ — Optimizer state partitioning. Split momentum/variance/(fp32 masters) into $N_d$ shards. Rank $i$ updates only shard $i$, then all-gather refreshed parameters. Large $N_d$: model-state memory $\approx 4\Psi$ → about 4× reduction vs full DP.

$P_{\mathrm{os}}+P_{\mathrm{g}}$ — Add gradient partitioning. Each rank only keeps reduced gradients for its parameter shard (bucketized reduce-scatter during backward). Large $N_d$: footprint $\approx 2\Psi$ → about 8× reduction.

$P_{\mathrm{os}}+P_{\mathrm{g}}+P_{\mathrm{p}}$ — Add parameter partitioning. Each rank stores only its parameter shard and broadcast/all-gathers remote shards just-in-time for the active layer. Memory falls linearly with $N_d$ (e.g. 64× at $N_d=64$) with about 1.5× communication volume vs baseline DP (~50% increase).

ZeRO-R (residual path)

After model states shrink, activations/buffers/fragments dominate:

• Partition activation checkpoints across MP ranks (remove accidental replication) and optionally offload checkpoints when arithmetic intensity hides PCIe.
• Keep temporary fused buffers constant-sized so they do not grow with $\Psi$.
• Defragment by parking long-lived checkpoints/gradients into preallocated contiguous arenas.

Algorithm walkthrough

One training step under full ZeRO-DP intuition:

1. Forward: For each layer, gather the parameter shard(s) you need; compute; release remote shards you no longer need.
2. Backward: Produce gradients; reduce each gradient bucket only to the owner rank for that shard; free local gradient memory after reduction.
3. Optimizer step: Owner ranks update their optimizer-state shards and corresponding parameter shards.
4. Publish weights: All-gather (or equivalent) so the next step sees updated parameters — or, under $P_{\mathrm{p}}$, keep parameters sharded and gather again on demand.

ZeRO-100B in the paper implements $P_{\mathrm{os}}+P_{\mathrm{g}}$ (+ ZeRO-R), still pairing with modest MP when activations demand it.

What to notice when reading

• Figure 1’s three bars are the whole mental model: same $\Psi$, shrinking per-GPU model-state stack as stages unlock.
• Table 1’s bold cells mark “fits on 32GB V100” combinations — including 1T parameters at $N_d=1024$ with $P_{\mathrm{os}+g+p}$ (~15.6GB).
• Communication is not free under $P_{\mathrm{p}}$, but the paper argues volume stays near DP (1.5×), unlike naive MP across nodes.
• ZeRO does not change Adam’s math — only where states live — so convergence behavior of the optimizer is preserved by design.

Results and evidence

Authors’ reported system results (ZeRO-100B = $P_{\mathrm{os}}+P_{\mathrm{g}}$ + ZeRO-R), not timeless absolutes:

Table 1 memory examples (mixed-precision Adam, per device): a 7.5B model drops from 120GB (DP) to 31.4 / 16.6 / 1.88GB under $P_{\mathrm{os}}$ / $P_{\mathrm{os}}+g$ / $P_{\mathrm{os}+g+p}$ at $N_d=64$.

Limitations

• Full trillion-parameter wall-clock training was still impractical on then-current flops even if memory fit; the evaluated path targets ~100B-class models.
• $P_{\mathrm{p}}$ (stage 3) was not in the ZeRO-100B evaluation slice at publication time — linear memory scaling to 1T is an analysis + roadmap claim in the paper.
• Activation memory can still dominate; ZeRO often still pairs with MP or checkpointing for huge models.
• Very large global batch from high DP can hurt convergence; MP remains a knob for batch-size control.
• Results are on V100-era clusters and DeepSpeed’s implementation details; absolute TFLOPS will not transfer unchanged to newer GPUs.

How to read the paper

1. Abstract + §1 — problem framing and the three-stage punchline.
2. §3 — the 16Ψ mixed-precision Adam accounting (why 1.5B needs ~24GB).
3. §5 — deep dive on $P_{\mathrm{os}}$, $P_{\mathrm{g}}$, $P_{\mathrm{p}}$ + Table 1.
4. §4 / ZeRO-R sections — activations, buffers, fragmentation.
5. §7 — dynamic communication schedule (why volume stays near DP).
6. Evaluation figures 2–3 — 100B/170B throughput and superlinear scaling.
7. Related-work contrasts with MP/PP/offload — what ZeRO refuses to trade away.

Knowledge check

Keep reading

1. Original paper (arXiv:1910.02054) — Figure 1, Table 1, Figures 2–3.
2. Megatron-LM tensor model parallelism — arXiv:1909.08053.
3. Adam optimizer — arXiv:1412.6980.
4. Related Fanout Daily: Switch Transformers, Chinchilla, Scaling Laws, PagedAttention (inference-side memory).

Sources

• Rajbhandari, Rasley, Ruwase, He, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC 2020 — arXiv:1910.02054
• Shoeybi et al., Megatron-LM, 2019 — arXiv:1909.08053
• Kingma & Ba, Adam, 2015 — arXiv:1412.6980
• DeepSpeed library (paper’s open-source vehicle) — github.com/microsoft/DeepSpeed