ZeRO++: Extremely Efficient Collective Communication for Large Model Training
Authors: Guanhua Wang
Heyang Qin
Sam Ade Jacobs
Xiaoxia Wu
Connor Holmes
Zhewei Yao
Samyam Rajbhandari
Olatunji Ruwase
Feng Yan
Lei Yang
Yuxiong He
Venue: ICLR 2024 poster (2024)
Keywords: low-precision LLM pretraining
2 bits
auto compression
low memory pretraining
Primary area: infrastructure, software libraries, hardware, etc.
TL;DR: Efficient collective communication design for large model training

Abstract: Zero Redundancy Optimizer (ZeRO) has been used to train a wide range of large language models on massive GPU clusters due to its ease of use, efficiency, and good scalability. However, when training on low-bandwidth clusters, and/or when small batch size per GPU is used, ZeRO’s effective throughput is limited due to communication overheads. To alleviate this limitation, this paper introduces ZeRO++ composing of three communication volume reduction techniques (lowprecision all-gather, data remapping, and low-precision gradient averaging) to significantly reduce the communication volume up to 4x that enables up to 2.16x better throughput at 384 GPU scale. Our results also show ZeRO++ can speedup the RLHF by 3.3x compared to vanilla ZeRO. To verify the convergence of ZeRO++, we test up to 13B model for pretraining with 8/6-bits all gather and up to 30B model for finetuning with 4/2-bits all gather, and demonstrate on-par accuracy as original ZeRO (aka standard training). As a byproduct, the model trained with ZeRO++ is naturally weight-quantized, which can be directly used for inference without post-training quantization or quantization-aware training.

Published as a conference paper at ICLR 2024
ZERO++: EXTREMELY EFFICIENT COLLECTIVE COMMUNICATION FOR LARGE MODEL TRAINING
Guanhua Wang"! Heyang Qin", Sam Ade Jacobs', Xiaoxia Wu', Connor Holmes?, Zhewei Yao², Samyam Rajbhandari', Olatunji Ruwase', Feng Yan', Lei Yang', Yuxiong He Microsoft DeepSpeed', OpenAI,,
Snowflake',
University of Houston",
University of Nevada, Reno"
ABSTRACT
Zero Redundancy Optimizer (ZcRO) has been used to train a wide range of large language models on massive GPU clusters due its ease of use, efficiency. and good scalability. However, when training on low-bandwidth clusters, and/or when small batch size per GPU is used, ZeRO's effective throughput limited by communication overheads. To alleviate this limitation, this paper introduces ZeRO++ composing of three communication volume reduction techniques (low- precision all-gathez data remapping, and low-precision gradient averaging) to sig- nificantly reduce the communication volume up to 4x that enables up to 2.16x better throughput at 384 GPU scale. Our results also show ZeRO++ can speedup the RLHF training by 3.3x compared to vanilla ZeRO. To verify the convergence of ZeRO++ test up 13B model for pretraining with 8/6-bits all gather and up to 30B model for finetuning with 4/2-bits all gather, and demonstrate on-par accuracy as original ZeRO (aka standard training). As byproduct, the model trained with ZeRO++ is naturally weight-quantized, which can be directly used for inference without post- post-training quantization or quantization-avae aware training.
INTRODUCTION
The size of deep learning (DL) models has increased from 100 million to over 500+ billion pa- rameters, ranging from BERT (Devlin al.] 2018) to Megatron- Turing NLG (Smith et al. (022). With the increase in model size. the memory and compute requirements for training have increased significantly beyond the capability of a single accelerator (e-g., GPU). Training these massive models requires the efficient aggregation of computing power and memory across hundreds or even thousands of GPU devices. There are two popular approaches to alleviate this, namely 3D paral- lelism (Narayanan et al. 2021 Team & Majumder, 2020) and Zero Redundancy Optimizer (ZeRO)
Rajbhandari et al. 2020).
Compared to 3D parallelism, 7eRO is easier to use without model code refactoring. ZcRO memory efficient variation of data paral- lelism (Ben-Nun & Hoefler 2019 Dean et al. 2012) where model states are partitioned across all the GPUs, instead of being replicated, and reconstructed using gather based communica- tion collectives on-the-fly during training. This allows ZeRO to effectively leverage the ag- gregate GPU memory across machines, at the expense of minimal communication overhead (50%) compared to standard data parallel train- ing (2M vs 3M for model size of M) (Rajbhandari et al., 2020), while still achieving excellent throughput scalability (Rajbhandari et al. 2021).
GPUS
(e) Communicatlo Bardwisth Performance
(b1 Batch Sioe per CPU
TELOPs pur dP2
GPtc
SPM
cer GHA
Figure 1: Training throughput constrained net- work bandwidth and batch size per GPU.
However, the communication overhead of ZeRO can limit throughput two important scenarios (1) bow-bandwidth cluster and (2) low-compute training scenario, e.g., the batch size per GPU is
equal contribution

---

Published as a conference paper at ICLR 2024
small. We demonstrate these two cases in Figure As can be seen, the bandwidth of the cluster becomes smaller and/or the training batch size decreases, the training efficiency of ZeRO significantly reduces.
To overcome the communication overhead of ZeRO, we here present a novel system of communi- cation optimizations collectively called ZeRO++, for which the contribution can be summarized:
Quantized Weight Communication for ZeRO (qwZ). First, in order to reduce parameter commu- nication volume during forward all-gather (which is M for ZeRO), adopt quantization on weights to shrink down each model parameter from FP16 to lower-precision data type (e.g.. 8/6-bits) before communicating. To preserve decent model training precision, we adopt block-based quantization [Dettmers et al. 2022, Yao et al. 2022), which conducts independent quantization on each subset of model parameters. There no existing implementation for high performance block-based quantiza- tion. Thus, we implement highly optimized quantization CUDA kernels from scratch. Another great byproduct of qwZ that it automatically converts the model weight from half- malf-precision (16-bits) to lower precision for efficient inference without any inference quantization methods. Particularly, as compared to the current popular post-training quantization methods (Frantar et al.) 2022). qwZ can push the weight-precision to 2-bits for finetuned model without significant accuracy drop.
Hierarchical Weight Partition for ZeRO (hpZ.). Second, to reduce the communication overhead of all-gather weights during backward (which is M for ZeRO), we trade GPU memory for com- munication. More specifically, instead of spreading whole model weights across all the machines, maintain full model copy within each node. At the expense of higher memory overhead, this al- lows us to replace the expensive cross- machine all-gather on weights with intra-machine all-gather, which is substantially faster due to much higher intra-machine communication bandwidth. hqZ reduces the all-gather for backward from volume M to O.
Quantized Gradient Communication for ZeRO (ggZ.. Third, reducing the communication cost of gradients using reduce- scatter (which is M for ZeRO) is even more challenging. Directly applying quantization to reduce communication volume is infeasible due the training accuracy drop. To bet- ter preserve the accuracy, we propose a novel hierarchical all-to-all based gradient scatter-reduction schedule, where an intra-node gradient reduction first applied and then a inter-node quantized gradient communication (e.g., 8/4-bits) is used resulting in 2/4x communication volume reduction. To achieve the best outcome, we incorporate pipelining intra-node and inter-node communication and conducting CUDA kernel fusion.
By incorporating all three components above, we reduce the cross-node communication volume by 4x from 3MM down to less than 7.M (< 0.5M for forward, 0 for backward, and 0.25. for gradient sduue--satterr. We extensively test ZeRO++'s system performance, and show (i) scalability GPT- 3 like models up 384 GPUs achieving over 45% of sustained peak throughput, (ii) consistent speedup of up to 2.4x over ZeRO across models ranging from 10- (0-138B parameters, (iii) comparing with baseline in 4x higher bandwidth cluster, ZeRO++ achieve better throughput in wow-bandwidth setting, (iv) 3.3x better throughput for RLHF training and generation.
To verify the convergence of 7eRO++, we tested it both pretraining and fine-nuning. Our results show that ZeRO++ can (1) achieve on-par training loss as ZeRO for up to 13B GPT-3 models on pretraining with both 8-bits and 6-bits qwZ, and (2) realize similar (slightly worse) model quality as ZeRO for OPT-30B finetuning with 4-bits (2-bits) qwZ. The 2-bits qwZ trained model can be directly served for inference, which is significantly better than current methods (PTQ) (Frantar et al., 2022) while ZeRO++ even saves training cost and PTQ cost.
quantization
BACKGROUND AND RELATED WORK
This section only describes the core background needed for understanding the techniques discussed in this paper. A broader discussion of the related work is provided in Appendix A
2.1 ZERO OPTIMIZER
ZeRO is memory-optimizee solution for data parallel training. ZeRO partitions and distributes all model states (i.c., parameters, gradients, optimizer states) among GPUs in use and recollects model states only when the layer needs to be computed. There are three different stages for using ZeRO
2

---

Published as a conference paper at ICLR 2024
to optimize on-device memory usage. In ZeRO stage 1 (ZeRO-1), only optimizer states split and spread across all GPUs in use. ZeRO stage 2 (ZeRO-2) partitions both optimizer states and gradients, and ZeRO stage 3 (ZeRO-3) splits all three components of model states as parameters, gradients, and optimizer states.
Algorithm 1: ZeRO algorithm
ZeRO-3 is the most memory efficient solution for model training at large scale, but at the cost of more collective communications. Algorithm[ iillus- trates the high-level pseudocode for ZeRO-3. --0. Dur- ing model training. ZeRO-3 lazily schedules the fetching of parameters until the computation needs to happen on particular layer. Before forward propagation, ZeRO launches an all-gather to collect the full model weights and then computes the for-
Input model.worrlSize
Output: model
while model not converged do
Mat.ggherPPrraacctttii Battheer_PPaammteerrsssse -...ommttrssuooorSssrrrr all model. forward().
partition portission.worrlSiieer (worldSize);
ali gather Porometerss
model. backward();
partition (worldiize)
reduce seatter, (rateerrrrrreetttadtiSSiirr optsmizer.step))
end while Return: model
ward pass (line 2-3) of Algorithm | Then ZeRO empties the all-gather weights buffer after forward computation completes (line 4). During backward, ZeRO re-collects all model weights again via second all all-gather (line 5) to calculate gradients (line 6). Once gradients are calculated on each GPU, ZeRO empties weights buffer again (line 7) and conducts reduce-scatter operation do gradient averaging and re-distribution (line 8). Model states and parameters are updated in the optimizer step (line 9).
11
In summary, to minimize the on-device memory footprint using ZeRO-3, at each iteration there three collective communication operations: two all-gather on weights and one reduce-scatter on gradients.
2.2 COMMUNICATION REDUCTION TECHNIQUES
Quantization is often used to reduce memory footprint, and data movement volume by using low precision to represent data (Dettmers) 2015 Dettmers et al. 2022]. However, the loss of information from representing high precision data with lower precision often comes with accuracy degradation. Related work aims to enhance quantization accuracy by addressing the challenges associated with differences in number ranges and granularity between high and low precision data. Some related work (Zhao et al. 2019) propose to filter the outliers in data to mitigate the gap in numerical ranges. Yet their accuracy hinges the quality of outlier filtering and it brings extra filtering overhead. Dettmers et al.] (2022) propose to use block based quantization on optimizer states to improve the quantization accuracy. However, it requires modifications to the model structure, thereby limiting its usability.
Gradient Compression techniques, including 1-bit SGD, 1-bit Adam, and 1-bit Lamb, have been developed to optimize gradient communication in distributed training by utilizing minimal bit repre- sentation (Seide et al. 2014 Tang et al. 2021 Li et al.) 2021). However, direct application of these methods to ZeRO-3 is infeasible as they assume full view of optimizer states (OS) across GPUs, which is not the case in ZeRO-3.
ZeRO Communication Reduction. Recent optimizations on 7eRO-3, like MICS ZZang et al. 2022cp and HSDP (Zhang et al. 2022a), divide the GPU cluster into sub-groups and leverage high bandwidth intra-node interconnect, or hierarchical communication to minimize communication vol- ume. hpz in ZeRO++ adopts a similar approach yet it performs only secondary partitioning on weights, while keeping all other model states partitioned across all GPUs. This allows hpZ to achieve significant communication reduction without the massive memory overhead of MiCS.
DESIGN
In this section, we elaborate on the design of our three key optimizations in ZeRO++ introduced in Section [ for reducing the communication overhead of ZeRO: i) Quantized Weight Communication for ZeRO ZqwZ), ii) Hierarchical Partitioning for ZeRO (hpZ), and iii) Quantized Gradient commu- nication for ZeRO (q We further discuss the end-to-end impact of these optimizations to reduce to total communication volume of ZeRO in Appendix B.4
3

---

Published as a conference paper at ICLR 2024
3.1 QUANTIZED WEIGHT COMMUNICATION FOR ZERO (qwZ)
As discussed in Section 2.1 ZeRO partitions the model weights across all the ranks (i.e., GPUs) and fetches the FP16 weights layer-by-layer right before they are needed in computation via all-gather for the forward and backward of each training iteration. To reduce the communication overhead of forward all all-gather on weights, qwZ,, quantizes FP16 weights to lower precision right during the all-gather, and dequantizes them back to FP16 on the receiver side, and then conducts layer computation.
While this reduces the communication volume of the all-gather by than 2x, doing naively results in two major issues: i) the lowering of precision results in significant accuracy degradation during training as discussed in Section 2.2 and ii) the quantization and dequantization overhead negates any throughput gain from communication volume reduction.
To improve quantization accuracy we use blocksize based quantization (Yao al. 2022 Shen al., 2020) (further details in Appendix [.I). To mitigate the quantization and dequantization overhead, we develop custom optimized implementation of qw (further details Appendix
Additional Benefits of qw?
i) Automatic Quantization qwZ automates parameter quantization during training, allowing higher compression (to 2 bits) with minimal accuracy loss, thereby obviating the need for post training
quantization (Frantar al. 2022 Yao et al. 2022) (see Section for more details),
ii) Memory footprint reduction For untrainable/frozen weights (e.g., in LoRA multimodal train- ing), qu alleviates the necessity of persistent FP16 storage, which reduces memory footprint. This reduction can allow for a) training with larger batch sizes for better throughput, b) inference or train- ing with fewer resources without running out memory, and c) speed up inference by reducing the amount of parameter data that needs to be read from memory.
In fact, given the automatic quantization and memory footprint reduction, qwz has a similar effect for fine-tuning as QLoRA (Dettmers al., 2023), in terms of weight quantization and memory requirement reduction, but addition ZeRO++ with qw can support significantly larger model sizes than QLoRA due to the weight partitioning from ZeRO.
3.2 HIERARCHICAL PARTITIONING FOR ZERO (hpZ)
With hp?, we eliminate the inter-node all- gather during the backward pass by holding secondary FP16 weights partition within each node. We do this by creating a hierarchical par- titioning strategy consisting of two partitions: first, all model states are partitioned globally across all devices in ZeRO-3, which call primary partition. Second, a secondary copy of FP16 parameters is partitioned at the sub- global level (c.g.. compute node, see Figure which we call secondary partition. This sec- ondary copy of FP16 parameters is replicated across multiple secondary partitions.
Machine
Machine
Machine
Machine
gather weight portitions (c)N)-5)
All secondary weight dell)+s)
Figure 2: hpZ removes node traffic in backward all-gather by holding secondary weight partitions on- device memory.
Consider a 64-node cluster, each node with 8 GPUs. Model weights are partitioned in two stages: i) across all 512 GPUs that we call primary partition, and ii) the same weights are also partitioned within a compute node across 8 GPUs, that we call secondary partition. In this example, for the secondary partition, each compute node in the cluster holds a full replica of FP16 weights partitioned among the GPUs within the node, and there are 64 of such replicas in total.
A training iteration with hpZ. During the forward pass of a training iteration, we all-gather weights based on the primary partition across all GPUs. However, once the weights have been used during the forward pass, these weights are then partitioned based on the secondary partition. Given the temporal consistency of model parameters between forward and backward passes, when the weights needed again during the backward pass, all-gather weights based on this secondary group.
4

