Quantized LLM Training at Scale with ZeRO++ // Guanhua Wang // AI in Production 2025

Communication overhead is a primary bottleneck in large-scale language model training, especially in two common scenarios: training with small batch sizes and training in environments with limited internode bandwidth. As model sizes grow exponentially, more GPUs are required for training, which forces a smaller micro-batch size per GPU for a fixed global batch size. This, in turn, increases the proportion of time spent on communication, sometimes accounting for nearly 44% of the training latency. ZeRO++, a project from the Microsoft DeepSpeed team, addresses this challenge by significantly reducing the communication volume required during training.

Understanding the ZeRO-3 Communication Volume

The standard ZeRO-3 optimizer involves three major collective communication operations in each training iteration, with the total data volume being three times the model size (3M):

1. Forward all-gather on weights (M): Before the forward pass of a layer, each GPU gathers the full weights for that layer.
2. Backward all-gather on weights (M): This is repeated before the backward pass to reconstruct the full weights needed for gradient calculation.
3. Backward reduce-scatter on gradients (M): After the backward pass, gradients are reduced and scattered across GPUs for the parameter update.

ZeRO++ introduces three components to tackle each of these communication steps, reducing the total volume from 3M to 0.75M—a 4x reduction.

The Three Pillars of ZeRO++

1. QWZ (Quantized Weights in ZeRO): This component targets the forward all-gather on weights. Instead of naively quantizing weights, which can harm model convergence, QWZ uses block quantization. The weight tensor is divided into smaller chunks, and quantization is performed within each chunk. This approach improves data precision by 3.3x over global quantization. By communicating weights in 8-bit format instead of the standard 16-bit, QWZ cuts the communication volume for the forward pass in half, from M to 0.5M.

2. HPZ (Heterogeneous Partitioning in ZeRO): HPZ eliminates the backward all-gather communication across nodes. It achieves this by holding a secondary replica of the model partitions within each machine (node). This is a memory-for-communication trade-off. Since the full weights for a layer can be reconstructed using data already present on the node, the cross-node all-gather operation is no longer necessary. This reduces the communication volume for the backward weight gathering from M to zero.

3. QGZ (Quantized Gradients in ZeRO): This component is designed to reduce the reduce-scatter volume for gradients. Directly quantizing gradients for a standard reduce-scatter is problematic because the sequential dequantization and requantization steps (one for each GPU in the ring) introduce significant latency and precision loss. QGZ replaces the ring-based reduce-scatter with a novel hierarchical all-to-all protocol that solves several challenges:

   • Challenge: Sequential quantization overhead.
     • Solution: A one-shot quantization is performed on each GPU, followed by an all-to-all communication, reducing the number of sequential quantization steps from N (number of GPUs) to just one.
   • Challenge: Communication volume blow-up with a simple all-to-all.
     • Solution: A two-hop protocol is used. First, an intra-node all-to-all and reduction occurs, followed by an inter-node all-to-all. This prevents each GPU from sending its full quantized gradient block to every other GPU in the system.
   • Challenge: Data misplacement after the hierarchical communication.
     • Solution: A tensor slice reordering step is performed before any communication. Local gradient chunks are swapped so they land on the correct target GPU after the two-hop all-to-all is complete.

Further optimizations like kernel fusion and overlapping communication collectives are also applied. QGZ effectively reduces the gradient communication volume from M to 0.25M.

Performance and Adoption

ZeRO++ demonstrates significant performance gains.

• LLM Training: In low-bandwidth (100 Gbps) settings with over 300 GPUs, ZeRO++ achieves up to a 2.16x speedup over the baseline ZeRO-3.
• RLHF Training: For models like Llama 2 70B and OPT 30B, it delivered speedups of 3.3x and 3x, respectively.
• Model Convergence: Validation loss curves for GPT models up to 13B parameters show that ZeRO++ converges correctly, matching the baseline. For most models under 200B parameters, 8-bit quantization does not result in any accuracy loss.

The technology is already battle-tested. LinkedIn reported that ZeRO++ helped reduce training time by over 50% in their production AI workloads. Additionally, Microsoft’s own Phi-3 series of small language models were trained using ZeRO++ to improve throughput.

• Tags:
• Feature
• LLM Training
• Microsoft
• Distributed Training
• DeepSpeed
• ZeRO++
• Quantization
• Communication Bottleneck