Current DeepSpeed ZeRO-3 partitions each parameter as flattened intra-layer slices across ranks

The ZeRO-3 Diagram Most People Remember Is Wrong

The ZeRO-3 Diagram Most People Remember Is Wrong Many engineers first learned ZeRO-3 from an animation that looked like pipeline parallelism. One GPU owned early layers, another GPU owned later layers, and the active layer block appeared to be broadcast to all other GPUs when needed. That picture is memorable. It is also not the right mental model for current DeepSpeed ZeRO-3 training. The current steady-state model is intra-layer partitioning. Each parameter is flattened, padded if necessary, and split across data-parallel ranks. Before forward or backward compute needs that parameter, ranks AllGather the full parameter. After gradients are produced, ranks ReduceScatter gradients back to the owning shards. That is the operational reading of ZeRO-3 from the ZeRO paper (arXiv:1910.02054) and the current DeepSpeed runtime. ...

June 16, 2025 · 8 min · Duo An
ZeRO stages partition optimizer state, gradients, and parameters

ZeRO: Partitioning Optimizer State, Gradients, and Parameters

ZeRO: Partitioning Optimizer State, Gradients, and Parameters Plain DDP is clean but wasteful. Every data-parallel rank stores the same parameters, gradients, fp32 master weights, and optimizer moments. At small scale that redundancy is convenient. At LLM scale it is the memory wall. ZeRO, the Zero Redundancy Optimizer, keeps data-parallel semantics and removes the redundant storage one category at a time. ZeRO-1 shards optimizer state. ZeRO-2 shards optimizer state and gradients. ZeRO-3 shards optimizer state, gradients, and parameters (arXiv:1910.02054). ...

April 27, 2025 · 9 min · Duo An