ZeRO-Offload: Democratizing Billion-Scale Model Training

Abstract

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.

References in corpus (7)

• Adam: A Method for Stochastic Optimization
• Language Models are Few-Shot Learners
• Scaling Laws for Neural Language Models
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Training Deep Nets with Sublinear Memory Cost
• Training Large Neural Networks with Constant Memory using a New Execution Algorithm
• Efficient Memory Management for GPU-based Deep Learning Systems

Cited by in corpus (15)

• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
• Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training
• M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
• Pre-Trained Models: Past, Present and Future
• Federated Fine-Tuning of LLMs on the Very Edge: The Good, the Bad, the Ugly
• M6-T: Exploring Sparse Expert Models and Beyond
• TensorTEE: Unifying Heterogeneous TEE Granularity for Efficient Secure Collaborative Tensor Computing
• ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training
• Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading
• Breaking the Memory Wall: A Study of I/O Patterns and GPU Memory Utilization for Hybrid CPU-GPU Offloaded Optimizers
• Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training
• Distributed Deep Learning in Open Collaborations
• Long-term series forecasting with Query Selector -- efficient model of sparse attention
• Pinpointing the Memory Behaviors of DNN Training
• JUWELS Booster -- A Supercomputer for Large-Scale AI Research