ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

Abstract

In the last three years, the largest dense deep learning models have grown over 1000x to reach hundreds of billions of parameters, while the GPU memory has only grown by 5x (16 GB to 80 GB). Therefore, the growth in model scale has been supported primarily though system innovations that allow large models to fit in the aggregate GPU memory of multiple GPUs. However, we are getting close to the GPU memory wall. It requires 800 NVIDIA V100 GPUs just to fit a trillion parameter model for training, and such clusters are simply out of reach for most data scientists. In addition, training models at that scale requires complex combinations of parallelism techniques that puts a big burden on the data scientists to refactor their model. In this paper we present ZeRO-Infinity, a novel heterogeneous system technology that leverages GPU, CPU, and NVMe memory to allow for unprecedented model scale on limited resources without requiring model code refactoring. At the same time it achieves excellent training throughput and scalability, unencumbered by the limited CPU or NVMe bandwidth. ZeRO-Infinity can fit models with tens and even hundreds of trillions of parameters for training on current generation GPU clusters. It can be used to fine-tune trillion parameter models on a single NVIDIA DGX-2 node, making large models more accessible. In terms of training throughput and scalability, it sustains over 25 petaflops on 512 NVIDIA V100 GPUs(40% of peak), while also demonstrating super linear scalability. An open source implementation of ZeRO-Infinity is available through DeepSpeed, a deep learning optimization library that makes distributed training easy, efficient, and effective.

References in corpus (13)

• Language Models are Few-Shot Learners
• Scaling Laws for Neural Language Models
• Mixed Precision Training
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Training Deep Nets with Sublinear Memory Cost
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
• Supporting Very Large Models using Automatic Dataflow Graph Partitioning
• PyTorch Distributed: Experiences on Accelerating Data Parallel Training
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• Distributed Hierarchical GPU Parameter Server for Massive Scale Deep Learning Ads Systems
• Memory-Efficient Pipeline-Parallel DNN Training
• Mesh-TensorFlow: Deep Learning for Supercomputers
• Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

Cited by in corpus (17)

• On the Opportunities and Risks of Foundation Models
• Sustainable AI: Environmental Implications, Challenges and Opportunities
• PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
• OneFlow: Redesign the Distributed Deep Learning Framework from Scratch
• Hanayo: Harnessing Wave-like Pipeline Parallelism for Enhanced Large Model Training Efficiency
• HongTu: Scalable Full-Graph GNN Training on Multiple GPUs (via communication-optimized CPU data offloading)
• Maximizing Parallelism in Distributed Training for Huge Neural Networks
• M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
• Pre-Trained Models: Past, Present and Future
• M6-T: Exploring Sparse Expert Models and Beyond
• 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
• 8-bit Optimizers via Block-wise Quantization
• A Short Study on Compressing Decoder-Based Language Models
• Varuna: Scalable, Low-cost Training of Massive Deep Learning Models
• Is the Number of Trainable Parameters All That Actually Matters?
• AutoFL: Enabling Heterogeneity-Aware Energy Efficient Federated Learning