ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

To understand ZeRO, we first need to understand what actually fills GPU memory during training. When training with mixed precision and the Adam optimizer, the memory consumed by a model with Ψ parameters breaks down into two categories: model states and residual states.

Model states include three components. First, the fp16 parameters themselves — 2Ψ bytes. Second, the fp16 gradients — another 2Ψ bytes. Third, the optimizer states: Adam maintains an fp32 copy of the parameters (4Ψ bytes), plus fp32 momentum (4Ψ bytes) and fp32 variance (4Ψ bytes) — totaling 12Ψ bytes. Combined, model states alone consume 16Ψ bytes. For a 1.5B parameter model like GPT-2, that is 24 GB — and we have not even counted activations yet.

Residual states include activations (which can consume 60 GB for GPT-2 at batch size 32), temporary buffers (around 6 GB for large models), and unusable fragmented memory. These compete with model states for the same GPU memory.

Before ZeRO, practitioners faced a painful trade-off between two parallelism strategies.

Data parallelism replicates the full model on every GPU. Each GPU processes a different data batch, computes gradients, and synchronizes them via all-reduce. This is efficient — every GPU does useful compute simultaneously — and easy to implement. But it wastes memory: every GPU holds a full copy of all model states (16Ψ bytes), even though they are all identical.

Model parallelism splits the model across GPUs, so each device holds only a fraction of the parameters. This solves the memory problem but creates new ones: splitting requires manual code refactoring per model, dependencies between layers create sequential bottlenecks that leave GPUs idle, and communication of activations between pipeline stages adds latency. In practice, model parallelism rarely achieves more than 50% compute efficiency.

ZeRO asks: what if we could partition the model states like model parallelism does, but keep the execution pattern of data parallelism?

ZeRO's core idea is elegantly simple: in standard data parallelism, every GPU maintains an identical copy of all model states. This is pure redundancy. ZeRO partitions these states across GPUs instead of replicating them, then uses collective communication to gather pieces when needed.

The insight is that not every GPU needs every piece of state at every moment. Optimizer states are only needed during the parameter update step. Gradients are only needed after the backward pass. And parameters, while needed during forward and backward passes, can be temporarily gathered from other GPUs.

ZeRO-DP introduces this partitioning in three cumulative stages, each one partitioning an additional type of model state. The stages are designed so that earlier stages give the biggest memory savings with the least communication overhead.

Stage 1 — Optimizer State Partitioning (P_os): Each GPU stores only 1/N of the optimizer states (where N is the number of GPUs). After each backward pass, gradients are reduced as usual with all-reduce. Then each GPU updates only its own partition of parameters using its partition of optimizer states. Finally, an all-gather broadcasts the updated parameters to all GPUs. Memory for optimizer states drops from 12Ψ to 12Ψ/N — a 4× savings on 64 GPUs with no additional communication.

Stage 2 — Gradient Partitioning (P_os+g): Building on Stage 1, each GPU also keeps only the gradients it needs — those corresponding to its optimizer partition. Instead of all-reduce (which gives every GPU all gradients), ZeRO uses reduce-scatter (which gives each GPU only its slice). After the update, all-gather synchronizes parameters as before. Memory drops from 14Ψ to (2 + 14/N)Ψ — an 8× savings on 64 GPUs, still with no additional communication.

Stage 3 — Parameter Partitioning (P_os+g+p): Now each GPU stores only 1/N of the parameters too. During forward and backward passes, each layer's full parameters are reconstructed on the fly via all-gather from the GPUs that own the relevant slices. After use, the gathered parameters are discarded. Memory drops to 16Ψ/N — a linear reduction. This stage adds 1.5× communication overhead (two extra all-gathers per step), but enables truly massive models.

Even after ZeRO-DP frees model state memory, residual memory — activations, temporary buffers, and fragmentation — can still be the bottleneck. ZeRO-R addresses each one.

Activation partitioning and CPU offload: Activation checkpointing already reduces activation memory by recomputing activations during the backward pass instead of storing them all. ZeRO-R goes further. It partitions checkpointed activations across GPUs, and for very large models, offloads them to CPU memory. Since activations are only needed during the backward pass, the offloaded data can be prefetched back to GPU just before it is needed.

Constant-size buffers: Large all-reduce operations are most efficient when fusing many small tensors into one large buffer. But naively allocating a buffer proportional to model size creates its own memory pressure. ZeRO-R uses a fixed-size buffer that is large enough for efficiency but does not grow with model size.

Memory defragmentation: During training, memory is allocated and freed in varying sizes, creating gaps (fragmentation) that cannot be used even though total free memory seems sufficient. ZeRO-R pre-allocates contiguous memory blocks for activations and gradients, avoiding fragmentation and ensuring memory is usable.

Memory partitioning is only useful if the communication cost remains manageable. ZeRO achieves this through careful choice of collective communication primitives.

In standard data parallelism, after the backward pass, an all-reduce operation synchronizes gradients across all GPUs. The all-reduce sends 2Ψ data elements per GPU — this is the baseline communication volume.

ZeRO Stage 1 replaces the all-reduce with a reduce-scatter (which gives each GPU only its partition of the reduced gradients) followed by an all-gather (which distributes updated parameters). The combined volume equals the all-reduce baseline: zero overhead.

ZeRO Stage 2 uses reduce-scatter for gradients, which is actually more efficient than all-reduce. Combined with the all-gather for updated parameters, total volume remains equal to the baseline.

ZeRO Stage 3 adds two additional all-gather operations — one before the forward pass and one before the backward pass — to reconstruct full parameters from their partitions. This increases total communication to approximately 1.5× the baseline. For high-bandwidth interconnects like NVLink, this overhead is acceptable for the dramatic memory savings it enables.

The ZeRO paper demonstrated remarkable scaling results. Using ZeRO-100B (a combination of ZeRO Stage 1 with model parallelism), the authors trained models with over 100 billion parameters on 400 NVIDIA V100 GPUs, achieving super-linear speedup and 15 petaflops of sustained throughput.

Critically, ZeRO also made large-scale training accessible without model parallelism. Using ZeRO-DP alone (no model parallelism), models up to 13 billion parameters could be trained — an order of magnitude larger than what standard data parallelism supports. This is especially significant because model parallelism requires specialized code for each model architecture, while ZeRO works as a drop-in replacement for standard data parallelism.

The paper's theoretical analysis showed that with 1024 GPUs, ZeRO Stage 3 could support training of models with over 1 trillion parameters — roughly 6× beyond what model parallelism alone could handle on the same hardware.

ZeRO's legacy extends far beyond any single paper or system. It demonstrated that the memory wall in deep learning training is not fundamental — it is an artifact of redundant data management. By treating model states as shared resources that can be partitioned, gathered, and released on demand, ZeRO turned the memory scaling problem from multiplicative (every GPU holds everything) to additive (each GPU holds a fraction). This insight underpins every modern large-scale training system, from Megatron-LM's 3D parallelism to PyTorch FSDP, and remains the foundation upon which trillion-parameter training is built.