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

By Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, Yuxiong He. 2021

548 citations recorded by the source corpus.

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 11DeepSpeed (https://www.deepspeed.ai/) is a deep learning optimization library designed to make distributed training easy, efficient, and effective. DeepSpeed has been extensively adopted by the DL community..

Subfields: Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

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

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

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).”
Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model (2022, methodology): “We refer to Rajbhandari et al. [56] for an analytical study of memory consumption during training.”
Pathways: Asynchronous Distributed Dataflow for ML (2022, methodology): “Large models (Brown et al., 2020) can be limited by available accelerator memory, requiring GPU memory virtualization (Rhu et al., 2016; Ausavarungnirun et al., 2018) or DRAM ofﬂoad (Rajbhandari et al., 2021).”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, methodology): “We note that we have only considered ZeRO-3 without tensor parallelism.”
Reducing Activation Recomputation in Large Transformer Models (2022, methodology): “An alternative to model parallelism is to combine a number of training techniques along with data parallelism to enable large scale model training [14, 17, 15, 16].”
Sustainable AI: Environmental Implications, Challenges and Opportunities (2021, background): “The strong performance scaling demand for ML motivates a variety of scale-out solutions [11], [18] by leveraging parallelism at scale with a massive collection of training accelerators.”
8-bit Optimizers via Block-wise Quantization (2021, background): “Optimizer sharding can also have significant communication overhead (Rajbhandari et al., 2021).”
XLS-R: Self-supervised Cross-lingual Speech Representation Learning at Scale (2021, background): “Training sets comprise be-tween 30 and 76 hours of annotated data.”
Pre-Trained Models: Past, Present and Future (2021, background): “Megatron-LM (Shoeybi et al., 2019; Narayanan et al., 2021) and Deep Speed (Rajbhandari et al., 2020, 2021) target at training large-scale NLP PTMs. Insight Face (ins, 2021) trains large-scale face recognition models.”
DeepSpeed- Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale (2022, methodology): “ZeRO-Inference builds on the offloading techniques of ZeRO-Infinity [35], and adapts them to inference.”
OneFlow: Redesign the Distributed Deep Learning Framework from Scratch (2021, background): “Megatron-LMs (Shoeybi et al., 2020; Narayanan et al., 2021) and DeepSpeed (dee, 2021; Rajbhandari et al., 2021; 2020) support model parallelism for pre-training large NLP models.”
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].”
DHEN: A Deep and Hierarchical Ensemble Network for Large-Scale Click-Through Rate Prediction (2022, methodology): “To solve this issue, we use fully sharded data parallel (FSDP [3, 32]) to remove memory redundancy in traditional data parallelism by further sharding weights to different GPUs, activation checkpointing to trade more compute for less peak memory usage [5], cpu offloading [31] to further reduce GPU m”
VeriGen: A Large Language Model for Verilog Code Generation (2023, methodology): “We set the training hyperparameters to their defaults as recommended for the ZeRO-3 [26] optimizer state implementation.”
Parallel Training of Pre-Trained Models via Chunk-Based Dynamic Memory Management (2021, methodology): “Although successful pre-training efforts exist on models of trillions of parameters [33, 34, 38], it still has a long way to go before making PTM accessible to everyone.”
nuQmm: Quantized MatMul for Efficient Inference of Large-Scale Generative Language Models (2022, background): “Correspondingly, even though new innovative memory architectures (e.g., HBM [39]) are proposed, the maximum memory capacity for a single GPU is still limited up to a few tens of gigabytes [16].”
Scaling Speech Technology to 1, 000+ Languages (2023, methodology): “To make efficient use of GPU memory, we use a fully sharded backend [Rajbhandari et al., 2021] as well as activation checkpointing [Chen et al.”