ZeRO and the Memory Tricks

It is October 2019. A team at Microsoft is training language models whose optimizer state alone needs sixteen bytes of memory for every parameter — replicated, byte for byte, on every one of hundreds of GPUs, before a single activation or training example has touched the chip. A gradient-checkpointing paper three years earlier had already attacked half of the memory bill, the activations; over the next four years, three more papers stop treating that per-GPU replication of model state as an unavoidable cost of data parallelism. One shards it, one offloads it to ordinary RAM, and one folds the whole trick natively into PyTorch — and between them they raise the largest model trainable on a fixed GPU budget by roughly two orders of magnitude.

Where We Were

Chapter 4 frames the general shape of the problem: GPU memory grows far more slowly than the appetite of the models researchers want to train, and beyond a point compute is not the binding constraint — bytes are. Chapter 8 and Chapter 9 cover the two classical answers to “the model doesn’t fit”: replicate the whole model on every worker and split the data (data parallelism), or split the model itself across workers and pay a communication tax for every layer boundary crossed (tensor and pipeline parallelism). By the late 2010s, both answers were running into the same wall from different directions.

Data parallelism’s problem was redundancy, not compute. Every data-parallel worker keeps a complete copy of the model’s parameters, its gradients, and — the part that is easy to forget — its optimizer’s internal state. Chapter 5 covers what Adam and AdamW actually compute; the detail that matters here is only the bookkeeping: Adam keeps a running first moment and a running second moment for every parameter, in addition to the parameter itself. Trained in the mixed-precision regime that had become standard practice after 2017, that bookkeeping compounds: an FP16 copy of the parameters and gradients each cost 2 bytes per parameter, and Adam’s FP32 master parameters, momentum, and variance each cost 4 bytes per parameter, for a total that Rajbhandari et al. would later name explicitly: 16 bytes of model-state memory for every parameter, full stop, before any activation is stored (Rajbhandari, Rasley, Ruwase, and He, 2019). A 1.5-billion-parameter model — GPT-2’s largest configuration — already needs 24 GB just for that bookkeeping, on top of activations, and every one of thousands of data-parallel GPUs was, as of 2019, holding an identical copy of it.

Model parallelism’s problem was the opposite: it split the redundant state, but only by cutting the model into pieces that talk to each other on every forward and backward pass, which is exactly the communication overhead Chapter 9 quantifies. That overhead scales with how finely the model is cut, which means model parallelism alone does not fit comfortably inside a single node’s fast interconnect once a model needs to be split across more than a handful of GPUs.

The other half of the memory bill came from a source unrelated to redundancy at all: activations. A forward pass through an n-layer network stores, in the naive implementation, every intermediate activation it computes, because the backward pass needs them to compute gradients. Deep learning had spent 2015 and 2016 pushing depth aggressively — ResNets already reached over 100 layers by late 2015 — and a network’s activation memory grows linearly with its depth, independent of how many parameters it has or how many GPUs are training it. Splitting the model or the batch across more devices does nothing to shrink that linear-in-depth activation footprint per device; it was a separate bottleneck that neither data nor model parallelism addressed.

The Breakthrough

The first fix predates ZeRO by three years and targets the activation problem specifically. Chen, Xu, Zhang, and Guestrin observed that a network does not need every intermediate activation held in memory simultaneously — it only needs whichever ones the backward pass is currently consuming (Chen, Xu, Zhang, and Guestrin, 2016). Their algorithm stores activations only at a sparse set of checkpoint layers, roughly one in every $\sqrt{n}$ layers, and discards everything in between; when the backward pass reaches a segment whose activations were discarded, it simply reruns the forward pass for that short segment to reconstruct them on demand, then proceeds with the backward computation as usual. The technique costs one additional forward pass per training step and, in exchange, drops activation memory from linear in depth to roughly the square root of depth. On a 1,000-layer residual network on ImageNet, the paper reports cutting memory from 48 GB to 7 GB at a cost of about 30% more running time — a trade any lab whose bottleneck was memory, not GPU-hours, would take.

Gradient checkpointing solved activations. It did nothing about the 16-bytes-per-parameter model-state redundancy that data parallelism was still paying on every GPU. That is the problem ZeRO — Zero Redundancy Optimizer — set out to eliminate, not by splitting the model the way tensor parallelism does, but by splitting the bookkeeping while leaving the model’s logical structure, and the training math it computes, completely untouched. Rajbhandari, Rasley, Ruwase, and He, at Microsoft, proposed partitioning the model states across the data-parallel group itself, in three cumulative stages of increasing aggressiveness, each retrieving via communication only the slice of state a given GPU needs at the moment it needs it, and discarding the rest immediately afterward (Rajbhandari, Rasley, Ruwase, and He, 2019). Every data-parallel worker still computes on the full model and the full batch it is assigned; it simply stops being the sole owner of a full private copy of everything the optimizer needs to remember. On 400 GPUs, the paper reports training models with over 100 billion parameters — combining ZeRO’s sharding with Megatron-style tensor parallelism kept within a single node — at an aggregate throughput above 15 petaflops, and separately reports training a 17-billion-parameter model, Turing-NLG, at a sustained 41.4 teraflops per GPU.

