ZeRO-Offload

Modeling mixed-precision Adam training on M parameters (16M bytes of model states: 2M each for p_16 and g_16, 4M each for p_32, m_32, and v_32) as a weighted directed cyclic data-flow graph where O(M B) FWD-BWD compute resides on GPU implies that any model-state offload cut must cross at least two directed cycle edges of weight >= 2M, yielding a theoretical minimum GPU-CPU communication volume of 4M bytes per step (2M g_16 offload + 2M p_16 swap); among all partitions achieving 4M communication, co-locating {g_16, p_32, m_32, v_32, ParamUpdate, float2half} on CPU and {p_16, FWD-BWD} on GPU uniquely achieves the maximum 8x GPU model-state memory reduction (from 16M to 2M steady / 4M buffered, Table 1).

By integrating with ZeRO Stage-2 across N_dp data-parallel GPUs, each GPU reduce-scatters and offloads only its 1 / N_dp partition of FP16 gradients to CPU and updates only its 1 / N_dp shard of FP32 optimizer states in parallel via a SIMD AVX512 / loop-unrolled / OpenMP CPU-Adam kernel before an all-gather on GPU, keeping aggregate GPU-CPU communication volume constant at 4M bytes regardless of N_dp while scaling total CPU optimizer compute linearly with N_dp.

To hide CPU optimizer overhead at small batch sizes, One-Step Delayed Parameter Update (DPU) trains the first N - 1 steps normally, skips the CPU update at step N, and for step i >= N + 1 executes the CPU parameter update using step i - 1 gradients concurrently with the GPU forward-backward pass using step i - 1 parameters (Figure 6), yielding 1.12x-1.59x higher throughput while preserving GPT-2 pretraining loss and BERT-Large SQuAD F1 (92.8), and enabling 13B-parameter training on a single V100 GPU at 40 TFLOPS and 70B-parameter training on a 16-GPU DGX-2 node.

Because ZeRO-Offload builds on ZeRO-2 and keeps a full replica of FP16 parameters p_16 on every GPU, its maximum trainable model size without model parallelism is bounded by the parameter capacity of a single GPU's memory, and at large GPU counts (e.g., 64-128 GPUs) where ZeRO-2 fits in aggregate GPU memory, ZeRO-Offload's PCIe transfer overhead makes it slightly slower than GPU-only ZeRO-2.