ZeRO-Offload: Democratizing Billion-Scale Model Training

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

• Model states—including parameters, gradients, and optimizer states—are the primary memory bottleneck in large-model training.Source: §2 · p. 3
• Mixed-precision Adam training stores fp16 and fp32 parameters, fp16 gradients, and fp32 optimizer states, producing memory requirements beyond flagship GPU capacity.Source: §2 · p. 3

Scale out large model training

• Scale-out methods use multiple GPUs by partitioning model and residual states across devices.Source: §2 · p. 3
• ZeRO partitions model states across GPUs instead of replicating them and does not require user-model changes.Source: §2 · p. 3
• ZeRO-Offload extends this approach by offloading model states to CPU memory, enabling larger single-GPU models and compatibility with ZeRO and model parallelism.Source: §2 · p. 3

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.Source: §2 · p. 3
• L2L manages GPU memory layer by layer but incurs extra communication overhead and lacks the same scaling path described for ZeRO-Offload.Source: §2 · p. 4

• Model-state offloading must balance CPU computation, GPU–CPU communication, and GPU memory savings.Source: §3 · p. 4, §3.1 · p. 4

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.Source: §3 · p. 4, §3.1 · p. 4

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.Source: §3.2 · p. 5

3.3 Minimizing Communication Volume

• The strategy minimizes CPU–GPU communication because PCIe bandwidth is slower than CPU and GPU memory bandwidth.Source: §3.3 · p. 5
• Fp32 states are co-located with parameter-update and float-to-half operations, forming the Update Super node in the reduced graph.Source: §3.3 · p. 5

3.4 Maximizing Memory Savings

• 8x maximum memory savings is achieved by offloading fp16 gradients and the Update Super node to the CPU.Source: §3.4 · p. 5

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.Source: §3.5 · p. 5
• The resulting strategy is unique because greater memory savings would require either more CPU computation or additional GPU–CPU communication.Source: §3.5 · p. 6

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.Source: §4.1 · p. 6
• Gradients are transferred individually or in small groups after computation, while forward propagation requires no CPU communication because fp16 parameters are already on GPU.Source: §4.1 · p. 6

Scaling to Multi-GPUs

• ZeRO-Offload preserves ZeRO Stage-2 optimizer-state and gradient partitioning while offloading partitioned data and parameter updates to CPU.Source: §4.2 · p. 6
• Partitioning assigns each data-parallel process a parameter subset, keeping aggregate GPU–CPU communication constant while parallel CPU resources reduce total update time.Source: §4.2 · p. 6
• ZeRO-Offload also combines with tensor-slicing model parallelism to train larger models than model parallelism alone.Source: §4.2 · pp. 6–7

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.Source: §5 · p. 7
• The CPU optimizer uses SIMD vector instructions, loop unrolling, and OpenMP multithreading to exploit CPU hardware parallelism and memory bandwidth.Source: §5.1 · p. 7
• Adam maintains gradients, first and second momentums, and FP32 master parameters on CPU while FP16 parameters support GPU forward computation.Source: §5.1 · p. 7
• The tiled CPU-to-GPU parameter copy overlaps Adam computation with transfer, reducing GPU idle time before the next training step.Source: §5.1 · p. 8

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.Source: §5.2 · p. 8
• Delayed updates introduced after a few dozen iterations achieve the same model training accuracy as ZeRO-Offload alone with higher training throughput.Source: §5.2 · p. 8

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.Source: §6.2.1 · p. 9
• 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.Source: §6.2.1 · p. 9

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.Source: §6.2.2 · p. 10
• 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.Source: §6.2.2 · p. 10

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.Source: §6.2.3 · p. 10

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.Source: §6.2.4 · p. 11

• ZeRO-Offload delivers high compute efficiency and near-linear throughput scalability for multi-billion-parameter model training.Source: §7 · p. 11
• ZeRO-Offload enables data scientists to train multi-billion-parameter models on a single GPU without model refactoring.Source: §7 · p. 11
• The authors open-sourced ZeRO-Offload through the DeepSpeed library to help data scientists harness massive deep-learning models.Source: §7 · pp. 11–12