Two years later, ZeRO-Offload asked a more aggressive question: if the optimizer states and gradients can be sharded across GPUs, why keep them on a GPU at all? Ren, Rajbhandari, Yazdani Aminabadi, Ruwase, Yang, Zhang, Li, and He moved the optimizer’s FP32 states, its momentum and variance, and the gradient reduction and update computation itself onto the host’s ordinary CPU and system RAM, keeping only the FP16 parameters and the forward/backward computation on the GPU (Ren et al., 2021). CPU RAM is far cheaper and far more abundant per dollar than GPU HBM, and a CPU, while much slower per operation, was fast enough to keep up with an Adam update once the data movement between GPU and CPU was minimized and overlapped with GPU compute. The result was a single 32 GB NVIDIA V100 training a 13-billion-parameter model — a 10x increase over the roughly 1.4 billion parameters PyTorch alone could fit on the same card — at 40 teraflops per GPU for a 10-billion-parameter model, scaling near-linearly up to 128 GPUs, and, combined with tensor-slicing model parallelism across a 16-GPU DGX-2 node, training a 70-billion-parameter model, a 4.5x increase over model parallelism alone on the same hardware.

By 2023, sharding model states across data-parallel workers was established enough as a technique that it needed to stop being a separate third-party library and become a first-class part of the framework itself. Zhao, Gu, Varma, and collaborators at Meta built Fully Sharded Data Parallel directly into PyTorch, implementing the same full-sharding idea ZeRO’s third stage had introduced — parameters, gradients, and optimizer states all partitioned across workers — but rebuilt around PyTorch’s own module and communication primitives rather than as an external wrapper, with careful engineering of when each shard is fetched and released to keep the extra communication mostly hidden behind ongoing GPU compute (Zhao et al., 2023). The paper reports training a 175-billion-parameter model on up to 512 A100 GPUs at 173 to 186 teraflops per GPU — 55 to 60% of the hardware’s peak — with per-GPU throughput degrading only about 7% when scaling a separate 11-billion-parameter model from 8 to 512 GPUs, evidence that the sharding-and-communication design held up at scale rather than merely working in a small demo.

The Mathematics

Every technique in this chapter is an answer to the same accounting question: how many bytes does training one parameter actually cost, and which of those bytes can be avoided, shrunk, sharded, or moved somewhere cheaper. Start with the baseline that motivated ZeRO — the memory a single data-parallel worker holds for its model states under mixed-precision training with the Adam optimizer.

$$M_{\mathrm{states}} = (2 + 2 + K)\,\Psi, \qquad K = 12$$

Term	Bytes per parameter	What it is
FP16 parameters	2	Used for the forward and backward pass
FP16 gradients	2	Produced by backpropagation
FP32 parameter copy	4	Master copy Adam updates for stability
FP32 momentum	4	Adam’s first-moment estimate
FP32 variance	4	Adam’s second-moment estimate
Total ($2+2+K$)	16	Model-state memory, replicated on every data-parallel worker before ZeRO

ZeRO’s contribution is to make that 16Ψ a function of how many data-parallel workers are sharing the burden, rather than a fixed tax paid identically by each one. Each of its three stages shards a progressively larger slice of that total across the data-parallel group of size $N_d$:

$$M_{\mathrm{states}}(p, N_d) = \begin{cases} 4\Psi + \dfrac{12\Psi}{N_d} & p = 1 \ \text{(shard optimizer states)} \\[6pt] 2\Psi + \dfrac{14\Psi}{N_d} & p = 2 \ \text{(shard optimizer states and gradients)} \\[6pt] \dfrac{16\Psi}{N_d} & p = 3 \ \text{(shard optimizer states, gradients, and parameters)} \end{cases}$$

Symbol	Meaning
$M_{\mathrm{states}}$	model-state memory resident on one GPU, in bytes per parameter
$p$	ZeRO stage, 1 through 3
$N_d$	number of data-parallel workers sharing the shard
$\Psi$	number of model parameters

