ZeRO-Offload: Democratizing Billion-Scale Model Training

By Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyang Yang, Minjia Zhang, Yuxiong He. 2021

618 citations recorded by the source corpus.

Large-scale model training has been a playing ground for a limited few requiring complex model refactoring and access to prohibitively expensive GPU clusters. ZeRO-Offload changes the large model training landscape by making large model training accessible to nearly everyone. It can train models with over 13 billion parameters on a single GPU, a 10x increase in size compared to popular framework such as PyTorch, and it does so without requiring any model change from the data scientists or sacrificing computational efficiency. ZeRO-Offload enables large model training by offloading data and compute to CPU. To preserve compute efficiency, it is designed to minimize the data movement to/from GPU, and reduce CPU compute time while maximizing memory savings on GPU. As a result, ZeRO-Offload can achieve 40 TFlops/GPU on a single NVIDIA V100 GPU for 10B parameter model compared to 30TF using PyTorch alone for a 1.4B parameter model, the largest that can be trained without running out of memory. ZeRO-Offload is also designed to scale on multiple-GPUs when available, offering near linear speedup on up to 128 GPUs. Additionally, it can work together with model parallelism to train models with over 70 billion parameters on a single DGX-2 box, a 4.5x increase in model size compared to using model parallelism alone. By combining compute and memory efficiency with ease-of-use, ZeRO-Offload democratizes large-scale model training making it accessible to even data scientists with access to just a single GPU.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

In Optimization & Training

Citations in the corpus

PaLM: Scaling Language Modeling with Pathways (2022, background): “Many other works aim to increase of the scale of models, while limiting communication overheads (Rajbhandari et al., 2020; Lepikhin et al., 2020; Li et al., 2020; Rasley et al., 2020; Rajbhandari et al., 2021; Ren et al., 2021; Narayanan et al., 2021a).”
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep learning (2021, methodology): “line for each experiment we use the relevant state-of-the-art method among 3D Parallelism [13], ZeRO [11], or ZeRO-Offload [12].”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, methodology): “Another way to train large models is to offload “non-active” parameters to the CPU, as in ZeRO-Offload [29].”
Reducing Activation Recomputation in Large Transformer Models (2022, methodology): “As we detail in Section 2, there are several other techniques to reduce the memory requirements of training large models, such as partitioning various data across the data parallel ranks or oﬄoading data to CPU memory [14, 17].”
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, methodology): “Zero-Redundancy parallelism [26, 27, 29] shards parameters as well but communicates parameters on-demand to recover the unsharded form and executes the model as if it were replicated on every device.”
Efficient Memory Management for Large Language Model Serving with PagedAttention (2023, methodology): “Swapping [23, 42, 55], recomputation [7, 24] and their combination [40] have been utilized to reduce the peak memory of training.”
Fast Inference of Mixture-of-Experts Language Models with Offloading (2023, methodology): “This technique works particularly well when processing large batches of data, during training Pudipeddi et al. (2020); Ren et al. (2021) or large-batch non-interactive inference Aminabadi et al. (2022); Sheng et al. (2023), where each layer processes a lot of tokens each time the layer is loaded…”
Pre-Trained Models: Past, Present and Future (2021, background): “As shown in Figure 9, some works such as ZeRO-Offload (Ren et al., 2021) design delicate strategies to schedule the swap between the CPU memory and the GPU memory so that memory swap and device computation can be overlapped as much as possible.”
TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models (2021, background): “TeraPipe can also be combined with various memory optimization techniques, e.g., gradient accumulation (Fan et al., 2020), rematerialization (Chen et al., 2016; Jain et al., 2019), or memory swapping (Ren et al., 2021).”
This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (neutral): “In addition to parallelization, there are other complementary techniques for training large-scale models, such as memory optimization [12, 14,21,23,28,46], communication compression [6,53], and low-precision training [35].”
FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models (2022, background): “ZeRO Offload [28] reduces memory consumption of data parallelism by swapping data to host memory, and the data is further off-loaded to disks [26].”
Gemma 3 Technical Report (2025, positive): “The optimizer state is sharded using an implementation of ZeRO-3 (Ren et al., 2021).”
A Survey of Large Language Models for Healthcare: from Data, Technology, and Applications to Accountability and Ethics (2023, positive): “Originally, with a single V100, a 1.4B model can be trained with PyTorch and the throughput is 30 TFLOPS, and with ZeRO-Offload augmentation, a 10B model can be trained and the throughput is 40 TFLOPS.”
Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging (2022, background): “Paging: Huang et al. (2020) and Ren et al. (2021) page activations off a memory-scarce GPU to the CPU when out of memory.”
What Makes Good Data for Alignment? A Comprehensive Study of Automatic Data Selection in Instruction Tuning (2023, methodology): “To facilitate parallel training, we employ DeepSpeed Zero-Stage 3 (Ren et al., 2021) and FlashAttention-2 (Dao, 2023).”
High-throughput Generative Inference of Large Language Models with a Single GPU (2023, background): “Within broader domains, memory optimizations and offloading have been studied for training (Huang et al., 2020; Ren et al., 2021; Steiner et al., 2022) and linear algebra (Jia-Wei & Kung, 1981; Demmel, 2013).”
LMFlow: An Extensible Toolkit for Finetuning and Inference of Large Foundation Models (2023, methodology): “LMFlow developed an easy-to-use inference interface for LLMs. Based on Deepspeed 3 , LMFlow supports parameter partitioning with zero-offload strategies as introduced by [28].”