---

Published as a conference paper at ICLR 2024
When the secondary partitioning is set to be inside compute node, this avoids any inter-node com- munication for this all all-gather. Finally, at the end of the iteration, during the optimizer step, all the model states, and the primary copy of the fpl6 parameters are updated based on the primary parti- tion. hp?. makes two changes to baseline ZeRO pseudocode in Algorithm [ i) in line 4, parameter partitioning is based secondary group size, ii) parameter all-gather preceding backward pass in line 5 is also based on secondary group size.
Our design of hp2 can flexibly support any secondary group size. The group size controls how many ranks (i.e., GPUs) are in the secondary partition. It is also measure of the memory-communicatinn trade-off of hp2 discussed in Appendix B.2
3.3 QUANTIZED GRADIENTS COMMUNICATION FOR ZERO (gg2)
qgZ is novel quantized reduce-scatter algorithm based on all-to-all collectives that enables 4x communication volume reduction of gradient reduce-scatter by replacing FP16 with INT4. qgZ has three components: 1) all-to-all-basee quantized gradient reduce-scatter, 2) reducing communication volume with hierarchical collectives (details in Appendix B.3), ). 3) tensor slice reordering for correct gradient placement.
3.3. ALL-TO-ALL BASED IMPLEMENTATION
A naive way to quantized reduce-scatter without precision loss is to apply quantization and dequanti- zation a ring-based based seduue-scatter directly as shown on the left of Figure We inject quantization and dequantization on each GPU. Once a GPU receives gradients from its predecessor, we dequan- tize it to full precision and conduct local reduction. After that we quantize local reduction output and pass it to its successor. To complete the whole reduce-scatter, the number of sequential quantiza- tion and dequantization kernels equal to the number of GPUs in use. Thus, applying quantization on existing ring based reduce-scatter collective will lead to high communication latency and low value precision due to multiple sequential quantization and dequantization steps. Although recent tree-based collectives like Blink(Wang et al. 2020) could reduce the number of sequential kernels from n to log(n), the long latency and low precision issue is not completely resolved.
To overcome this, completely abandon ex- isting ring-based reduce-scatter approach and incorporate 1-hop all-to-all collective for our gradient communication. As shown on the right of Figure 3 we first apply quantization on given tensor, then we conduct all-to- all communication among all the GPUs. Af- ter all-to-all, we apply another dequantization to recover the data precision and then reduce on high-precision values to get the final gra- dient reduction output. By replacing ring- based solution with our all-to-all collective, we reduce the number of sequential quantiza- tion+dequantizzttoon kernels from the number of GPUs to 1. We further reduce cross node communication volume by incorporating hierarchi- cal collectives named h-hop all-to-all, which is detailed in Appendix B.3
Quantiaaiion
Docuanttation Roduction
Machine
Nochine 1
Macine
Michine
G2
3
NCOL Ring based reduce scater (ZeR0-3 3] If aenuuaaraall OHD em of
1-hop al (ng2 Z6RD+e) #cfasquantia DVO -.
Figure 3: Comparison between ZeRO-3 ring-based reduce-scatter and qgZ I-hop all-to-all.
3.3.2 TENSOR SLICE REORDERING FOR CORRECT DATA PLACEMENT
With our 2-hop all-to-all, inter-node communication volume is expected, however, this introduces a gradient misplacement issue. We deseribe this issue using a 2x2 example, where we have ma- chines and each machine has 2 GPUs. As shown in Figure the correct final gradient placement shown as green boxes in the figure, where GPU 0 holds final gradient partition 1, GPU holds gradient partition 2, on and forth.
Our 2-step all-to-all communication works as follows, first we divide all gradients on each GPU into 4 chunks, then conduct our intra-node all-to-all. After intra-node all-to-all finishes, GPU0 (i.e., GO) holds partial aggregated gradient partitions 1,2 whereas GI holds gradient partitions 3,4. Same thing happens G2 and G3. Since G1 does not have gradient partition 2 (which is supposed to be
5

---

Published as a conference paper at ICLR 2024
Madhine
Machoe
Mathine
Machied
Mammeedd
Machine
3-
Initar grad. sioes
Nongton
Gred. alter into comm
Mtoplacemerg
Correet finall grad.
Grad. after intar-comm
Step intra-node al-iolll
Stop mode
2: intrectode albtoval
Stop 2. Fiter mode al
Figure 4: Gradient partition misplacement when ap- plying hierarchical all-to-all in qgZ.
Figure 5: Tensor slices reordering correct gradient misplacement in qgz.
held by GI) while G2 does have gradient partition 3, after inter-node all-to-all, there is gradient misplacement issue on both G1 and G2.
We address this with tensor slice reordering. In Figure before intra-node all-to-all, first swap tensor slice order of slice and 3, which is shown as orange arrows. Then after intra-node all-to-all completed, Gl now has gradient 2 while G2 has gradient 3. Therefore, after inter-node all-to-all, all GPUs get the correct gradient placement. Mathematically, given X GPUs per node and Y nodes total, each GPU holds X*Y gradient slices initially. Tensor slice reordering works as follows:
before: : [0, 1, 2, 3, YX YX Y,
after: [0, X, 2x, (Y 1)x; x + 1. 11.. 1]
Based on Eq. D and 2
we can map each original tensor slice position (i.e., Eq. D new tensor slice position (i.c., Eq. on each GPU to correct final gradient misplacement issue.
3.4 OPTIMIZED IMPLEMENTATION
To optimize ZeRO++, we implemented key optimizations: i) overlapping compute streams and, ii) optimizing CUDA kernels. Overlapping involves concurrett execution of quantization compu- tation and communication during weight all-gathering, leveraging asynchronous quantization and execution order tracking. The CUDA kernels were optimized for quantization, dequantization, and tensor slice reordering to fully utilize device memory bandwidth and minimize memory traffic by 9x in Z. Additionally, hierarchical approach in gradient communication and generalized ten- sor slice reordering scheme are introduced to minimize latency and address the needs of different pipeline stages. These optimizations allow ZeRO++ efficiently leverage 4x communication vol- ume reduction, enhancing its throughput. Please refer to the Appendix C for implementation details and communication analysis.
EVALUATION
This section evaluates ZeRO++ in three areas. First, shows end-to-end throughput scalability and speedup over baseline for standard and RLHF training across different models, model sizes, hardware configurations and cluster settings, demonstrating consistent speedup (up 3.3x) across the board. Second, it shows convergence properties of ZeRO++ for both pre-training and fine- tuning demonstrating its robustness and tolerance extreme quantization all the way down 2-bits. Third, it shows ablation studies demonstrating the impact of each component of ZeRO++ and the effectiveness of our kernel optimizations.
4.1 METHODOLOGG
Hardware: 24 NVIDIA DGX-2 nodes where each with 16 V100 SXM3 32 GB GPUs. The nodes are connected by InfiniBand (IB) with NVIDIA SHARP support, achieving total inter-node bandwidth of over 800 Gbps. To evaluate ZeRO++ in clusters under different network environments, show the performance of ZeRO++ running with different cross-node bandwidths by enabling from to 8 IB connections (i.e., 100 Gbps to 800 Gbps).
We use ZeRO-3 as the baseline given its ease-to-use for training giant models at large scale. To evaluate the performance of our optimized kernels, we also implemented ZeRO++ with Py- Torch quantization and non-fused kernels as baselines for ablation study.
Baseline:

---

Published as a conference paper at ICLR 2024
Baselin
Table 1: End-to-end speedup of ZeRO++ on V100 and A100 GPUs.
D.T,Mor
(a) 364 V100-GUU
33
TIE
Cotniction
IIB Connectios
Medel Size
Tokon GPU
Buseline TELOPs
ZeRLS++
TELOPs
Speciuu
Baselin:
TPLOEs
Z:RU3
TELOPs
Specllp
1330
2K
19.39
27.50
90%
47.35
55.30
1388
1K
11.25
21.81
9416
34.19
44.38
30%
915
2K
19.99
3570
5016
47.74
56.35
15%
91B
11.27
21.93
95%
34.49
44.36
29%
2K
20005
1908
9016
4502
56.21
17%
T7L29
49B
11.27
21.95
95%
34.54
44.06
29%
188
2K
25.8
46.40
70%
47.31
33.05
13%
18B
14.15
30.57
1168
31.27
37,87
218
(b) 32 A100 GPUS
1B Conncctinn
4 Connection
Figure 6: Scalability on up to 384 GPUs of 18B model with different numbers of InfiniBand con- nections and tokens per GPU.
Mealel
Siec
Tokan
GPU
Bancline
TFLOPs
ZeRD5+
TRLOPs
Spocdlu 71.8% 83.8%
Sasclue TRLOPs
Z,R03++
TELOP,
Sprattp
186
64.99 34.52
111.66 65.2
134,85 85.15
134.94 95.12
0.2%
18R
9.9g
Model Configurations: We use transformer models for evaluation including GPT, OPT, andLLaMA for various model sizes. Given Regatron-TTurngg 2K GPUs using 2K tokens per GPU (i.e., micro batch size), we evaluate ZeRO++ with the same 2k tokens per GPU setting. We also evaluate on IK tokens per GPU to test ZeRO++ with more extreme scale scenario. The number of layers and hidden sizes are adjusted to have models of different sizes. Please refer to the appendix and our open-sourced evaluation scripts for hyperparameters and other training details.
(Smith et al., 2022) training 530B model on
4.2 E2E SYSTEM EVALUATIONS
We evaluate ZeRO++ end-to-end performance and present an ablation study here. One key metric we use here the percentage of peak performance, which is shown as Eq.
peak.performance Schieved,TFLLOP TORPs/maxTFFLOPP
E.g. when we use V100 GPU, its max TFLOSS is 120 TFLOPs (VVIDIA V10O datasheet) for mixed precision computation. Thus, our reported peak performance refers to the percentage number of achieved.TFLLOP//II
Scalability upto 384 GPUs In Figure 6 we present ZeRO++ scalability evaluation from 64 to 384 GPUs with 18B model on both low (1 IB) and high (8 IB) bandwidth clusters. On low bandwidth cluster, ZcRO++ achieves 30% and 38.3% of peak performance (120 TFLOPs) even at 384 GPUs for 1K and 2K batch sizes, which is much higher compared to 12.5% and 21.6% baseline peak performance. This presents up to 2.4x better throughput. On high bandwidth cluster, despite having significantly more bandwidth, ZeRO++ still enables up 1.29x better throughput, and can achieve up 45% of sustained peak throughput at 384 GPUs. ZeRO++ significantly speed up large scale training for low bandwidth clusters while archiving decent speedup even on high bandwidth clusters.
Throughput for different model sizes and GPU architectures Table TTa compares training throughput for models of 18B-138B on 384 GPUs between ZeRO++ and baseline on both low and high bandwidth clusters. On low bandwidth cluster, ZeRO++ consistently achieves over 31.5% and 18.1% peak performance for 2K and 1K batch sizes on all models. Compared with the baseline peak performance of 16.6% and 9.3%, the speedup is up to 2.1 On high bandwidth cluster, ZeRO++ peak performances are 44.7% and 31.5%, which is 1.3x over the baseline peak performance of 31.5% and 26.0%.
Table Tb) illustrates the performance of ZeRO++ on 32 A100 GPUs, demonstrating that ZeRO++ surpasses the bascline even on smaller scale cluster. Specifically, in low-bandwidth clusters, ZeRO++ exceeds the baseline by 88.8%, and in high-bandwidth clusters, achieves a 9.3% im- provement over the baseline. ZeRO++ exhibits robustness and uniform speedup across varying model and batch sizes, as well as differing network bandwidths and GPU architectures.
Democratization for large scale training Figure 7 compares the throughput of ZeRO++ a low cross-node bandwidth (200 Gbps as 2 IB) cluster with the baseline running on 800 Gbps high- bandwidth (8 IB) cluster. For small model of 18B, ZeRO++ achieves higher peak performance of

