PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, Yuxiong He

arXiv:2104.07857v1cs.DCcs.AIcs.LGcs.PF

TL;DR

Rapid model growth is nearing the GPU memory wall, while existing large-model training requires costly clusters and substantial model refactoring. ZeRO-Infinity uses heterogeneous GPU, CPU, and NVMe memory with memory and bandwidth optimizations to expand scale and preserve efficiency. It demonstrates unprecedented model scale and argues that slow but abundant memory can support efficient training on current GPU clusters.

  • Problem

    Model size has grown over 1000x while single-GPU memory grew 5x, creating GPU-memory and accessibility limits for large-model training.

  • Method

    ZeRO-Infinity combines GPU, CPU, and NVMe memory with memory-centric tiling, bandwidth-centric partitioning, overlap-centric execution, and an ease-inspired implementation.

  • Results

    ZeRO-Infinity demonstrates unprecedented model scale and excellent training efficiency on current-generation GPU clusters.

  • Takeaways & Limitations

    Large-model training need not fit entirely in expensive GPU memory; heterogeneous CPU and NVMe memory can provide the aggregate bandwidth needed for efficient training.

  • Takeaways & Limitations

    External parameters used across module boundaries make automated data movement difficult for some model architectures.

Abstract

from arXiv · show

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.

1 EXTENDED INTRODUCTION

Rapid model growth is approaching GPU-memory and usability limits, while 3D parallelism demands substantial resources and model refactoring. ZeRO-Infinity addresses these challenges through heterogeneous memory, tiling, bandwidth-aware execution, and automated implementation.

  • 1000x model growth has outpaced single-GPU memory growth of 5x, from 16 GB to 80 GB.
  • 3D parallelism can scale to over one trillion parameters on 800 NVIDIA V100 GPUs but requires aggregate GPU memory and complex parallelism.
  • The GPU memory wall makes trillion-parameter training and fine-tuning inaccessible to many researchers and companies without massive clusters.Fine-tuning GPT-3 would require over 8 DGX-2 nodes with 3D parallelism, although one node has enough compute for reasonable-time fine-tuning.
  • 3D parallelism also requires major model-code refactoring, tensor slicing, and load-balanced pipeline stages.
  • ZeRO-Infinity combines CPU and NVMe offload with memory-centric tiling to support massive models on limited GPU resources.Memory-centric tiling handles individual layers that do not fit in GPU memory without model parallelism.
  • Bandwidth-centric partitioning, communication overlap, high-performance NVMe access, and automated data movement target efficient training without manual model refactoring.

2 BACKGROUND AND RELATED WORK

Prior large-model systems combine parallelism and memory optimization to scale training, but they impose refactoring, GPU-memory, or batching constraints. ZeRO-Infinity builds on these approaches while targeting heterogeneous memory limitations.

  • Data, model, and pipeline parallelism split training across devices, while 3D parallelism combines all three for trillion-parameter scaling.
  • 3D parallelism requires model-code refactoring, load-balanced pipeline stages, and sufficient total GPU memory.
  • ZeRO partitions optimizer states, gradients, and parameters across data-parallel processes to remove memory redundancy while retaining computational granularity.
  • ZeRO-Offload stores gradients and optimizer states in CPU memory but retains replicated parameters in GPU memory, limiting model scale.It also requires large batch sizes because of suboptimal data partitioning and limited PCIe bandwidth.
  • Activation checkpointing reduces activation memory by trading it for 0.33x additional recomputation.
  • Adam maintains first- and second-order statistics for each parameter and gradient, creating a significant memory footprint.
  • Mixed-precision training performs forward and backward propagation in FP16 and parameter updates in FP32.This exploits tensor-core acceleration on modern GPUs.

3 MEMORY REQUIREMENTS

Training memory comprises model states, residual activation states, and GPU working memory, with requirements that grow sharply for very large Transformer models. The analysis motivates offloading and techniques for handling large operators and activations.

  • Transformer training memory includes model states, residual activation states, and minimum GPU working memory needed during computation.
  • Adam mixed-precision training requires 20 bytes per parameter for model states.
  • 64 GPUs are required to fit model states for a 100B-parameter model, while a trillion-parameter model requires over 512 GPUs.
  • Activation checkpointing reduces activation memory at the expense of 0.33x additional recomputation.
  • Activation working memory becomes large beyond 10 trillion parameters even with one Transformer block between checkpoints.
  • Model State Working Memory is the GPU memory needed for the largest operator after model states are offloaded.It must hold that operator’s parameters and gradients for backward propagation.
  • Beyond 100 billion parameters, model-state working memory can require multiple gigabytes of contiguous memory.Memory-centric tiling is presented as an alternative to splitting operators across GPUs.
  • Activation working memory consists of activations between checkpoints and avoids contiguous-memory issues when its total fits in GPU memory.

