ZeRO++: Extremely Efficient Collective Communication for Giant Model Training
This addresses efficiency issues for researchers and practitioners training giant models on distributed systems, representing an incremental improvement over ZeRO.
The paper tackles the high communication volume bottleneck in ZeRO for training large language models on low-bandwidth or large-scale GPU clusters, resulting in a 4x reduction in communication volume and up to 2.16x better throughput at 384 GPU scale.
Zero Redundancy Optimizer (ZeRO) has been used to train a wide range of large language models on massive GPUs clusters due to its ease of use, efficiency, and good scalability. However, when training on low-bandwidth clusters, or at scale which forces batch size per GPU to be small, ZeRO's effective throughput is limited because of high communication volume from gathering weights in forward pass, backward pass, and averaging gradients. This paper introduces three communication volume reduction techniques, which we collectively refer to as ZeRO++, targeting each of the communication collectives in ZeRO. First is block-quantization based all-gather. Second is data remapping that trades-off communication for more memory. Third is a novel all-to-all based quantized gradient averaging paradigm as replacement of reduce-scatter collective, which preserves accuracy despite communicating low precision data. Collectively, ZeRO++ reduces communication volume of ZeRO by 4x, enabling up to 2.16x better throughput at 384 GPU scale.