---

Published as a conference paper at ICLR 2024
Size:
Model GPU:
$4
TR.0
10
8:0
2000 pal1
pe6 Dkm
Figure 7: ZeRO++ achieves high bandwidth cluster performance with much lower bandwidth.
Figure 8: ZeRO++ achieves up to 3.3x better throughput for RLHF training.
Validation UM losa vs. Staps
Validation LM loss vs. Staps
Table 2: Perplexity after -unint different mod- els with various quantization bits in ZeRO++.
2080+
ZekGe- Zek (28e1)
OPT-1338
OPT-138
LLaMA-300 30B
FPI6 (Baseline)
L804
1,698
1.490
ZeRO++ 6-bits
L809
1.705
1,500
ZeRO++ 4-bits
1830
1,705
1.494
ZeRO++ 2-bits
2.218
1.809
1.544
Figure 9: Convergence curves: Pretraining GPT: 350M and GPT-13B on the Pile dataset.
41.6% compared with baseline peak performance of 39.1% despite running with 4x lower cross-node bandwidth. For large model of 138B, ZeRO++ and baseline achieve the same peak performance of 40%, but baseline runs at 4x higher cross-node bandwidth. This evaluation shows that ZeRO++ makes large scale training more accessible by significantly decreasing the minimum cross-node bandwidth requirement for efficient training. Furthermore, it demonstrates that optimized ZeRO++ implementation effectively translates the 4x communication reduction of ZeRO++ into real end-to- end system throughput gain.
Speedup for RLHF Reinforcement Learning from Human Feedback (RLHF) is unique and com- monly employed scenario in LLM training. In Figure [ we present a comparison between ZeRO++ and the baseline for OPT-30B and LLaMA-2 2-00BB demonstrating that ZeRO++ 3.3x and 2.97x improved throughput 2o-LLMM--2-708 and OPT-30B respectively. The throughput measured step3 of RLHF training where there are both training and generation. It's noteworthy that the results were obtained with 8 InfiniBand connections on V100 GPUs, and we anticipate an even larger gap lower bandwidth clusters. These findings underscore the versatility and efficacy of ZeRO++ across diverse training scenarios and various model families.
achieve up to
4.3 MODEL CONVERGENCE ANALYSIS
This section assesses ZeRO++'s impact on the convergence of large models during both pretraining and fine-tuning stages.
Pretraining. We analyze the pretraining of GPT-350M and GPT-13B models on the Pile dataset Bii derman et al. 2022), employing ZeRO++ with non-blocked quantization, ZeRO++ (with blocked quantization), and ZeRO-3 the bascline. To maintain fairness, all hyperparameters remain con- sistent between baseline and ZeRO++ trainings. Convergence measured by the validation LM loss.
Figure 9 displays the comprehensive pretraining trace. Training with 8-bit non non-blocked quantization diverged initially, rendering no visible data. Conversely, ZeRO++ with 8-bit blocked quantization aligns closely with the baseline, reinforcing our prior analysis that block-hased quantization achieves superior quantization accuracy. Furthermore, ZeRO++ convergence remains closely aligned with the baseline even when using 6-bit quantization.
Fine-tuning. We fine-tuned the pre-trained models: OPT-1.3B/-11 Zhang al. 2022b) and LLaMA-30B (Touvron et al., 2023), with ZeRO++ using FP6, INT4, and INT2 precision for qu Z and INT-4 for 7 on the high-quality open-source instruction dataset We kept all training hyper- parameters the same across all setups. The evaluation relies on the metric perplexity (the lower, the better).
We include the huggingface datasets: Dahoas/mm-stttic, Dahoas/full-hh-rlh Dahoos/yynthetic-nnssuut. gptj-pairwise, wwt.tnngciiilf---eeataaaeeeee
8

