ZeRO: How DeepSpeed Made Trillion-Parameter Training Practical

Training a billion-parameter language model with standard data parallelism requires every GPU in the cluster to hold a complete copy of the model, its gradients, and the optimizer states. For the Adam optimizer in mixed-precision training, those three categories together consume exactly 16 bytes per parameter: 2 bytes for the fp16 copy of the weights, 2 bytes for the fp16 gradients, and 12 bytes for the fp32 master weight copy and the two Adam moment estimates. A GPT-style model with 7 billion parameters therefore demands 112 gigabytes per GPU before any activations or intermediate buffers are counted. An 80 GB A100 cannot fit it alone. A 175 billion parameter GPT-3-class model would need 2.8 terabytes. The Zero Redundancy Optimizer (ZeRO), introduced in arXiv:1910.02054 by Microsoft researchers in October 2019 and presented at SC20 (Supercomputing 2020), solves this by replacing replication with partitioning. Instead of every GPU holding the same 16 bytes per parameter, each GPU holds only its assigned slice, with the rest reconstructed on demand via collective communication.

The paper and its authors

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models was submitted to arXiv on October 4, 2019 by Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He, all at Microsoft. The paper was revised and expanded through May 13, 2020 and published at SC20. All four authors are members of the Microsoft DeepSpeed team. Rajbhandari is the first author and has continued to lead ZeRO research through subsequent extensions. The paper was accompanied by the open-source release of DeepSpeed (github.com/microsoft/DeepSpeed), a PyTorch extension library that implements ZeRO alongside other distributed training optimizations. DeepSpeed has since been adopted by Hugging Face Transformers, Megatron-LM, and PyTorch Lightning as an optional training backend.

Where the 16 bytes per parameter comes from

The 16-byte figure deserves a precise accounting because it is the foundation for all of ZeRO's memory reduction claims. Standard large-model training uses mixed precision: forward and backward passes run in fp16 to maximize GPU tensor core throughput, while optimizer updates run in fp32 to preserve numerical stability. The Adam optimizer requires five values per parameter: the fp16 parameter itself (2 bytes), the fp16 gradient produced during the backward pass (2 bytes), an fp32 master copy of the parameter used to apply updates (4 bytes), the first moment estimate (a running mean of gradients, 4 bytes in fp32), and the second moment estimate (a running variance of gradients, 4 bytes in fp32). The optimizer-state component alone is 12 bytes per parameter, three-quarters of the total. For a 1.5 billion parameter model such as GPT-2 XL, the model states require 24 gigabytes, and that is without any activations. On a 32 GB V100, even a 1.5 billion parameter model fits in memory only if the 24 GB of model states leaves enough room for activations and buffers, which it often does not.

• fp16 parameters: 2 bytes per parameter (used in forward and backward passes).
• fp16 gradients: 2 bytes per parameter (accumulated during the backward pass).
• fp32 master parameters: 4 bytes per parameter (the numerically stable copy updated by Adam).
• fp32 first moment (momentum): 4 bytes per parameter (running mean of gradients).
• fp32 second moment (variance): 4 bytes per parameter (running mean of squared gradients).
• Total model states: 16 bytes per parameter, replicated on every GPU in standard data parallelism.

Three stages of partitioning

ZeRO implements memory reduction as three incremental stages, each partitioning an additional category of model states across the N GPUs in the data-parallel group. In Stage 1, only the optimizer states (the 12-byte-per-parameter component) are partitioned. Each GPU stores 1/N of the optimizer states and updates only its slice. The fp16 parameters and gradients remain replicated. With N GPUs, per-device memory for model states drops from 16 bytes per parameter to approximately 4 plus 12/N bytes. For large N, this approaches 4 bytes, a 4-fold reduction. In Stage 2, the gradients are also partitioned: each GPU accumulates and stores only the gradients for its slice of optimizer states, discarding the rest after the reduction step. Per-device memory approaches 2 plus 14/N bytes, close to a 8-fold reduction at large N. In Stage 3, the model parameters are also partitioned. Each GPU holds only 1/N of the fp16 weights, gathering the weights for a layer before running the forward or backward pass for that layer and immediately discarding them. Memory per device approaches 16/N bytes, a reduction that scales linearly with the number of GPUs. The paper demonstrates this on 400 V100 GPUs achieving 15 petaflops of throughput on a 100 billion parameter model, an 8-fold increase in trainable model size compared to the prior state of the art.

