ZeRO Optimization › ZeRO Optimization: Optimizer, Gradient, and Parameter Partitioning

Modern language models have crossed a threshold where no single GPU can hold all training state in memory. GPT-3 with 175 billion parameters requires roughly 700 GB of memory at 32-bit precision just to store its weights. Add the optimizer states and gradients that training requires, and you need several terabytes. A single A100 with 80 GB of HBM cannot hold even the parameters alone.

The standard response has been to distribute training across many GPUs. As we covered in the data parallelism chapter, each device holds a full model replica and processes a different data shard, then synchronizes gradients. This works well for moderate model sizes, but it scales memory requirements in a fundamentally wasteful way: 100 GPUs each holding a full copy of a large model means 99 copies of every tensor are pure redundancy. Adding more hardware does not reduce the per-device memory burden, it simply multiplies the number of devices carrying the same redundant data.

ZeRO (Zero Redundancy Optimizer) was introduced by Rajbhandari et al. at Microsoft Research in 2019 and implemented in the DeepSpeed library. Its central insight is that the redundancy in data-parallel training is not just wasteful, it is unnecessary. The parameters, gradients, and optimizer states held by each GPU are identical across all replicas because every device applies the same update to the same model. ZeRO eliminates this redundancy by partitioning each of these components across the fleet of GPUs, so each device owns a unique shard rather than a full copy.

This might seem to require each GPU to constantly request tensors from other devices, creating a communication bottleneck. ZeRO's elegance is that it achieves this partitioning while keeping communication volume equal to or only modestly above standard data parallelism. The three stages of ZeRO progressively partition more state while following a carefully designed communication schedule that hides most of the cost.

This chapter covers how ZeRO works mechanically through its three stages, the communication patterns each stage introduces, and the quantitative memory savings they achieve. We also examine ZeRO-Offload, ZeRO-Infinity, and ZeRO++, which extend the core idea to CPU memory, NVMe storage, and quantized communication respectively.