---

Published as a conference paper at ICLR 2024
Table End-to-end performance when using ZeRO++ w.) wo.
upc
w.o optimized kernels.
Qw2
hp2
Number Bt:
Number of IB:8
TPLOPs per GPU
Optimized Quantization
Oplimized Fusion
4344343 44
TFLOPs
Kemel
Kernel
20
Baseline
NIA
N/A
15
ZeRO++
No
No
19.73
Pick 10122
Narr BAOS
Pic
90121
ZeRO++
No
Yes
21.6
Figure 10: Throughput of 18B models on GPUs with ZeRO+t, qwZ, qgZ. hpZ., and baseline vari- ous InfiniBand connections.
ZcRO++
Yes
No
31.40
ZeRO++
36.16
raveal the robustness of ZeRO++, well as its ability create low- wew-preciioon models during fine-tuning without requiring post-quantization. Notice, the validation perplexity of ZeRO++ varies only slightly from the baseline, with mere 0.27%-0.41 % deviation, even when quantized to 4-bits.
4.4 THROUGHPUT BREAKDOWN AND ANALYSIS
Impact of Individual Techniques. In Figure TO we show the individual and combined impact of qwZ, hpZ. and qgZ, on the throughput of 18B model on 128 GPUs. On low bandwidth clusters, each of these techniques enables a speedup ranging from 1.3-1.4x compared with baseline, while achieving an aggregated speedup of up to 2.26x. Note that our TFLOPs throughput is calculated from wall-clock time measurement, ZeRO++ aggregated throughput gain is not equivalent to sum of qgZ, qwZ, hpZ gain. We can validate the theoretical speedup with composition our techniques by accumulating the speedup multiplicatively: 1.4 1.26 1.3 2.29, which is very near to what we achieved as 2.26x.
For high bandwidth clusters, the individual speedup ranges between 1.13-1.16x, for combined speedup of up to 1.3x. The figure demonstrates that each these techniques has a similar impact towards throughput improvement and they compose effectively with larger aggregated speedup.
Impact of Kernel Optimizations. We evaluate our optimized kernels on ZeRO++ throughput using 18B model running 64 GPUs.
Quantization Kernel: As shown in Table compared with the baseline that uses Py Torch quanti- zation, our optimized quantization kernels can achieve up 1.67x speedup terms of end-to-end throughput. Also, the baseline implementation suffers performance degradation as the group number increases which means the throughput gap will be larger when used with larger models.
Kernel Fusion: As described in Appendix C.2, kernel fusion is one of our key optimizations to improve memory throughput when executing sequences CUDA kernels. Our fusion includes 1) tensor-reorder and quantization fusion 2) intra-node dequant, intra-node reduction and inter-node quant fusion. As shown in Table 3 we achieve up 1.15x speedup on the end-to-end throughput.
CONCLUSION
This paper presents ZeRO++, an efficient collective communication solution for large model training using ZeRO stage-3. We optimize both model weights and gradients communication in the forward and backward pass of each training iteration. To reduce communication volume of model weights forward propagation, we adopt block-based quantization and data pre-fetching. To remove cross- node communication of weights during backward pass, we hold secondary model partition on each node to trade memory for communication. To minimize gradient communication during backward propagation, design and implement novel all-to-all based gradient quantization and reduction scheme.
By incorporating all the three optimizations above, improve system scalability and throughput for pre- traininnnng, fine-tuning and RLHF training with speedup of up to 3.3x. Furthermore, ZeRO+ can automatically quantize the parameters to ultra-low precision making the resulting model inference ready without framework for training large models at scale.
quantization. We envision ZeRO++ as the next generation of easy-to-use

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We would like thank all DeepSpeed members for their helpful discussion and support to this work, and anonymous reviewers for their constructive feedback. This work was sponsored in part by NSF grant CAREER-2048044
REFERENCES
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD: communication--efiiiine SGD via gradient quantization and encoding. In Isabelle Guyon, Ul- rike von Luxburg. Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Confer- ence on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA. pp. 1709-1720, 2017. URL http::tttpttpcceea///aacee aa .tmaaattttmmml.
neurips. cc/paper/2017/
Abstract
Tal Ben-Nun and Torsten Hoefler. Demystifying parallel and distributed deep learning: An in-depth concurrency analysis. ACM Computing Surveys (CSUR), 52(4):1-43, 2019.
Stella Biderman, Kieran Bicheno, and Leo Gao. Datasheet for the pile, 2022.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao. Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivan- shu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. GPT-NeoX-20B: An open-source autoregressive language model. 2022.
Jeffrey Dean, Greg Corrado. Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc' aurcloo Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. Advances in neural information processing systems, 25, 2012.
Tim Dettmers. 8-bit approximations for parallelism in deep learning.
arXiv preprint
arXiv:151104561, 2015.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. The Tenth International Conference on Learning Representaatinns ICLR 2022, Virtual Event, April 29. 2022. OpenReview.net, 2022. URL URLlttppss//ppeeeviiee net//
forum? 1 w=w.mm/dssshpppppxxxx.x
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. arXiv preprint arXiv:2305. 14314, 2023.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04885, 2018.
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint 22Xi::2210.17322. 2022.
Eduard Gorbunov, Filip Hanzely, and Peter Richtárik. unified theory of SGD; variance reduction, sampling, quantization and coordinate descent. In Silvia Chiappa and Roberto Calandra (eds.),
The 23rd International Conference on Artificial Intelligence and Statisties, AISTATS 2020, 26-28 August 2020, Online [Palermo, Sicily, Italy,, volume 108 of Proceedings of Machine Learning Research, pp. 680-690. PMLR, 2020. URL http://pooooeiinnnn.m...eeeeee/. gorbunov20a... html
mlr wwwps/eee8le
Yanping Huang, Yonglong Cheng, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. ArXiv, aws//811..0990000 2018.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
10

---

Published as a conference paper at ICLR 2024
Conglong Li, Ammar Ahmad Awan, Hanlin Tang, Samyam Rajbhandari, and Yuxiong He. 1-bit LAMB: communication efficient large-scale large- batch training with lamb's 's convergence speed. CoRR, abs/2104.06065 2021. URL wwwpvvaaivvvvvva///iv//...aa..iv wwwpwbg..//.ooooo...tt...t...ert 20tp/22144... 06069
Ilia Markov, Adrian Vladu, Qi Guo, and Dan Alistarh. Quantized distributed training of large models with convergence guarantees, arXiv preprint arXiv:2302.023390 2023.
Microsoft.
Turing-nly
A 17-billion-parameter language
model by mi-
crosoft.
www. tuur g- bi www.oooaaaaeeeeeeaaannneeeoeeeee oor.o/ooooooaaaaaaaaaaaaaaaaroaa crosofft/
ht http://wwwww.ioooooffoooo//oeees
2020.
Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Granger, Phil Gibbons, and Matei Zaharia. Pipedream: Generalized pipeline parallelism for dan training. In ACM Symposium on Operating Systems Principles (SOSP 2019), October 2019.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vi- jay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on gpu clus- ters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC '21, New York, NY, USA, 2021. Associa- tion for Computing Machinery. ISBN 9781450384421. doi: 10.1145/3458817:34460000 URL
https:://do http://o.oorro/////////o S1t..1151111..3377888888
NVIDIA DGX-2. NVIDIA DGX-2.
nvidia www.ii.cooddooooooo..... ceoom//nnusddata-ceettt
dgx- 2/ 2018.
NVIDIA
SHARP.
Nvidia infiniband
adaptive routing
tech-
nology.
https: //vvdam widen. net/s//////////e hetts//wmszwwfrtte wwchmswwfre
oo 2021.
NVIDIA V100 datashect.
Nvidia tesla v100 gpu accelerator.
https://wwa www.
2017.
PyTorch quantization. Quantization -pytii
pytorch documentation. hntt::://pytorch.orr////
wttblee/uantizatiittt html 2023.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language
models are unsupervised multitask learners. 2019.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Perfor- mance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of the Interna- tional Conference for High Performance Computing. Networking, Storage and Analysis, SC '21, 2021.
Ali Ramezani- Kamnzaan-Keebyy, Fartash Faghri, Ilya Markov, Vitali Aksenov, Dan Alistarh, and Daniel M. Roy. NUQSGD: provably communication--fficcienn data-parallel SGD via nonuniform quanti- zation. J. Mach. Learn. Res., 22:114:1-114:44 2021. URL
v22 222/220255..0 html
Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient descent and its application to data-pun aata--arallll distributed training speech dnns. In Fiffeenth annual conference of the international speech communication association, 2014.
Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W. Mahoney, and Kurt Keutzer. Q-BERT: hessian based ultra low precision quantization of BERT. In The Thirty- Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Secone Innova- tive Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12,
2020, pp. 881 1821. AAAI Press, 2020.
11

---

Published as a conference paper at ICLR 2024
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari,, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deep- speed and megatron to train megatron-turinng nlg 530b, large-scale generative language model. arXiv preprint arXiv:2201.1 /990, 2022.
Hanlin Tang, Shaoduo Gan, Ammar Ahmad Awan, Samyam Rajbhandari, Conglong Li, Xiangru Lian, Ji Liu, Ce Zhang, and Yuxiong He. 1-bit adam: Communication efficient large-scale training with adam's 8 convergence speed. In International Conference on Machine Learning, pp. 10118- 10129, PMLR, 2021.
DeepSpeed Team and Rangan Majumder.
DeepSpeed: Extreme-scale model train-
ing for everyone. https: www mi mwwoorioott...oooooooooomooooooe wtt.:///eee////eeeee///aeaaa deepspeed
foreouo--ooreveeeeeeeeeeeeo 2020.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09220 2023.
Guanhua Wang, Shivaram Venkataraman, Amar Phanishayee, Jorgen Thelin. Nikhil R. Deva- nur, and Ion Stoica. Blink: Fast and generic collectives for distributed ML. In Inderjit S. Dhillon, Dimitris S. Papailiopoulos, and Vivienne Sze (eds.), Proceedings of Machine Learn- ing and Systems 2020, MLSys 2020, Austin, TX, USA, March 2-4, 2020. mlsys.org, 2020. URL mlsys www.ooggg.oooooooooooooo 29t.k//99.pp.. pdf
Jaxxiang Wu, Weidong Huang. Junzhou Huang, and Tong Zhang, Error compensated quantized sgd and its applications large-scale distributed optimization. In International Conference on PMLR, 2018.
Machine Learning, pp.
Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. Advances in Neural Information Processing Systems, 15:27168-277 2022.
Buyun Zhang, Liang Luo, Xi Liu, Jay Li, Zeliang Chen, Weilin Zhang. Xiaohan Wei, Yuchen Hao, Michael Tsang, Wenjun Wang. Yang Liu, Huayu Li, Yasmine Badr, Jongsoo Park, Jiyan Yang, Dheevatsa Mudigere, and Ellie Wen. DHEN: A deep and hicrarchical ensemble network for large- scale click-through rate prediction. CoRR, abs/2203. 1014, 2022a. doi: 10448550/RXII.2220. 11014. /tt.tttss///.....ooooo/////////
doi http:/oo.oooo...oo//ooo/o 485.0 wtw.s////rrivv.. 2ww.:xii...2.311111.. 11014
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li. Xi Victoria Lin, al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01066 2022b.
Zhen Zhang, Shuai Zheng, Yida Wang. Justin Chiu, George Karypis, Trishul Chilimbi, Mu Li, and Xin Jin. Mics: Near-linear scaling for training gigantic model on public cloud, 2022c.
Ritchie Zhao, Yuwei Hu, Jordan Dotzel, Christopher De Sa, and Zhiru Zhang. Improving neural network quantization without retraining using outlier channel splitting. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 7543-7552. PMLR, 2019. URLhttp: mlr. wwt.ppers//v///heeee html
APPENDIX
A FURTHER BACKGROUND AND RELATED WORK
A.1 DATA, TENSOR AND 3D PARALLELISM
Data parallelism (DP), pipeline parallelism (PP), and tensor parallelism (TP) are three forms of par- allelism used train large models across multi-GPU clusters. (Dean al. 2012 Narayanan al.,
12

---

Published as a conference paper at ICLR 2024
(a) Baseline vs. Blocked Quantization
(b) Quantization Error
0.631 6.529
9.220
0.631
9.220
500
Baseline
3.044
-2.475
4.690
-3.044
-2475
4.690
Euclidean Distance
8.441 2.901
5.707
88.441
2.901
5.707
Quantize
400
127
-128
128
-18
127
-116
-128
127
-64
300
Dequantize
0.653
-6.534
-92220
0.631
6.53
9.220
3.049
2,468
4.719
-3.033
-2.461
4.690
200
-8421
-5.735
8.44
2.901
5.729
Quantization Error: 2.937E-04
Quantization Error: 9.005E-05
64 5122048
Number of Blocks
Figure 11: Illustration & example of block based quantization vs. baseline
2021 2019 Huang et al. 2018 DP is commonly used when model size fits within single GPU memory. In DP, each GPU holds full copy of model weights and trains separate input data. MP is orthogonal to DP, and is often used in cases where model size cannot fit into a single GPU's memory. Instead of splitting input data, model parallelism partitions full model into pieces and assigns each model piece onto a GPU. There are mainly two approaches for model parallelism: i) pipeline parallelism (PP) and ii) tensor parallelism (TP). PP (Huang et al., 2019 Narayanan al.) 2019 Huang et al. 2018) splits models vertically, creating sequential stages consisting of a contigu- ous subset of layers. While there is sequential dependency between stages for an input micro- batch, the stages can be executed in parallel across micro-batches. In contrast, TPP (Narayanan et al. 2021) splits each layer across multiple GPUs, where each GPU works on different part of the layer for the same input.
3D parallelism (Smith et al. 2022 Team & Majumder 2020) refers to combination of Data Par- allelism. Pipeline Parallelism and Tensor Parallelism (Dean et al. 2012 Narayanan et al.) 2021 2019 Huang et al., 2018). and is capable of achieving excellent throughput and scalability, and has been used train wide range of large language models (Microsoft) 2020 Narayanan al. 2021 Radford et al. 2019 Black et al. 2022). Despite being highly efficient, 3D parallelism is severely limited by the fact that it requires complete rewrite of model and training pipeline make them compatible with 3D parallelism (Smith et al. 2022).
A.2 OTHER COMMUNICATION REDUCTION TECHNIQUES
Quantization In addition the quantization approaches discussed in the main text, another line of research focus on error-feedback techniques to achieve efficient quantization without significant accuracy loss, such as (Tang al. 2021 Wu et al. 2018). We view error-feedback as an orthogonal approach to 7eRO++.
ZeRO Communication Reduction Recent work such QSDP (Markov at. 2023) shows promis- ing results using quantization of weights and gradients. Compared to these approaches, ZeRO++ uses a novel communication design to reduce quantization error, overlapping communication and quantization hide quantization overhead. and optimized weight placement to reduce inter-node communication, which is reflected in the non- non-triviae triviri speedup the evaluation.
B
DESIGN DETAILS
B. BLOCK-SIZE BASED QUANTIZATION
As illustrated in Figure II each weight tensor is divided into smaller chunks, and converted into INT8 by symmetric quantization, using independent quantization scaling coefficient. By keeping the quantization granularity small, we significantly mitigate the gap in number ranges and granular- ity.
13