The reason this scheme is adopted so readily rather than treated as a clever but impractical trick is that the extra communication it costs is small relative to what it saves. Baseline data-parallel training already moves $2\Psi$ bytes of gradients per step through a ring all-reduce; ZeRO’s first two stages reorganize that same traffic into a reduce-scatter plus an all-gather and move exactly the same $2\Psi$, with zero communication overhead over the baseline. Only the third stage, which also shards the parameters themselves, adds an extra all-gather to reconstruct each layer’s full weights just before it is used, bringing total traffic to $3\Psi$ — a 1.5x increase in communication bought for a reduction in memory that scales with the entire cluster.

Gradient checkpointing and ZeRO shard two different, additive parts of the same memory bill — activations and model states respectively — which is why the two techniques compose rather than compete:

$$M_{\mathrm{total}} \approx \frac{16\Psi}{N_d} + O(\sqrt{n})$$

ZeRO-Offload’s contribution does not change either formula; it changes which processor’s memory the terms live in. Rather than sharding the FP32 states across GPUs, it keeps the FP16 parameters and the forward/backward compute on the GPU and moves the FP32 states, the gradient reduction, and the Adam update itself onto the host CPU’s RAM — trading a slower processor for one with a far larger and far cheaper pool of memory, and overlapping the GPU-to-CPU data movement with ongoing GPU compute so the CPU’s lower throughput does not become the new bottleneck.

What It Enabled

	Baseline data parallelism	ZeRO stage 3 / FSDP
Model-state memory per GPU	$16\Psi$ bytes, full replica	$16\Psi / N_d$ bytes, sharded
Extra communication vs. baseline all-reduce	none (this is the baseline)	1.5x ($3\Psi$ vs. $2\Psi$)
Largest model on a single 32 GB V100	~1.4B parameters (PyTorch alone)	13B parameters (ZeRO-Offload, 10x larger)
Largest model demonstrated	bounded by single-GPU memory	175B parameters on 512 A100 GPUs (FSDP), 170B+ on 400 GPUs (ZeRO)
Where the extra bytes go	replicated on every GPU	sharded across GPUs, or offloaded to CPU RAM

Stage	Partitions	Memory formula	Max reduction	Communication
Baseline DP	none	$16\Psi$	1x	$2\Psi$
ZeRO-1, $P_{os}$	optimizer states	$4\Psi + 12\Psi/N_d$	~4x	$2\Psi$, same as baseline
ZeRO-2, $P_{os+g}$	+ gradients	$2\Psi + 14\Psi/N_d$	~8x	$2\Psi$, same as baseline
ZeRO-3, $P_{os+g+p}$	+ parameters	$16\Psi/N_d$	linear in $N_d$	$3\Psi$, 1.5x baseline

The practical effect of this arc was to decouple “how large a model I can train” from “how much memory a single GPU has,” replacing it with “how much memory my cluster has, or how much system RAM I can attach to each node.” That decoupling is what let the Chapter 9 toolkit of tensor and pipeline parallelism become optional rather than mandatory for a growing range of model sizes: a lab could reach for ZeRO or FSDP first, and only add model parallelism once even a fully sharded model still would not fit, rather than the reverse. It is also the reason Chapter 11 treats ZeRO-style sharding, offload, and checkpointing as defaults stacked together in essentially every published large-model training recipe rather than as alternatives a team chooses between.

The same accounting shows up again wherever a training or fine-tuning setup is memory-constrained rather than compute-constrained. Chapter 13 covers mixture-of-experts models, whose sparse routing multiplies total parameter count far faster than active compute per token, making state-sharding close to mandatory rather than an optimization. Chapter 14 covers an entirely different answer to the same 16Ψ bill: instead of sharding the optimizer state across more hardware, freeze almost all of $\Psi$ and train a tiny fraction of it, shrinking the bill at its source rather than distributing it. Chapter 20 revisits this chapter’s arc as one instance of a recurring pattern — a hardware constraint (GPU memory capacity growing slower than model size) provoking a software response (sharding, offload) that buys a few more years of runway on the same silicon generation.

Key Papers

Paper	Authors	Year	Link	Why it matters
Training Deep Nets with Sublinear Memory Cost	Chen, Xu, Zhang, and Guestrin	2016	arXiv:1604.06174	Gradient checkpointing; cut activation memory from linear to roughly $\sqrt{n}$ in network depth at the cost of one extra forward pass
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models	Rajbhandari, Rasley, Ruwase, and He	2019	arXiv:1910.02054	Sharded optimizer, gradient, and parameter state across data-parallel workers in three stages, with little to no added communication for the first two
ZeRO-Offload: Democratizing Billion-Scale Model Training	Ren et al.	2021	arXiv:2101.06840	Moved optimizer states and the optimizer update itself to CPU RAM, fitting a 13B-parameter model on one 32 GB GPU
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel	Zhao et al.	2023	arXiv:2304.11277	Rebuilt ZeRO stage 3’s full-sharding idea as a native PyTorch primitive, trained at 175B parameters on 512 GPUs

PM Lens