ZeRO Optimization › ZeRO++ and Communication Efficiency

ZeRO++ (published in 2023) addresses Stage 3's 3x communication overhead through three complementary improvements. Each targets a specific phase of the training step and uses different compression techniques.

The All-Gather during the forward pass communicates the parameters for each layer before computing with them. ZeRO++ compresses these parameters from 16-bit to 8-bit before transmission, halving the communication volume for this phase. After reception, each GPU dequantizes the parameters back to 16-bit (or bf16) before computation. The dequantization step takes negligible time compared to the actual layer computation.

The quantization uses a block-wise scheme where parameters are divided into small blocks, each with its own scale factor. This limits the quantization error relative to a global scale factor scheme, preserving model accuracy during the forward pass. The gradients computed from quantized parameters accumulate small errors, but these errors are comparable in magnitude to the noise already present from stochastic gradient descent and mini-batch sampling, so training stability is maintained.

Standard ZeRO Stage 3 distributes parameters uniformly across all $N_d$ GPUs, regardless of physical topology. This means that when a layer's All-Gather is performed, parameters must travel across slow inter-node InfiniBand links for any parameters owned by GPUs on other nodes.

ZeRO++ introduces hierarchical partitioning that exploits the two-level topology of modern clusters. Within a single node, NVLink provides 600 GB/s of bandwidth. Between nodes, InfiniBand provides 200-400 Gb/s. Under hpZ, parameters are partitioned only within each node: each GPU owns $1/(N_d / N_\text{nodes})$ of the parameters within its node, but every node holds a complete copy. The forward-pass All-Gather is then a within-node operation using NVLink, avoiding the slow inter-node link entirely.

This does reintroduce some parameter replication (each node stores a complete copy), which somewhat reduces the memory savings of Stage 3. However, on multi-node training runs where inter-node bandwidth is the bottleneck, hpZ can dramatically improve throughput. The memory tradeoff is often acceptable because within-node NVLink bandwidth is used much more efficiently, and the remaining memory pressure can be addressed through other means.

The Reduce-Scatter for gradients transfers $2\Psi$ bytes per step. ZeRO++ applies 1-bit quantization to these gradients before the Reduce-Scatter, reducing the gradient communication by up to 16x (from 16-bit to 1-bit). The 1-bit quantization uses the same error-feedback mechanism as 1-bit Adam: quantization errors are accumulated and added to the next gradient, so no information is permanently lost, just deferred.

This is the most aggressive of the three ZeRO++ optimizations and the one with the most potential for convergence impact. In practice, the gradient quantization tends to work well when the gradient signal is strong (early training) and can be problematic when gradients become very small (fine-grained fine-tuning). The qgZ component is typically used selectively, activated during phases where gradient magnitudes are large.

Together, qwZ, hpZ, and qgZ can reduce Stage 3's communication overhead by 4x or more, largely closing the gap with Stages 1 and 2 in terms of achieved hardware utilization. ZeRO++ makes Stage 3 practical on a wider range of cluster configurations, particularly those with slower inter-node interconnects.