4 BANDWIDTH REQUIREMENTS

The section quantifies how arithmetic intensity and available bandwidth determine training efficiency for model states and activation checkpoints. It finds sharply different bandwidth requirements across these data types.

  • Efficiency model: The efficiency metric estimates training efficiency from peak throughput, data-movement bandwidth, and arithmetic intensity without compute–communication overlap.Arithmetic intensity is total computation divided by required data movement; higher intensity lowers the required bandwidth.
  • Efficiency model: Model states and activation checkpoints have different bandwidth requirements because their arithmetic intensity depends on different workload parameters.Model-state requirements depend on batch size and sequence length, whereas activation-checkpoint requirements depend on checkpoint frequency and hidden dimension.
  • Scope: The bandwidth analysis is a single NVIDIA V100 DGX-2 data point that isolates each state or activation while assuming infinite bandwidth for the others.Results vary with the chosen achievable peak throughput.
  • Bandwidth requirements: Over 70 GB/s for parameters and gradients achieves over 50% efficiency even at the smallest batch size.At this bandwidth, data movement can theoretically be completely overlapped with computation.
  • Bandwidth requirements: Nearly 1.5 TB/s is required for 90% efficiency with batch size 2 per GPU for optimizer states.Optimizer states require nearly four times the bandwidth needed by parameters and gradients for 50% efficiency because their updates cannot overlap with computation.
  • Bandwidth requirements: With activation checkpointing, 2 GB/s sustains over 50% efficiency at hidden size 2K, while the requirement falls below 1 GB/s above hidden size 8K.These estimates use an achievable peak throughput of 70 TFlops/GPU for NVIDIA V100 DGX-2 workloads.

5 ZERO-INFINITY DESIGN OVERVIEW

ZeRO-Infinity combines heterogeneous GPU, CPU, and NVMe memory with partitioned state movement to overcome GPU-memory limits. Its design supports much larger models while using offload and memory-centric execution.

  • Architecture: ZeRO-Infinity exploits heterogeneous GPU, CPU, and NVMe memory to transcend the GPU memory wall.Modern clusters provide NVMe storage over 50x larger than GPU memory and nearly 20x larger than CPU memory.
  • Scale: One NVIDIA V100 DGX-2 node supports one trillion parameters, a 50x increase over 3D parallelism.The comparison is presented as the maximum achieved model size in Figure 1.
  • Offload engine: The infinity offload engine can keep partitioned model states on GPU or offload them to CPU or NVMe according to memory requirements.ZeRO-Infinity is built on ZeRO-3, which partitions optimizer states, gradients, and parameters.
  • Execution flow: During backward propagation, partitioned parameters move from slow memory to GPU, form the full layer, and then gradients are aggregated, repartitioned, and offloaded.Figure 4 depicts this sequence across four data-parallel ranks.
  • Activation offload: Activation checkpoints can also be offloaded to CPU memory, allowing models with hundreds of trillions of parameters to fit their checkpoints.The paper notes that 3 TB of checkpoints for a 100-trillion-parameter model is within reach of next-generation CPU memory.
  • Memory-centric tiling: Memory-centric tiling breaks large operators into sequential smaller tiles, reducing working memory in proportion to the number of tiles.With ZeRO-3, each tile’s parameters and gradients can be fetched and released one at a time.

5.2 Design for Excellent Training Efficiency

