ZeRO Optimization › ZeRO-Offload and ZeRO-Infinity

The original three ZeRO stages partition state across GPU memory. Two extensions push this further by moving state off the GPU entirely, exploiting the memory hierarchy of modern servers.

ZeRO-Offload extends Stage 2 by offloading optimizer states and the optimizer step computation to CPU memory and CPU compute. The central observation is that the optimizer step, the Adam update computing new parameter values from gradients and moments, is a simple element-wise operation that does not require the GPU's massive parallelism. A CPU can execute it efficiently at lower cost. CPU RAM is also far more abundant: a modern server might have 1-2 TB of DRAM while the attached GPUs have 80-160 GB total.

The challenge with naively offloading to CPU is that the data transfer between GPU and CPU over PCIe is slow, typically 16-32 GB/s bidirectional, versus 2-4 TB/s for GPU HBM. ZeRO-Offload addresses this through pipelining: while the GPU runs the forward and backward passes for training step $t+1$, the CPU runs the optimizer step for step $t$. If the optimizer step takes longer than the GPU computation, there is a wait at the end of each step, but for large models where the optimizer state is large but the per-parameter update is cheap, the CPU work finishes before the GPU needs the updated parameters.

ZeRO-Offload reduces GPU memory requirements by the full $12\Psi$ optimizer state bytes, replacing them with CPU RAM. This allows training larger models on a single GPU or small cluster than would otherwise be possible. For a single A100 (80 GB), ZeRO-Offload can train models significantly larger than what fits in 80 GB of GPU HBM alone.

ZeRO-Infinity extends the storage hierarchy further, to NVMe SSD storage. Modern NVMe drives offer 4-8 GB/s of sequential bandwidth and multiple terabytes of capacity. ZeRO-Infinity uses a three-tier heterogeneous memory hierarchy: NVMe for the largest tensors (optimizer states and inactive parameters), CPU RAM as a staging buffer for data moving between NVMe and GPU, and GPU HBM for active computation.

The key algorithmic challenge for ZeRO-Infinity is bandwidth management. GPU HBM operates at 2-4 TB/s; CPU RAM at 50-100 GB/s; NVMe at 4-8 GB/s. The system must carefully schedule which tensors live at which tier and when to prefetch data up the hierarchy to avoid stalling the GPU. ZeRO-Infinity uses a bandwidth-optimal data movement engine that analyzes the compute graph to schedule prefetches with enough lead time that the GPU is never blocked waiting for data from NVMe.

ZeRO-Infinity can train models with trillions of parameters on hardware configurations that would otherwise require a datacenter-scale GPU cluster. The original paper demonstrated training a 1 trillion-parameter model on a single DGX-2 node (16 V100 GPUs with 32 GB each), a configuration with roughly 0.5 TB of total GPU memory. Without ZeRO-Infinity, this would be impossible.

The tradeoff in both offloading approaches is latency and throughput. NVMe access is orders of magnitude slower than HBM, and the achievable training throughput depends critically on how well the prefetching hides the storage latency. In practice, ZeRO-Infinity achieves good throughput for compute-bound workloads (large models with large batch sizes) but can stall on memory-bound operations or when storage bandwidth becomes the bottleneck.