PaperlayerView exact revision on arXiv

Source-linked AI summary

ZeRO-Offload: Democratizing Billion-Scale Model Training

Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, Yuxiong He

arXiv:2101.06840v1cs.DCcs.LG

TL;DR

Large-model training is difficult to access because existing methods require expensive multi-GPU systems and heterogeneous approaches lack suitable support for large attention-based workloads. ZeRO-Offload partitions training between GPU and CPU, then combines offloading with ZeRO to scale across GPUs. It reports 10x larger models on a single GPU, including up to 13B parameters, while preserving efficiency and avoiding model refactoring.

  • Problem

    Large-model training requires expensive multi-GPU systems, while prior heterogeneous methods lack support for large attention-based workloads and efficient multi-GPU scaling.

  • Method

    ZeRO-Offload offloads selected model states, optimizer computation, and updates to CPU memory and compute while keeping forward/backward computation on GPUs and integrating with ZeRO data parallelism.

  • Results

    10x larger models can be trained on a single GPU, with up to 13B parameters, while maintaining compute efficiency and near-linear multi-GPU throughput scalability.

  • Takeaways & Limitations

    ZeRO-Offload makes multi-billion-parameter training accessible on single-GPU systems without requiring model refactoring.

  • Takeaways & Limitations

    CPU computation can become a bottleneck with very small batch sizes, requiring one-step delayed parameter updates to overlap CPU and GPU computation.

Abstract

from arXiv · show

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.

1 Introduction

ZeRO-Offload targets the growing memory and accessibility barriers of large-model training by using CPU memory and compute without model refactoring. Its design combines efficiency, scalability, and usability to support substantially larger models on single or multiple GPUs.

  • 175B parameters marks the scale reached by attention-based language models, whose growth has been accompanied by continued accuracy improvements.
  • Large-model training remains costly because existing distributed approaches require multiple GPUs and complex system technologies to exceed single-device memory limits.
  • Existing heterogeneous-training methods primarily target smaller CNN workloads, offload memory without CPU computation, and lack clear multi-GPU scaling paths.
  • ZeRO-Offload uses CPU memory and compute while combining with ZeRO-powered data parallelism for efficient scaling.
  • 13B parameters can be trained on one NVIDIA V100 GPU at 40 TFLOPS, versus 1.2B parameters at 30 TFLOPS without CPU offloading.
  • ZeRO-Offload preserves efficiency through CPU-optimizer optimizations and delayed updates, while its PyTorch implementation requires only a few lines of code changes.
  • The contributions include a unique offload strategy, scalable multi-GPU and model-parallel integration, optimized CPU execution, open-source PyTorch support, and extensive evaluation.

2 Background and Related Work

Large-model training is constrained by model-state memory requirements that exceed single-GPU capacity. Prior approaches distribute or compress computation and memory, but heterogeneous methods have limited scope or scalability.

  • Model states—including parameters, gradients, and optimizer states—are the primary memory bottleneck in large-model training.
  • Mixed-precision Adam training stores fp16 and fp32 parameters, fp16 gradients, and fp32 optimizer states, producing memory requirements beyond flagship GPU capacity.
  • Scale out large model training: Scale-out methods use multiple GPUs by partitioning model and residual states across devices.
  • Scale out large model training: ZeRO partitions model states across GPUs instead of replicating them and does not require user-model changes.
  • Scale out large model training: ZeRO-Offload extends this approach by offloading model states to CPU memory, enabling larger single-GPU models and compatibility with ZeRO and model parallelism.
  • Scale up large model training: Scale-up methods reduce activation memory through checkpoint recomputation, reduce state and activation memory through precision changes, or use external memory.
  • Scale up large model training: L2L manages GPU memory layer by layer but incurs extra communication overhead and lacks the same scaling path described for ZeRO-Offload.

3 Unique Optimal Offload Strategy

ZeRO-Offload derives a unique GPU–CPU partitioning of the training data-flow graph that limits CPU computation, minimizes communication, and maximizes GPU memory savings. It keeps forward/backward computation and fp16 parameters on the GPU while offloading gradients, fp32 states, and updates to the CPU.

  • Model-state offloading must balance CPU computation, GPU–CPU communication, and GPU memory savings.
  • 3.1 DL Training as a data-flow graph: ZeRO-Offload uses first-principles analysis of a weighted data-flow graph to partition computation and data between CPU and GPU.
  • 3.2 Limiting CPU computation: Forward and backward propagation remain on the GPU because each has O(MB) complexity, while lower-complexity computations may be offloaded to the CPU.
  • 3.3 Minimizing Communication Volume: The strategy minimizes CPU–GPU communication because PCIe bandwidth is slower than CPU and GPU memory bandwidth.
  • 3.3 Minimizing Communication Volume: Fp32 states are co-located with parameter-update and float-to-half operations, forming the Update Super node in the reduced graph.
  • 3.4 Maximizing Memory Savings: 8x maximum memory savings is achieved by offloading fp16 gradients and the Update Super node to the CPU.
  • 3.5 A unique and optimal offload strategy: ZeRO-Offload stores fp32 states and fp16 gradients on the CPU while keeping fp16 parameters and forward/backward computation on the GPU.
  • 3.5 A unique and optimal offload strategy: The resulting strategy is unique because greater memory savings would require either more CPU computation or additional GPU–CPU communication.