---

Published as a conference paper at ICLR 2024
Machine0
Machine1
Memory Consumed
K-12
M-1008
gpuo
gpul
gpu2
gpu3
P-1024
DP
(2 K) M
1.6TB
ZeRO-3
1.56GB
K+2a)-M
14.06GB
hpZ
Parameters (Primary)
Gradients
Optimizer
Parameters
States
(Secondary)
Per-device memory consumption analysis of standard data parallel (DP), ZeRO stage 3 (ZeRO- 3) and proposed hierarchical partitioning of ZeRO parameters (hp2). K denotes the memory multiplier of optimizer states, M represents the number of trainable parameters, P is the data parallel group size world size, and the number of secondary groups or ratio of world size the number of ranks the secondary group. A typical real world scenario example is provided the last column. We assume a model size of 100B trained on 1024 V100 GPU DGX cluster (64 compute nodes, 16 GPUs per node).
Figure 12:
CPUS sodo
Model dibe
Quandicaiio compresiionnn 2
GPUS node 3
Midel liee
Quern tohion compression 1916:
Machene
Match
14/2
MA/prrN (A777N)
Gi
Reduce scatter
1-hap all-tolll hoTMAZZ
Dach comm, valume reductions MAN210
reduction:
Doos-aade gradient comm, vtlume [(ano-3]
Drow-rade gradient comm volume (ggd SelDDee
Step Inco aode-
Step Intor mode allio
Figure 13: Compare the communication volume of
Figure 14: qgZ apply hierarchy all-to-all to reduce cross node traffic.
ZeRO-3 reduce-scatted scatter with the qgZ. 1-hop all-to-all.
We show an example of the quantization error of performing block based quantization vs. the non-blocked quantization baseline in Figure aa).. Figure [T(bb shows case study of weights quantization on BERT model, where block based quantization reduces the quantization error by 3x. More in-depth convergence evaluations are shown in Section
B.2 MEMORY USAGE ANALYSIS OF hpZ
By design, hpZ trades memory for communication efficiency. is important to analyze this tradeoff. Recall that standard data parallel DNN (DP) replicates model parameters across data parallel ranks, ZeRO-3 on the other hand partitions parameter across data parallel ranks. A midway approach is model parameters partitioned across subset of devices as long as model parameters fit.
Figure provides concrete memory usage estimate of a typical large language model size of 100B parameters, with primary group size of 1024 GPUs and secondary group size of 16 GPUs (e.g.. DGX-2 V100 node). As shown in Figure [2] with our proposed method, hpZ consumes 8.9z more memory than ZeRO-3, our approach is still 14x less memory requirement than standard DP. This marginal increase in memory usage is compensated for by efficient intra-node communication schedule. By eliminating reducing inter-node communication for backward pass, hpZ reduces the end-to-end communication of ZeRO by 1.5x, while supporting model training with hundreds of billions of parameters.
B.3 REDUCING INTER-NODE COMMUNICATION VOLUME
Although replacing reduce-scatter with all-to-all achieves single-shot quantization and dequanti- zation, introduces a new problem; the inter-node communication volume increases instead of decreasing despite the quantization of data. We elaborate this in Figure
14

---

Published as a conference paper at ICLR 2024
Given that mntra-machine often have high bandwidth interconnects (e.g., NVLink, NVSwwtch), cross- machine communication links are often the bottleneck. Given this, we analysis cross node commu- nication volumes and ignore intra-node communication volumes.
Here we assume model size of M, GPU per node is N, gradient compression ratio as Z. Reduce- seatter, reduces the data during transmission over the ring. thus the total amount of data for cross- node communication is M. However, when using our 1-hop all-to-all approach, even though the data are compressed before communication (i.e., M/Z), each GPU needs to send out M/Z amount of data to GPUs on the other nodes. Therefore, each machine will generate N * M/Z amount of cross-node communication data, which is much bigger than reduce-scatter communication volume,
To address this, we do a hierarchical 2-hop all-to-all instead of l Ihoop a) first intra-node all-to-all and b) followed by inter-node all-to-all, which is shown as Figure [4] First, with banhdhh links among GPUs inside a machine, we conduct intra-node all-to-all on quantized data, then dequantize data and reduce on dequantized data. After intra-node quantization, all-to-all, dequantization, and reduction, we reduce the data size per GPU from M/Z to M/(Z N). After intra-node all-to-all completed, we conduct the inter-node all-to-all communication, which is similar to I I-hop all-to- all we described above. Given that now each GPU only needs to send out M/(Z N) data, the communication volume per machine is now M/ * N) * N = M/Z. By adopting this hierarchical resolve the communication volume blow-up issue in
all-to-all communication as 2-hop approach, our 1-hop scheme perfectly. Note that even though the total communication volume is doubled (one intra-node, the other inter-node), intra-node communication introduces negligible overhead given NVLink/NVswitch high bandwidth, and cross-node traffic has been significantly reduced, which is the major bottleneck in gradient communication.
B.4 ZERO++ COMMUNICATION VOLUME ANALYSIS
Comm. Volume
forward all-gather
backward all- all-gather
backward reduce-scatter
ZeRO-3
M
M
M
ZeRO++
0.5M
0.25M
Table Communication volume comparison between ZeRO-3 ZeRO++.
Table illustrates theoretical communication volume comparison between ZeRO-3 and ZeRO++. We assume the model size of M. As described in Section 2 during ZeRO-3 there 3 collective calls: all-gather on weights in forward pass, then all-gather on weights in backward pass and last is reduce-scatter on gradients in the backward. And each collective communicates M volume of data.
With ZeRO-3, in total we need to communicate 3M data per each training iteration. Given that intra- node communication is fast with NVLink and NVSwitch, we ignore intra-node communication and focus on cross-node traffic only. For all all-gather in the forward pass, by incorporating our quantized weights communication, we reduce communication volume from M 0.5M. During the all- gather in the backward pass, by holding secondary weights partition within each node, we completely removed node traffic. For reduce-scatter in the backward pass, by replacing reduce- -catter with our novel quantized gradient communication protocol, we reduce cross-node traffic from M to 0.25M. Therefore, compared with ZeRO-3, ZeRO++ reduces communication volume from 3M down to 0. 0.75M for each training iteration.
C OPTIMIZED IMPLEMENTATION
In this section, we discuss two key optimizations that enable ZeRO++ to fully realize the potential of 4x communication volume reduction to improve throughput without getting limited by imple- mentation overheads: i) overlapping different communication and compute streams, when doing so enables better resource utilization, and ii) optimized CUDA kemnlss for quantization, dequantization, and tensor slice reordering operators, and kernel fusion across these operators when appropriate to minimize the memory traffic overhead. Below discuss the two lines of optimization in detail.
15

