Papers
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He · SC 2020 (arXiv 1910.02054), 2020
The problem: where memory goes during training
The authors start from a paradox: a GPT-2 with 1.5 billion parameters takes up only 3 GB of weights in 16-bit precision, yet it does not train on a 32 GB GPU with TensorFlow or PyTorch. They split the consumption into two groups. The first is the model states: optimizer states, gradients, and parameters. The second is the residual states: activations, temporary buffers, and fragmented memory. For Adam in mixed-precision training, the paper accounts, for Ψ parameters, for 2Ψ bytes of fp16 parameters and 2Ψ bytes of fp16 gradients. In addition, the optimizer keeps an fp32 copy of the parameters, the momentum, and the variance, 4Ψ bytes each, which the paper groups as KΨ with K = 12. The total is 2Ψ + 2Ψ + KΨ = 16Ψ bytes, that is, at least 24 GB for the 1.5B model, and the optimizer states are usually the largest share. Among the residuals, the paper cites that the activations of that same model, with sequence length 1K and batch size 32, require about 60 GB. Activation checkpointing reduces this to about 8 GB at the cost of 33% recomputation. A fused fp32 buffer for 1.5B parameters takes up 6 GB, and fragmentation, which makes a memory request fail when there is no large enough contiguous block, was found to cause out-of-memory errors with more than 30% still free in extreme cases.
The limits of data parallel and model parallel
In basic data parallelism (DP), each device holds a full replica of the model; the mini-batch is split, each process computes forward and backward on its slice, and the gradients are averaged. The result is good computation and communication efficiency, but no per-device memory reduction: according to the paper, basic DP runs out of memory on 32 GB GPUs with models above 1.4B parameters. Model parallelism (MP) cuts each layer vertically across devices and thereby partitions the states; in exchange, the computation becomes too fine-grained and requires heavy communication between layers. It works well within a node, where the bandwidth between GPUs is high, and degrades outside it. The authors tested a 40B model with Megatron-LM on two DGX-2 nodes and observed about 5 Tflops per V100 GPU, less than 5% of peak. Later they cite the bandwidth drop from 300 GB/s per link (NVSwitch) to 12.5 GB/s (Infiniband EDR) when leaving the node. Pipeline parallelism, which slices the model horizontally by layers, also has costs: GPipe requires a batch proportional to the number of partitions, and PipeDream keeps stale copies of parameters. From there the paper draws three observations. DP scales better than MP. DP wastes memory by replicating everything. And not all state is needed all the time: a layer's parameters are only required during its forward and backward passes.
ZeRO-DP: three stages of partitioning the model states
The central idea of ZeRO-DP is to eliminate the redundancy of DP by partitioning the model states across the Nd data-parallel processes instead of replicating them, while keeping the computation granularity of DP. The three stages are cumulative. In the first (Pos), the optimizer states are split into Nd equal parts; each process stores and updates only its own and, at the end of the step, an all-gather distributes the updated parameters. Consumption drops from 16Ψ to 4Ψ + 12Ψ/Nd, close to 4Ψ for large Nd: a reduction of up to 4x. In the second (Pos+g), each gradient is reduced only on the process that owns that partition, via reduce-scatter with buckets, and then released, reaching about 2Ψ, a reduction of 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; memory drops from 16Ψ to 16Ψ/Nd, linear in the DP degree. In the paper's example, a 7.5B model with Nd = 64 needs 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. With Nd = 64, the stages accommodate models of up to 7.5B, 14B and 128B; with Nd = 1024 and all stages, the paper states that 1 trillion parameters fit. For that trillion, the total state would be about 16 TB, or 16 GB per GPU across 1024 GPUs.
The communication cost: free up to Pos+g, 1.5x with Pos+g+p
The natural question is whether ZeRO trades memory for communication. The paper answers by measuring the volume of data moved per process at each step. In standard DP, the all-reduce of the gradients is done in two stages, reduce-scatter and all-gather, each moving Ψ elements, for a total of 2Ψ per step. With Pos+g, the gradient needs only a reduce-scatter (volume Ψ), and then each process updates its partition and an all-gather collects the updated parameters (another Ψ). The total remains 2Ψ, the same as standard DP, even with a memory reduction of up to 8x. With Pos+g+p, the all-gather of the parameters has to happen twice: once spread across the forward pass, with the parameters discarded after use, and once in the backward pass, in reverse order. Added to the reduce-scatter of the gradients, the total is 3Ψ, that is, 1.5x the baseline, in exchange for memory proportional to Nd. The paper also analyzes ZeRO-R: in Megatron-LM with activation checkpointing, each transformer block performs six all-reduces (two in the forward, two in the recomputation, and two in the backward) and moves 12 × seq × hidden, and activation partitioning (Pa) adds only one all-gather of seq × hidden per block, less than 10% of the original MP volume. Since Pa reduces activation memory in proportion to the MP degree, the batch can grow and the DP communication volume, which is inversely proportional to the batch, can fall.
ZeRO-R, results and legacy
For residual states, ZeRO-R has three pieces. The first is partitioning activation checkpoints across the MP GPUs, which in conventional MP are replicated; in the paper's example (100B, batch 32, sequence 1024, MP 16), the cost of about 33 GB per GPU drops to about 2 GB, and those 2 GB can still be offloaded to the CPU. The second is constant-size buffers, which keep the fused buffer from growing with the model (for 3B parameters, a 32-bit buffer would need 12 GB). The third is defragmentation during training, with contiguous blocks preallocated for checkpoints and gradients. In the implementation, called ZeRO-100B (Pos+g plus ZeRO-R, in PyTorch), the paper reports, on 400 V100 GPUs, models of up to 170B parameters combined with MP, above 15 Petaflops aggregate and more than 38 Tflops per GPU at 100B. It 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. Without MP, models of up to 13B can be trained, whereas PyTorch's DDP runs out of memory at 1.4B. Pos+g+p was still planned work, and the authors admit that a trillion parameters would run into a lack of compute (more than a year of training). The paper cites Turing-NLG, at 17B, as a model made possible by ZeRO, and presents it as part of the open-source DeepSpeed library.
See the ideas at work
Practice the designs that use this paper in the System Design school.