Zero Redundancy Optimizer (ZeRO)

In plain data parallelism every device keeps a full copy of the model weights, the gradients, and the optimizer states (for Adam, the running mean and variance — twice the parameter count in fp32). That is enormously wasteful: across a hundred devices you store a hundred identical copies of everything. ZeRO's insight is that those copies are redundant. Instead of replicating the full optimizer state, ZeRO partitions it so each rank owns only its slice, and the ranks reconstruct what they momentarily need through communication.

ZeRO comes in three cumulative stages. Stage 1 shards the optimizer states across data-parallel ranks; stage 2 additionally shards the gradients; stage 3 also shards the parameters themselves, gathering each layer's weights with an all-gather just before its forward and backward pass and releasing them afterward. Memory per device thus falls roughly as the model state divided by the data-parallel degree, at the cost of extra all-gather and reduce-scatter traffic. ZeRO keeps the simple programming model of data parallelism — no model surgery — while reaching parameter counts that would never fit replicated.

ZeRO trades memory for communication volume, so its efficiency hinges on fast interconnects and on overlapping the gathers with compute. Stage 3 in particular issues an all-gather per layer, which can stall on slower fabrics. It composes with offload (moving shards to CPU or NVMe) and with tensor or pipeline parallelism inside each model-parallel group, and it is the conceptual core of FSDP.