Rajbhandari et al. (2021) — ARXIV-2104.07857
Key takeaways
ZeRO-Infinity extends the ZeRO family with an infinity offload engine that simultaneously leverages GPU, CPU, and NVMe memory to train models with tens of trillions of parameters on current GPU clusters—50x larger than 3D parallelism—without requiring model code refactoring.
- Bandwidth-centric partitioning scaling. Effective NVMe/CPU-to-GPU bandwidth scales linearly with data-parallel degree: 48 GB/s aggregate at 16-way DP vs. 12 GB/s for single-PCIe broadcast. Holds when: 16-GPU data-parallel setup; CPU offloading; allgather vs. broadcast comparison.
- Extreme-scale training capacity. 32 trillion parameter models can be trained on 512 V100 GPUs, 50x larger than what 3D parallelism supports on identical hardware. Holds when: 32 DGX-2 nodes (512 V100 GPUs); ZeRO-Infinity with full NVMe offloading.
- Throughput at scale. Over 25 petaflops (40% of peak) sustained throughput on 512 V100 GPUs during extreme-scale training. Holds when: 512 V100 GPUs; large transformer models with high arithmetic intensity.
What the evidence does not cover
- NVMe bandwidth is still the ultimate bottleneck for the optimizer update step; achieving >1.5 TB/s aggregate NVMe bandwidth requires many nodes.
- Training at tens-of-trillions scale requires enormous compute (exa-flop class systems) even if memory is no longer the binding constraint.
- Activation checkpoint offloading reduces efficiency for small hidden sizes (<8K); the technique is only bandwidth-efficient for large transformer layers.
- The overlap-centric design assumes a relatively stable operator execution order; dynamic computation graphs may reduce prefetching effectiveness.
Standing in the anthology
Read — the reading is NOTE-122. Arrived in the imported batch, which brought in the memory-partitioning branch.