Communication overhead: the key tradeoff

Partitioning model states requires communication that baseline data parallelism does not. Standard data parallelism performs a single AllReduce to sum gradients across all GPUs at the end of each backward pass. ZeRO Stage 1 has the same communication volume as standard data parallelism: each gradient is still reduced once, and optimizer states require no extra communication because each GPU updates only its partition. Stage 2 adds a ReduceScatter for gradients in place of the standard AllReduce, which has the same total bytes transferred. Stage 3 adds AllGather calls to reconstruct parameters before each layer's forward and backward pass. The paper's analysis shows that Stage 3 increases communication volume by 50 percent relative to standard AllReduce: a backward pass now includes one AllReduce worth of gradient communication plus half an AllReduce worth of parameter reconstruction, totaling 1.5 AllReduces of bandwidth. For most training jobs on interconnects faster than 100 gigabits per second, this 50 percent overhead does not reduce throughput because the compute time for large layers is long enough to overlap with the communication for earlier layers.

ZeRO-Infinity: offloading to CPU and NVMe

ZeRO-Infinity, published in a follow-up paper (arXiv:2104.07857, April 2021) by the same Microsoft team, extends Stage 3 by offloading partitioned parameters and optimizer states from GPU memory to CPU DRAM or NVMe solid-state drives. A GPU communicating via NVMe can access hundreds of gigabytes of storage at up to 7 gigabytes per second on modern PCIe 4.0 drives. CPU DRAM offers lower latency and higher bandwidth than NVMe, typically 50 to 100 gigabytes per second per CPU socket. ZeRO-Infinity introduces memory-centric tiling, which breaks individual large weight matrices into tiles that fit in GPU memory, so even a single GPU with 40 GB of HBM can train a model whose total weight footprint far exceeds 40 GB by streaming tiles from CPU memory. The paper reports training a 32 trillion parameter language model (a mixture-of-experts variant) on 512 NVIDIA V100 GPUs, a scale unreachable by any prior training system. For practitioners, ZeRO-Infinity is available in DeepSpeed as the offload parameter of the zero_optimization block in the DeepSpeed config JSON, requiring no model code changes.

Adoption in major training runs

ZeRO's first major public showcase was Turing-NLG, a 17.2 billion parameter autoregressive language model trained by Microsoft and announced in February 2020, at the time the world's largest language model by parameter count. Turing-NLG used ZeRO Stage 1 combined with model parallelism across 256 V100 GPUs. Since then, ZeRO has been used in or cited by the training of Megatron-Turing NLG 530B (Microsoft and NVIDIA, 2021, 530 billion parameters trained on 280 A100 GPUs), BLOOM (BigScience, 2022, 176 billion parameters trained on 384 A100 80 GB GPUs using a combination of ZeRO Stage 1 and Megatron-style pipeline and tensor parallelism), and Falcon 180B (TII, 2023). Hugging Face's Accelerate library integrates DeepSpeed's ZeRO as a one-line configuration option. The PyTorch team released a native ZeRO-style implementation called Fully Sharded Data Parallel (FSDP) in PyTorch 1.11, described by Meta researchers as inspired directly by ZeRO Stage 3, enabling parameter sharding without the DeepSpeed dependency.

The practical impact of ZeRO is that it converted the question of whether a model could be trained from a hardware-acquisition problem into a configuration problem. Before ZeRO, training a 100 billion parameter model required either a small number of expensive servers with specialized high-bandwidth interconnects and custom model-parallel code, or it was simply not feasible. After ZeRO, a cluster of commodity A100 GPUs with fast ethernet could train the same model by partitioning states across more devices, with no changes to the model's forward or backward pass code. The Stage 1 and Stage 2 configurations preserve nearly the same communication volume as standard data parallelism and require only a JSON config change. Stage 3 and ZeRO-Infinity trade bandwidth for memory at configurable ratios. PromptingIndex covers ZeRO alongside related training infrastructure including Flash Attention, grouped-query attention, and the Chinchilla scaling laws, all of which address different bottlenecks in bringing large language models to production at scale.