ZeRO

Data parallelism keeps a complete copy of the optimizer state on every worker, which is the same numbers stored N times. ZeRO shards them instead: each worker holds a slice, and the pieces are gathered when needed. The memory per worker falls almost linearly with the number of workers.

Three stages shard the optimizer state, then the gradients, then the parameters themselves. Each stage saves more memory and adds more communication, so the choice is a memory-for-bandwidth trade rather than a free win. Stage 3 makes models far larger than one device trainable without model parallelism.

With Ψ parameters, N workers and K bytes of optimizer state per parameter, memory per worker goes from 2Ψ + 2Ψ + KΨ to roughly 2Ψ + 2Ψ + KΨ/N at stage 1 and (2 + 2 + K)Ψ/N at stage 3.