---

Published as a conference paper at ICLR 2024
Deccarttifedatteen Coant
Deceantrfolaccion
No Nodkrr
Intrennde A7A
Data Chunkil
TTrr rode
A2A
Intervndee ADA
Labency Reduction
Figure 15: Pipelining and overlapping intra-node communication with inter-node communication qgZ.
C.I OVERLAP COMPUTE AND COMMUNICATION
To reduce end-to-end communication time, we overlap quantization computation with communica- tion for all-gathering of weights in both forward and backward passes. For the hierarchical all-to-all based reduce-scatter implementation of gradients, we overlap the intra-node communication with inter-node communication.
C.1.1 COMMCNICATION-COMUUNN OVERLAPPING ON WEIGHTS
For all-gather on weights, we enable commuuiccationcooputalli overlap using two key features i) we track the execution order of model layers to get the sequence they will be fetched. ii) we guarantee asynchronous quantization execution. Specifically, the call to the quantization kernel is non-blocking and we further avoid operations that involve explicit/implicit CUDA synchronization (e.g. tensor concatenation), making the quantization a non-blocking operation that can be launched asynchronously.
With this two features, as ZeRO fetch parameters for each layer, the communication of the current layer and the quantization of the next layer be launched the same time different CUDA streams. When the quantized data are needed for the next layer, ZeRO++ synchronizes the quantiza- tion stream to make the quantized data ready. This approach hides the quantization cost of the next layer under the communication time span of the current layer which hides the quantization overhead.
C.1.2 HIERARCHICAL COLLECTIVES FOR GRADIENT COMMUNICATION
As discussed in Section B.3 our all-to-all based gradient communication broken into two stages: first intra-node communication followed by inter-node communication. The inter-node communica- tion depends on the results of the intra-node communication, therefore, with naive implementation, inter-nodes links are idle during intra-node communication and vice versa. To reduce latency by leveraging both inter-node and intra-node links in parallel, we chunk our input gradient tensor and pipeline transfer between intra-node communication and inter-node communication. As shown in Figure 15] compared with "no pipeline" case on the top, simply adopting "2-stage pipeline" trans- fer achieves the amount of end-to-end latency reduction shown as the red arrow-line Figure [5] By overlapping intra-node and inter-node communication, the end-to-end latency of gradient com- munication is significantly reduced.
Doing this pipeline correctly has implications our tensor slice reordering process. The more pipeline stages we have, the more fine-grained tensor slices are needed for reordering. Therefore, also propose generalized tensor slices reordering scheme algorithm [ which covers both w/ and w/o pipelining data transfer cases. Here stages refer to the number of pipeline stages we have, nodeSize is the number of GPUs per node and nodes is the number of nodes.
Next, we discuss how we optimize our CUDA kerness to further reduce all quantization related overhead.
16

---

Published as a conference paper at ICLR 2024
Algorithm 2: Generalized tensor slice reordering (g92)
Constants:
slages, node.Size, nodes pprrlition D)
Input
Output : :mappedPart aariition
Devices & mode:Szze * nodes;
stagel D + partition % stages;
chaukk1 1D + partitionID.
stages
pipeline Of + stagel D Deviccs;,
chunkOf fee + modeSSize'
clonk1D.
chunk Base + DchunkID % nodeSize) x nodes;
Return: pipeline f set + chunk Base + chunkOf fset;
C.2 CUDA KERNELS
As existing quantization implementations are unable capture the combination of data mapping and high throughput necessary to minimize kernel overhead, we implement and optimize custom CUDA kernels implement these primitives. In particular, these kernels aim (I) saturate device memory bandwidth and (2) minimize the total traffic via fusion.
core quantization and dequantization library of compos- able operators was developed as the foundation for ZeRO++. The core primitives leverage efficient vectorized memory accesses at the maximum granularity given GPU architecture supports. In order to satisfy the alignment requirements these instructions have, model state is partitioned such that quantization granularities will be 16B aligned. Additionally, we leverage instruction level par- allelism to overlap multiple memory transactions with each other. In practice, the combination of vectorized accesses and instruction level parallelism enables the quantization library to achieve full GPU memory bandwidth utilization.
Maximizing Bandwidth Utilization:
Multiple techniques are used to reduce the total memory traffic for quantization kernels. First, the size of each quantization block is tuned as express sufficient parallelism to schedule across a GPU's streaming multiprocessors and cache values not quantized yet in the register file while calculating the quantization scale and offset for the block. Second, we fuse tensor reshaping and quantization into the same kernel to avoid redundantly loading data from global memory. For example, the tensor slice reordering (i.e., orange arrow-lines in Figure 5 is realized within a fused quantization and remapping kernel. This fused kernel achieves the same level of performance as a single quantization kernel working with contiguous data. Finally, we fuse sequential dequantization, reduction, and quantization operations into single kernel implementation, which reduces total memory traffic by 9x in ggZ.
Minimizing Total Traffic:
D ADDITIONAL EVALUATION & DISCUSSION
D.I EXTENDED CONVERGENCE DISCUSSION
In this section, we present the raw data of our convergence evaluation for both pretraining and fine- tuning Tan Please note the quantization bits shown here refer to the weight quantization. For the gradient, we use INT4 for all the experiments.
For theoretical convergence analysis, one thing to note is the unique design of ZeRO. In ZeRO paradigm, there are two kinds of weights: the temporary weights, which are all- all-gathered from all GPUs in order to do forward/backward on a particular model layer and then discarded; and the local weight shard, which is used to populate the temporary weights, are permanent and will be updated at each optimizer step. The ZeRO++ quantization only happens when we communicate the temporary weights. Thus, the local shard of the weights will be free of the quantization error. For the gradients that are computed from the quantized temporary weights, there will be minor errors for 8-bits and even 4 bits.
17