ZeRO-Infinity addresses the bandwidth cost of heterogeneous offload through bandwidth-centric partitioning, communication overlap, and optimized NVMe access. These mechanisms target the distinct bandwidth demands of parameters, optimizer states, and activations.

  • Bandwidth challenge: CPU memory bandwidth is an order of magnitude below GPU bandwidth, while NVMe bandwidth is another order of magnitude below CPU bandwidth.Reading and writing these memories from GPU is slower still, making efficient offload difficult.
  • Bandwidth requirements: Training requires over 70 GB/s, 1.5 TB/s, and 1–4 GB/s for parameters and gradients, optimizer states, and activation checkpoints, respectively.These thresholds are stated for a DGX-2 system.
  • Existing limitations: A single GPU’s 12 GB/s PCIe bandwidth is insufficient for heterogeneous training at scale under existing parameter-broadcasting approaches.Such approaches require large per-GPU batches, which can exceed activation memory and produce overly large effective batch sizes.
  • Bandwidth-centric partitioning: ZeRO-Infinity uses bandwidth-centric partitioning to retrieve offloaded parameters and gradients through parallel data mapping and access.This strategy is intended to provide virtually unlimited heterogeneous memory bandwidth.
  • Communication overlap: An overlap-centric design overlaps GPU–GPU computation with NVMe–CPU and CPU–GPU communication over PCIe.This reduces the impact of limited heterogeneous-memory bandwidth during training.
  • NVMe optimization: The infinity offload engine approaches peak NVMe bandwidth while overlapping NVMe reads, NVMe writes, and optimizer computation.It also reuses temporary buffers to reduce memory fragmentation and maintain efficiency with modest or small batch sizes.
  • Activation offload: On DGX-2, activation checkpoints can be offloaded at about 3 GB/s per GPU while retaining over 80% efficiency for hidden size above 8K.Checkpoint communication can also be overlapped with forward and backward computation.

5.3 Design for Ease of Use

ZeRO-Infinity removes the need for manual model refactoring by automating parameter movement and model partitioning within PyTorch. This supports models that exceed individual GPU or CPU memory.

  • Model adaptation: Memory-centric tiling avoids requiring tensor-slicing model parallelism for large individual layers.It reduces GPU memory requirements while preserving the operator’s mathematical equivalence.
  • Automation: ZeRO-Infinity’s PyTorch implementation eliminates manual model code refactoring when scaling to trillions of parameters.The implementation provides automated data movement and automated model partitioning during initialization.
  • Automated data movement: Pre- and post-forward/backward hooks automatically gather required parameters, then partition parameters and gradients and optionally offload them.These hooks perform the required data movement around PyTorch submodule execution.
  • Automated model partitioning: During initialization, parameters are partitioned and offloaded immediately after creation so the full model is never instantiated on one data-parallel process.This allows models exceeding single-GPU or single-CPU memory to be initialized without manual partitioning.

6 EFFICIENCY OPTIMIZATIONS

ZeRO-Infinity addresses heterogeneous-memory bandwidth limits through partitioned parameter retrieval and overlap-centric data movement. These optimizations keep communication from becoming a bottleneck, including on smaller configurations, while DeepNVMe approaches device bandwidth limits.

  • 6.1 Bandwidth-Centric Partitioning: Partitioning individual parameters across data-parallel processes lets ZeRO-Infinity use allgather instead of broadcast for CPU or NVMe-resident data.This activates all PCIe links in parallel, unlike the broadcast-based approach that first moves a complete parameter through one PCIe link.
  • 6.1 Bandwidth-Centric Partitioning: 48/25 GB/s is the effective CPU/NVMe-to-GPU bandwidth achieved by allgather on a DGX-2, versus about 12 GB/s for broadcast with 16-way data parallelism.The values correspond respectively to CPU and NVMe bandwidth and are limited by aggregate PCIe and per-node NVMe bandwidth.
  • 6.1 Bandwidth-Centric Partitioning: Over 3 TB/s of CPU memory bandwidth and over 1.5 TB/s of NVMe bandwidth are available to ZeRO-Infinity across 64 DGX-2 nodes.The bandwidth grows linearly with additional nodes, providing heterogeneous-memory bandwidth for large-scale training.
  • 6.2 Overlap-Centric Design: An overlap engine concurrently hides GPU communication, NVMe-to-CPU transfers, CPU-to-GPU transfers, and computation.A dynamic prefetcher prepares parameters before forward or backward use, while gradient movement overlaps with backward computation.
  • 6.3 DeepNVMe: DeepNVMe approaches peak sequential NVMe read and write bandwidth through parallel I/O, scheduling, copy avoidance, and memory pinning.Its asynchronous bulk requests can overlap storage operations with communication or computation.

7 EASE INSPIRED IMPLEMENTATION

