Concepts
ZeRO: partitioning training state across GPUs
Intuition
Imagine a team of 64 people in which each person carries a complete copy of the same encyclopedia, even though each one only consults one page at a time. Instead of everyone carrying everything, each person can keep just one piece and ask colleagues for the page they need at the moment of use. Training a large model with many GPUs suffers from the same waste: each GPU holds a whole copy of things it only uses in part, and the space runs out before the model fits.
How it works
When training with the Adam optimizer in mixed precision, the paper accounts for, for Ψ parameters, 2Ψ bytes of fp16 parameters, 2Ψ bytes of fp16 gradients and KΨ bytes of optimizer states (fp32 copy of the parameters, momentum and variance, with K = 12), for a total of 16Ψ bytes; the optimizer states are usually the largest share. In basic data parallelism (DP), each device keeps a complete replica of all this, so DP is efficient in computation and communication but does not reduce per-GPU memory. ZeRO-DP eliminates this redundancy by partitioning the model states across the Nd data-parallel processes, in three cumulative stages. In the first (Pos), the optimizer states are split into Nd parts, each process updates only its own and an all-gather distributes the updated parameters, which reduces memory to 4Ψ + 12Ψ/Nd (up to 4x). In the second (Pos+g), each gradient is reduced only on the process that owns the partition, via reduce-scatter, and then released, reaching about 2Ψ (up to 8x). In the third (Pos+g+p), the parameters are also partitioned and whoever needs a partition owned by another process receives it by broadcast, so memory drops to 16Ψ/Nd, linear in the DP degree. As for communication, standard DP moves 2Ψ per step (reduce-scatter plus all-gather, each with Ψ), and Pos+g keeps the same 2Ψ; Pos+g+p, however, needs two parameter all-gathers (forward and backward) plus the gradient reduce-scatter, adding up to 3Ψ, that is, 1.5x the baseline. The residual states (activations, buffers, fragmentation) are handled separately by ZeRO-R, which, for example, partitions the activation checkpoints (used in activation checkpointing) across the model-parallelism GPUs.
In practice
In the paper's example, a model with 7.5B parameters and Nd = 64 takes up 120 GB per device in standard DP, about 31.4 GB with Pos, 16.6 GB with Pos+g and about 1.9 GB with Pos+g+p. The evaluated implementation, ZeRO-100B (Pos+g plus ZeRO-R, in PyTorch), trained models of up to 170B parameters on 400 V100 GPUs combined with model parallelism, with more than 38 Tflops per GPU at 100B; without model parallelism, it trained models of up to 13B, whereas PyTorch's DDP runs out of memory at 1.4B. The paper also reports super-linear scaling from 64 to 400 GPUs on a 60B model, because less memory spent on states allows larger batches per GPU, and presents ZeRO as part of the open-source DeepSpeed library. The Pos+g+p stage was still planned work at the time, and the authors acknowledge that training 1 trillion parameters would run into a lack of compute.
Use it in a design
Practice the cases that rely on this concept in the System Design school.