---

Published as a conference paper at ICLR 2024
Table 5:
Convergence in terms of validation Im loss
Table 6: Convergence in torms of evaluation per- plexity when fine-tuning with different quantization bits.
when pre-training with different quantization bits.
Model: GPT 350M
Model: OPT-1.3b
Iteration
FP16
INT8
Epoch
PP16
FP6
INT4-bit
INT2
2400
2.891633
2.89641
2.018
2.021
2.034
2,430
4800
2.595669
2.600295
1.972
1.975
1.989
2.355
1.934
1.938
1.953
2.320
7100
2.50512
2.509532
1.906
1.910
1.924
2.291
9500
2.409789
2.417458
5
1.882
1.888
1.905
2.273
11800
2.355472
2.366978
6
1.863
1.868
1.887
2.267
14100
2.331281
2.345854
7
1.850
1.856
1.871
2.248
16500
2.286052
2.303065
8
1.842
1.841
1.861
2.241
9
1.822
1.830
1,850
2.228
18800
2.279778
2.301165
1.822
1.833
1,844
2.226
21200
2.244547
2.26685
1.818
2.224
23400
2.220746
2.246421
1.809
1.817
2.223
13
1.805
1.811
1.832
2.218
14
1.804
1.809
1.835
2.221
Model: GPT 13B
15
1.807
1.809
1.832
2.222
16
1.809
1.812
1.829
2.216
Iteration
FP16
INT8
FP6
2400
3.346583
3.340124
3.327485
Model: OPT-13b
4800
2.794749
2.787855
2.789709
Epach
FP16
FP6
INT4-bil
INT2
7100
2.56617
2.561561
2.562001
1.813
1.821
1.821
1.842
9500
2.398518
2.399067
2.39818
2
1.698
1.705
1.705
1.705
11800
2.370753
2.358937
2.357589
3
1.724
1.717
1.736
1.734
14100
2.310235
2.33077
2.329953
16500
2.275294
2.277154
2.275939
Model: LLaMA-30b
18800
2.180487
2.239978
2.238353
Epoch
FP16
FP6
INT4-bit
INT2
21200
2.181295
2.194993
2.133342
1.57
1.58
1.580
2.524
23400
2.181079
2.184269
183004
1.49
1.50
1.494
1.664
1.64
1.63
1.601
1.544
18

---

Published as a conference paper at ICLR 2024
ZeRO
hpZ TELOPs
MICS
Model Size
Token Size
TFLOPs
TFLOPs
7.5B
IK
36.99
38.39
38.96
7.5B
2K
53.3
54.4
52.72
18B
IK
51.47
52,42
OOM
18B
2K
60.94
61.44
OOM
Table 7: hpZ vs MiCS evaluation on 4 node cluster (16 V100 GPUs per node)
64 V100 GPUs
Ethernet 10Gbps
Model
Tokens
Baseline
ZeRO++
Speedup
Size
per GPU
TFLOPs
TFLOPs
18B
2K
9.72
31.94
229%
18B
1K
4.98
19.16
284%
18B
512
2.51
10.03
299%
7.5B
2K
10.07
24.51
143%
7.5B
1K
4.81
15.90
231%
7.5B
512
2.48
7.32
195%
Table 8:
Evaluation of 7eRO++ pretraining throughput on Ethernet with 10Gbps bandwidth
In terms of the theoretical convergence with gradient noise/error, existing research such as (Gor- bunov et al., 2020 -Ramezani-Kenn Kebrya et al. 2021; Alistarh et al., 2017) prove minor disturbance on gradients will not impact the overall convergence (the convergence rate follows O(1/sqr(T) for T iterations), which we have empirically show.
We also crror-feedback frror-feebbcck technique such as (Tang et al., [2021] Wu et al.) 2018) could help especially for 1- bit/2-bit training for smaller model, which is a part of our future work.
kind of orthogonal approach with ZeRO++ and we believe adding
D.2 COMPARING HPZ WITH MICS
As previously discussed in Appendix 2 closely related to hierarchical weight partition for ZeRO (hp2) MiCSiZhang al. 2022c). The key difference between the two methods what data are replicated in the secondary group: only model weights are replicated in hp7, while entire model states are replicated in MiCS.
the per-GPU throughput of both methods for different model and token size configurations. The table also shows that given secondary partition size of single node (16 V100 GPUs), hpZ can support 18 billion parameter model whereas MiCS reports out-of-memory (OOM) this scale.
D.3 ETHERNET CLUSTER EVALUATION
In the main text, we present low-bandwidth evaluation with 100Gbps bandwidth. Here we fur- ther extend our evaluation 10Gbps Ethernet scenario to show how inter-node bandwidth affects ZeRO++ performance in a more comprehensive manner. The pretraining setting is the same, where weights are quantized into INTS and gradients are quantized into INT4. As Table ZeRO++ can outperform baseline by up to 4x in terms of throughput on 10Gbps Ethernet, and the perfor- mance pattern is consistent with previous evaluation with InfiniBand. This aligns with our evalua- tion analysis in the main text that the performance boost of ZeRO+ increases network bandwidth decreases, making +eRO++ a valuable tool for training in oow-banwwitth clusters.
19


