ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
Abstract
Large deep learning models offer significant accuracy gains, but training billions to trillions of parameters is challenging. Existing solutions such as data and model parallelisms exhibit fundamental limitations to fit these models into limited device memory, while obtaining computation, communication and development efficiency. We develop a novel solution, Zero Redundancy Optimizer (ZeRO), to optimize memory, vastly improving training speed while increasing the model size that can be efficiently trained. ZeRO eliminates memory redundancies in data- and model-parallel training while retaining low communication volume and high computational granularity, allowing us to scale the model size proportional to the number of devices with sustained high efficiency. Our analysis on memory requirements and communication volume demonstrates: ZeRO has the potential to scale beyond 1 Trillion parameters using today's hardware. We implement and evaluate ZeRO: it trains large models of over 100B parameter with super-linear speedup on 400 GPUs, achieving throughput of 15 Petaflops. This represents an 8x increase in model size and 10x increase in achievable performance over state-of-the-art. In terms of usability, ZeRO can train large models of up to 13B parameters (e.g., larger than Megatron GPT 8.3B and T5 11B) without requiring model parallelism which is harder for scientists to apply. Last but not the least, researchers have used the system breakthroughs of ZeRO to create the world's largest language model (Turing-NLG, 17B parameters) with record breaking accuracy.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
ZeRO introduces a systematic memory partitioning strategy that eliminates redundancy in data-parallel training, fundamentally reshaping distributed ML infrastructure and enabling the practical training of trillion-parameter models. By rigorously formalizing state, gradient, and parameter partitioning, the work bridges the gap between algorithmic simplicity and extreme-scale efficiency, establishing a new standard for distributed training that has been universally adopted across the field and directly catalyzed the rapid scaling of modern foundation models.
Methodology
The paper introduces ZeRO (Zero Redundancy Optimizer), a principled framework that systematically partitions optimizer states, gradients, and model parameters across data-parallel workers instead of replicating them. By formalizing three progressive stages (ZeRO-1, ZeRO-2, ZeRO-3), the methodology eliminates the O(N) memory redundancy inherent in standard data parallelism while preserving its computational simplicity and avoiding the complex graph partitioning required by model parallelism. The authors provide rigorous memory footprint equations and communication volume analysis, demonstrating that partitioning can be scheduled to overlap with computation and maintain high arithmetic intensity. The approach is architecturally clean, mathematically grounded, and elegantly bridges the gap between naive data parallelism and complex tensor/pipeline parallelism.
Results & evaluation
The evaluation is extensive and well-calibrated for systems research. Benchmarks span up to 400 GPUs, demonstrating super-linear scaling and sustained throughput of 15 PFLOPS. The paper compares ZeRO against Megatron-LM and baseline data/model parallelism, showing an 8x increase in maximum trainable model size and 10x performance gains. Real-world validation includes training Turing-NLG (17B) and synthetic 100B+ parameter models. The scaling curves closely match theoretical predictions, and the ablation studies effectively isolate the communication overhead introduced by each ZeRO stage. While some results rely on proprietary cluster configurations, the empirical rigor and scale of the experiments are exemplary.
Reproducibility
High. The methodology is fully integrated into the open-source DeepSpeed library, which provides production-ready implementations, detailed documentation, and extensive configuration examples. The paper supplies precise algorithmic descriptions, memory equations, and communication scheduling logic. Although exact hardware (V100 clusters) and network topologies influence absolute throughput, the core partitioning logic is hardware-agnostic and has been successfully reproduced and extended across diverse GPU/TPU infrastructures by both academia and industry.
Limitations
The primary limitation is communication overhead at extreme scales; ZeRO-3's parameter partitioning increases all-to-all communication volume, which can become a bottleneck on clusters with limited interconnect bandwidth. The paper does not deeply address mixed-precision training stability, checkpointing overhead for partitioned states, or dynamic workload balancing. Additionally, the super-linear speedup claims are partially dependent on specific cache and network effects that may not generalize to all cluster topologies. Later community extensions (e.g., ZeRO-Infinity, ZeRO-Offload) were required to address CPU/NVMe offloading and MoE integration.
Broader impact
ZeRO fundamentally democratized large-scale model training by removing the steep engineering barrier of manual model parallelism. It directly enabled the modern LLM era by making 10B–100B+ parameter training feasible on commodity GPU clusters. The methodology has become a foundational primitive in nearly all contemporary training frameworks (PyTorch FSDP, Megatron-DeepSpeed, Hugging Face Accelerate), accelerating research velocity across academia and industry while reducing the carbon and financial costs of large-scale training. ZeRO introduces a systematic memory partitioning strategy that eliminates redundancy in data-parallel training, fundamentally reshaping distributed ML infrastructure and enabling the practical training of trillion-parameter models. By rigorously formalizing state, gradient, and parameter partitioning, the work bridges the gap between algorithmic simplicity and extreme-scale efficiency, establishing a new standard for distributed training that has been universally adopted across the field and directly catalyzed the rapid scaling of modern foundation models.
Curator note
Rajbhandari et al.; Microsoft; partitioned optimizer state / gradients / params
Author of this paper? Send a correction.