4 ZeRO-Offload Schedule

ZeRO-Offload places GPU-resident parameters and CPU-resident optimizer data to minimize communication, then combines this schedule with ZeRO partitioning for scalable multi-GPU training and model parallelism.

  • Single GPU Schedule: ZeRO-Offload stores fp16 parameters on GPUs while offloading fp16 gradients, optimizer states, and parameter updates to CPU memory.The single-GPU schedule transfers gradients after backward computation and copies updated parameters back to GPU memory.
  • Single GPU Schedule: Gradients are transferred individually or in small groups after computation, while forward propagation requires no CPU communication because fp16 parameters are already on GPU.
  • Scaling to Multi-GPUs: ZeRO-Offload preserves ZeRO Stage-2 optimizer-state and gradient partitioning while offloading partitioned data and parameter updates to CPU.
  • Scaling to Multi-GPUs: Partitioning assigns each data-parallel process a parameter subset, keeping aggregate GPU–CPU communication constant while parallel CPU resources reduce total update time.
  • Scaling to Multi-GPUs: ZeRO-Offload also combines with tensor-slicing model parallelism to train larger models than model parallelism alone.

5 Optimized CPU Execution

ZeRO-Offload accelerates CPU-side optimizer execution with a fast CPU Adam implementation and delayed updates that overlap CPU work with GPU computation, while preserving accuracy when needed.

  • CPU Optimizer: Two optimizations accelerate CPU parameter updates: a fast CPU Adam optimizer and one-step delayed updates that overlap CPU and GPU computation.The delayed schedule hides CPU execution time when enabled.
  • CPU Optimizer: The CPU optimizer uses SIMD vector instructions, loop unrolling, and OpenMP multithreading to exploit CPU hardware parallelism and memory bandwidth.
  • CPU Optimizer: Adam maintains gradients, first and second momentums, and FP32 master parameters on CPU while FP16 parameters support GPU forward computation.
  • CPU Optimizer: The tiled CPU-to-GPU parameter copy overlaps Adam computation with transfer, reducing GPU idle time before the next training step.
  • One-Step Delayed Parameter Update: Very small batch sizes can make CPU computation a training bottleneck, so delayed parameter updates overlap CPU and GPU work to hide this overhead.The schedule delays updates by one step and was evaluated as not affecting final accuracy.
  • One-Step Delayed Parameter Update: Delayed updates introduced after a few dozen iterations achieve the same model training accuracy as ZeRO-Offload alone with higher training throughput.

6 Evaluation

The evaluation measures model scale, throughput, scalability, and optimized CPU execution across single- and multi-GPU settings. ZeRO-Offload trains substantially larger models while maintaining strong throughput, near-linear scalability, and convergence behavior comparable to baseline training.

  • 6.2.1 Model scale: 13B parameters can be trained on one GPU with ZeRO-Offload, compared with 1.4B using PyTorch DDP.ZeRO-Offload offloads optimizer states and most gradients to CPU memory to maximize GPU memory savings.
  • 6.2.1 Model scale: 70B parameters can be trained on 16 GPUs in one DGX-2 node with ZeRO-Offload combined with model parallelism.This increases model scale by 50X, 4.5X, 7.8X, and 4.2X over PyTorch, Megatron, ZeRO-2, and L2L, respectively.
  • 6.2.2 Training throughput: 14% average throughput improvement over L2L is achieved by ZeRO-Offload, reaching up to 22% higher throughput.Its CPU-GPU communication volume is 4M rather than L2L’s 28M for a model with M parameters.
  • 6.2.2 Training throughput: 1B–15B models achieve up to 1.33X, 1.11X, and 1.64X higher speeds than PyTorch, ZeRO-2, and Megatron, respectively.Offloading optimizer states enables larger micro-batch sizes and higher throughput.
  • 6.2.3 Throughput Scalability: Near-perfect linear aggregate-throughput speedup is achieved through 128 GPUs while running above 30 TFlops per GPU.ZeRO-Offload makes training feasible from 1 to 16 GPUs where ZeRO-2 runs out of memory.
  • 6.2.4 Optimized CPU execution: 1.12–1.59 times higher throughput is achieved with one-step delayed parameter updates for small micro-batches, without hurting convergence or final accuracy.The delayed update overlaps CPU optimizer updates with the next GPU forward computation; BERT-Large reaches the same final F1 score of 92.8 as baseline.

7 Conclusions

ZeRO-Offload combines compute efficiency and near-linear throughput scalability to train multi-billion-parameter models on a single GPU without model refactoring. Its open-source DeepSpeed release aims to democratize large model training.

  • ZeRO-Offload delivers high compute efficiency and near-linear throughput scalability for multi-billion-parameter model training.
  • ZeRO-Offload enables data scientists to train multi-billion-parameter models on a single GPU without model refactoring.
  • The authors open-sourced ZeRO-Offload through the DeepSpeed library to help data scientists harness massive deep-learning models.
Loading 2101.06840v1…