Reviews:
[{"review_id": "XF18q9U6NY", "reviewer": "Reviewer_kJiX", "summary": "This paper proposes ZeRO++, a communication-efficient extension of the commonly used memory-efficient training framework for LLMs, ZeRO. The proposed techniques include three specific techniques to reduce communication costs, improve the throughput and accelerate the training process.", "questions": "1.\tWhat is the specific quantization methods in qwZ and qgZ?\n2.\tHow to choose the quantization ratio during training and fine-tuning? I am interested in the inner effect of ratio on the convergence of different tasks.\n3.\tIn the real low-bandwidth environment, does the trade-off between memory and communication still work without NVLink?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper proposes quantized weight communication, hierarchical partitioning, and quantized gradient communication techniques to reduce the communication costs and achieves efficient implementation of overlap computing and communication. The parallelism of training is the focus.\n\n2. The experiments presented are quite convincing, including large-scale language model training and fine-tuning on a large GPU cluster. The convergence accuracy and throughput reflect great improvements.\n\n3.\tTraining on low-bandwidth clusters is an important question and this paper solves it to some degree.", "weaknesses": "1.\tIn my opinion, the main insufficient of this paper is short of novelty. In fact, this paper is more like a technical report instead of an academic paper. The techniques of quantization and organizing the communication into a new topology are not new ideas. However, considering the overlap realization and trade-off between memory and communication are difficult and expensive, this paper is still significant for LLM community.\n\n2.\tAnother problem is that this paper is lack of specific introduction to used quantization techniques and corresponding theoretical analysis. For example, how the quantization error affects the convergence rate and whether the error-feedback technique works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yFoB1KJZxE", "reviewer": "Reviewer_TkZG", "summary": "The paper proposes an extension to the existing memory optimal training technique for LLMs, *ZeRO*. The proposed enhancements address the expensive communication problems of the state-of-the-art method, especially when trained on low bandwidth networks. ZeRO++ with the new quantization enabled communication strategies achieves reasonably good performance over ZeRO while keeping the performance of LLMs intact.", "questions": "- One simple experiment (even on a network of RasberryPi machines) with low-badwidth networks will obviously greatly strengthens the paper claims/contributions. This can be fine-tuning does not even have to be pre-training. \n- A feasibility study on the scalability of the method, especially because of the GPU memory tradeoff within the node, would help.\n- Upon some or all of the above concerns being addressed, would be happy to raise the score. \n- Also, wonder, if there is any open implementations for this complex custom communication MPI primitives?\n\n## Post Rebuttal \n\n- Raising the score to above the acceptance threshold", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written, detailed and easy to understand for both DL and HPC communities. (assuming readers have some basic understanding of both the worlds).\n- Has a good coverage of the existing literature and positions the paper with-in the body of work.\n- The problem is probably impactful for low-resource devices such as low band-width networks.", "weaknesses": "- The paper shows the effectiveness of the proposed method on powerful DGX GPUs with infiniBand interconnects. Noway, we can consider this interconnect as low-bandwidth. \n- Intra-node weight storage sounds interesting and empirically showed to work for upto 30B model fine-tuning and 13B model pre-training. Wonder, if this can be scaled if there is any ablation on this dimension?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PZiSSdYETJ", "reviewer": "Reviewer_V3tv", "summary": "This work, ZeRO++, proposes a collective communication optimization for ZeRO3 of DeepSpeed library. Three elements are introduced: 1) quantized AllGather of weight, 2) hierarchical partition, and 3) quantized ReduceScatter of gradient with AlltoAll. When combined together, ZeRO++ reduces communication volume and network hops to mitigate the communication bottleneck of ZeRO3 training for large models.", "questions": "*. Does this work provide open-source code?\n\n*. Is this work ever used in production for industrial models?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+. Replaced Ring-ReduceScatter with AlltoAll for lower latency and less quantization/dequantization times\n\n+. Optimized intra-node and inter-node AlltoAll", "weaknesses": "-. **Limited contribution**\n> qwZ quantizes FP16 weights to lower precision right during the all-gather, and dequantizes them back to FP16 on the receiver side, and then conducts layer computation.\n\n* Why cast FP16 weight into INT4 is a novel contribution? Just use \"torch.Tensor.to()\"?\n* What if model is being trained FP8 on the modern H100? so the communication benefit will be halved?\n* Will cast FP16 into INT4 occupy 25%+ extra GPU memory (assuming zero fragmentation)?\n\n> hpZ removes cross node traffic in backward all-gather by holding secondary weight partitions in on device memory.\n\n* How is this scheme different from HSDP of PyTorch (Intra-Machine FSDP to shard weight + inter-Machine DDP to replicate weight)? HSDP was published in VLDB'2023.\n\n> qgZ is a novel quantized reduce-scatter algorithm based on all-to-all collectives that enables a 4xcommunication volume reduction of gradient reduce-scatter by replacing FP16 with INT4.\n\n* Same question as weight quantization\n\n-. **Limited Convergence Analysis**\n\n> Pretrain: We analyze the pretraining of GPT-350M and GPT-13B models on the Pile dataset\n\nHow does this work apply to different datasets for convergence behavior, considering extreme quantization?\n\nBesides Figure 9, does this work provide accurate numbers for loss diffs for pertaining?\n\n> Fine-tuning. Table 2\n\nIt seems this work always worse the perplexity compared with baseline.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tdFFV2JreR", "reviewer": "Reviewer_ZpeK", "summary": "This paper proposed ZeRO++, which includes three main techniques to accelerate communication in distributed training of foundation models: i) low-precision AllGather for parameter collection, ii) hierarchy communication optimization, and ii) low-precision gradient aggregation. Empirical studies was conducted to verify the performance of the proposed method.", "questions": "Would it be possible to provide some theoretical analysis about the convergence given the quantized communication of the weights distribution?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Optimizing the communication in ZeRO or more generally speaking, optimizer parallelism is an important problem to scale out / speed up foundation model training.\n\n- The proposed design and implementation are reasonable. \n\n- The evaluation is conducted on a large production-level GPU cluster to evaluate models at the state-of-the-art scales.", "weaknesses": "- My main concern about the paper is lack of novelty as a research paper instead of a software report, concretely, w.r.t each technique contribution, I have the following comments:\n  - Low-precision gradient averaging: such a technique has been studied for decades in standard data parallel training; integrating them in an existing system probably cannot be viewed as a novelty in a top ML volume (i.e., ICLR).    \n  - Organizing the communication as a hierarchy topology is not a new approach; a similar implementation was introduced in PyTorch-FSDP about six months ago.\n  - Distributing the model weights with lower precision is an interesting point. However, there is a lack of concrete theoretical analysis about this approach, i.e., now the model parameter includes some quantization error; how would this influence the convergence behavior?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JWdiIZR7I6", "reviewer": "Reviewer_QbwT", "summary": "This paper proposes zero++, which contains 3 communication reduction techniques (low-precision parameter all-gather, data remapping, and low-precision gradient averaging), in order to reduce the communication overhead and acceleration the training.", "questions": "1. Does the gradient and parameter communication compression use the same quantizer/quantization algorithms? What are the specific quantization algorithm used in the experiments?\n\n2. In experiments, some details are missing. For example, in Section 4.3 \"Pretraining\" (or Figure 9), it is mentioned that 8-bit or 6-bit quantization is used. However, it is unknown whether is 8-bit for both parameters and gradients? or 8-bit for parameters and fp16 for gradients? or fp16 for parameters and 8-bit for gradients? I strongly recommend the authors to specify what (parameter or gradient) is quantized wherever quantization/precision configuration is mentioned.\n\n3. In the experiments,  Section 4.3 \"Fine-tuning\" (or Table 2), it is mentioned that \"with ZeRO++ using FP6, INT4, and INT2 precision for qwZ and INT-8 for qgZ ...\" I wonder why not also conduct some experiments of lower precision (int4, int2) on qgZ? Also, is there any ablation experiments on other configurations, such as int4 qwZ + fp16 qgZ?\n\nBTW, there is \"FP6\" in Section 4.3 (see the sentence I quoted above), is it a typo and the authors actually mean \"FP16\" here?\n\n4. What's the difference between Zero++ and QSDP (in system design), except for the hierarchical partitioning in Zero++ (hpZ)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. This paper proposes zero++, which contains 3 communication reduction techniques (low-precision parameter all-gather, data remapping, and low-precision gradient averaging) with sufficient details provided in the appendix.\n\n2. Additional to the system design, an optimized implementation is also proposed.\n\n3. The experiments show good performance on real-world training tasks.", "weaknesses": "1. Some part of the paper is a little bit vague and confusing, please refer to the \"Questions\" below for more details.\n\n2. The main reason I'm holding back from a better score is that zero++ seems to have huge overlap with QSDP (Markov, I., Vladu, A., Guo, Q., & Alistarh, D. (2023). Quantized Distributed Training of Large Models with Convergence Guarantees. ICML 2023). I understand that QSDP was officially published on ICML not very long ago and we reviewers should not take the release date on arxiv as a reason to reject this paper. However, this is a delicate situation for me and I guess we will need the AC to judge. \nFor now, I still need the authors to give a detailed comparison between QSDP and zero++, which will help a lot for my paper review.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)