ZeRO-Infinity automates tensor movement through PyTorch integration, while handling cross-module parameters and initialization of models too large for traditional replication. It preserves a standard data-parallel usage model and reduces initialization memory requirements.

  • 7 EASE INSPIRED IMPLEMENTATION: ZeRO-Infinity is designed for PyTorch use without model code refactoring, similarly to standard data-parallel training.It coordinates partitioning, offloading, GPU residency, and later re-partitioning of parameters, gradients, and optimizer states.
  • 7.1 Automating Data Movement: Recursive module hooks gather parameters before forward computation and automate their movement; the backward pass is handled similarly.The hooks block until parameters are available and rely on overlap-centric communication to minimize stalls.
  • 7.1 Automating Data Movement: External parameters cross module boundaries, as when GPT shares embedding weights at the network’s beginning and end, complicating automatic gathering.ZeRO-Infinity supports manual registration and mechanisms that automatically detect and register such parameters.
  • 7.2 Partitioned Initialization: A 500 billion parameter model would require 8 TB of aggregate CPU or GPU memory for initial replication on an eight-GPU node.Traditional data parallel initialization replicates the model before partitioning, exceeding the node’s available memory.
  • 7.2 Partitioned Initialization: Immediate per-layer partitioning during initialization reduces the 500 billion parameter example to 1 TB of aggregate CPU memory.Only individual submodules are fully initialized, so the full model is never replicated across data-parallel processes.

8 EVALUATION

The evaluation shows that ZeRO-Infinity substantially expands trainable model size while maintaining strong throughput, scalability, and usability on limited GPU resources. Its system features jointly enable trillion-scale training without model parallelism or code refactoring.

  • Model Size and Speed: 32T parameters versus about 650B with 3D parallelism yields a 50x increase in model scale.ZeRO-Infinity trains up to 20T-parameter models on 512 GPUs, reaching up to 49 TFlops/GPU.
  • Superlinear Scalability: Super-linear scalability is achieved from 64 to 512 GPUs when training a 1T-parameter model.The result leverages increased aggregate PCIe and NVMe bandwidth and additional CPU compute.
  • Democratizing Large Model Training: On one 16-GPU DGX-2 node, models up to 100B parameters exceed 40 TFlops/GPU, while 3D parallelism cannot scale beyond 20B.ZeRO-Infinity can therefore support fine-tuning models such as GPT-3 on a single node.
  • Impact of System Features on Model Scale: On one DGX-2 node, successive partitioning and offloading steps increase capacity from 1.4B with data parallelism to 1T with NVMe offload.The final increase is 700x relative to data parallelism alone.
  • Impact of System Features on Model Scale: Memory-centric tiling increases the trainable hidden size from 8K to 64K using a tiling factor of 16.This avoids the need for model parallelism for large hidden sizes.
  • Impact of System Features on Training Speed: At 64 GPUs, aggregate PCIe bandwidth produces nearly 2x faster back propagation than ZeRO-Offload for an 8B-parameter model.Prefetching and communication overlap are crucial at small per-GPU batch sizes, while activation-checkpoint offloading has minimal impact at 32K and 64K hidden sizes.

9 CONCLUSION & FUTURE IMPLICATIONS

ZeRO-Infinity uses GPU, CPU, and NVMe memory to overcome GPU-memory limits while preserving accessible and efficient large-model training. Future systems still need substantially more compute and device-to-device bandwidth for reasonable-time training at extreme scale.

  • Conclusion: ZeRO-Infinity leverages GPU, CPU, and NVMe memory to achieve unprecedented model scale and training efficiency.Its design uses slow but large memory in parallel across devices to provide aggregate bandwidth.
  • Future Implications: With accelerators offering 10x V100 compute, 512-device clusters require 30 GB/s between each accelerator and slow memory to remain efficient.The paper notes that this bandwidth is already possible with current NVLink technology.
  • Limitations and Future Requirements: Training tens or hundreds of trillions of parameters in reasonable time still requires massive compute increases and proportional device-to-device bandwidth.ZeRO-Infinity removes accelerator device memory as a model-scale limitation, but does not remove these future resource requirements.

A APPENDIX

The appendix provides model configurations for the evaluation figures, covering model size, hidden size, GPU count, parallelism, layers, attention heads, and batch sizes. The listed configurations span models from millions or billions to 50B parameters and include GPU ranges up to 128.

  • Figure 6(a): Figure 6(a) configurations vary hidden size, GPU count, model-parallel layers, model size, attention heads, and batch sizes.Listed model sizes range from 900M to 50B parameters.
  • Figure 6(b): Figure 6(b) configurations include an 8B model with hidden size 8192, 10 layers, 16 attention heads, and GPU ranges from 4 to 64.The configurations use batch size 2 and total batch sizes from 8 to 128.
Loading 2104.07857v1…