ZeRO Optimization › Summary

ZeRO eliminates the memory redundancy that makes naive data-parallel training wasteful at scale. Rather than replicating all training state across every GPU, it partitions the state so that each GPU permanently owns a $1/N_d$ shard and fetches the rest on demand. The total model state in mixed-precision Adam training is $16\Psi$ bytes, of which standard data parallelism replicates every byte across all devices. ZeRO progressively removes this redundancy layer by layer.

The three stages target progressively larger memory categories, each building on the previous:

• Stage 1 partitions optimizer states (Adam moments and 32-bit master parameters), reducing the largest memory component from $12\Psi$ to $12\Psi / N_d$ per GPU, at no extra communication cost over baseline data parallelism.
• Stage 2 additionally partitions gradients using Reduce-Scatter, saving another $2\Psi$ bytes per GPU with the same total communication volume as Stage 1 and baseline data parallelism.
• Stage 3 partitions the parameters themselves, achieving a perfect $16\Psi / N_d$ per-GPU memory footprint with 3x the communication of Stages 1 and 2. Each GPU fetches parameters on demand via All-Gather before each layer's computation.

Extensions like ZeRO-Offload and ZeRO-Infinity push state to CPU RAM and NVMe storage respectively, enabling models with trillions of parameters on practical hardware. ZeRO++ addresses Stage 3's communication overhead through quantized weight transmission, hierarchical partitioning that exploits within-node NVLink bandwidth, and quantized gradient communication.

The key practical insight is that ZeRO produces mathematically identical results to standard data parallelism. The partitioning is a pure implementation detail that changes where tensors live and what each GPU computes, without affecting the final parameter values or training dynamics. This correctness guarantee, combined with the linear memory reduction with GPU count under Stage 3, makes ZeRO the standard foundation for large-scale distributed training. FSDP in PyTorch brings these ideas natively into the PyTorch ecosystem, making ZeRO-style training accessible without requiring the full DeepSpeed infrastructure stack.