[Paper Review] ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

Samyam Rajbhandari, Olatunji Ruwase+3 more|arXiv (Cornell University)|Apr 16, 2021

Advanced Neural Network Applications30 references57 citations

TL;DR

ZeRO-Infinity offloads model states and activations to CPU/NVMe with memory-centric tiling to train trillion-parameter models on limited GPU resources, achieving high throughput and accessibility.

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.

Motivation & Objective

• Explain the GPU memory wall challenges for extreme-scale DL and motivate a system that increases reachable model size without refactoring.
• Introduce ZeRO-Infinity’s heterogeneous memory offload and tiling techniques to fit massive models on existing clusters.
• Demonstrate training efficiency and scalability on very large models, including trillion-parameter scales.
• Assess usability improvements that reduce or remove the need for model refactoring or complex parallelism.“
• Provide open-source implementation details and discuss implications for future hardware design.

Proposed method

• Propose infinity offload engine to simultaneously utilize GPU, CPU, and NVMe memory.
• Introduce memory-centric tiling to break large operators into tiles for reduced working memory.
• Define bandwidth-centric partitioning to exploit aggregate memory bandwidth across devices.
• Develop overlap-centric design to hide communication with computation across heterogeneous memory.
• Offer ease-inspired implementation to avoid model refactoring while training arbitrary architectures.
• Characterize memory and bandwidth requirements for large-model training with Transformer-based models.

Experimental results

Research questions

• RQ1 How can we break the GPU memory wall to train models with hundreds of trillions of parameters on practical GPU clusters?
• RQ2 Can we make large models accessible to researchers who lack massive GPU clusters by leveraging CPU and NVMe memory?
• RQ3 Is it possible to train trillion-parameter models without refactoring model code or adopting complex 3D parallelism?
• RQ4 What are the memory and bandwidth characteristics required to sustain efficient training with heterogeneous memory?
• RQ5 How does ZeRO-Infinity perform in terms of throughput, scalability, and accessibility on current hardware?

Key findings

• ZeRO-Infinity can train models with 32 trillion parameters on 32 NVIDIA DGX-2 nodes (512 GPUs).
• It sustains over 25 petaflops of throughput on the same hardware (40% of peak).
• It demonstrates superlinear scalability for a trillion-parameter model.
• It enables fine-tuning up to a trillion parameters on a single DGX-2 node without model parallelism or refactoring.
• Memory-centric tiling reduces per-operator working memory, enabling extremely large layers to fit without model parallelism.
• An open-source implementation of ZeRO-Infinity is available in DeepSpeed.