CO2: Efficient Distributed Training with Full Communication-Computation Overlap
Authors: Weigao Sun
Zhen Qin
Weixuan Sun
Shidi Li
Dong Li
Xuyang Shen
Yu Qiao
Yiran Zhong
Venue: ICLR 2024 spotlight (2024)
Keywords: Distributed Training
Data Parallelism
Local Updating
Asynchronous Communication
Primary area: optimization

Abstract: The fundamental success of large language models hinges upon the efficacious implementation of large-scale distributed training techniques. Nevertheless, building a vast, high-performance cluster featuring high-speed communication interconnectivity is prohibitively costly, and accessible only to prominent entities. In this work, we aim to lower this barrier and democratize large-scale training with limited bandwidth clusters. We propose a new approach called CO2 that introduces local-updating and asynchronous communication to the distributed data-parallel training, thereby facilitating the full overlap of COmmunication with COmputation. CO2 is able to attain a high scalability even on extensive multi-node clusters constrained by very limited communication bandwidth. We further propose the staleness gap penalty and outer momentum clipping techniques together with CO2 to bolster its convergence and training stability. Besides, CO2 exhibits seamless integration with well-established ZeRO-series optimizers which mitigate memory consumption of model states with large model training. We also provide a mathematical proof of convergence, accompanied by the establishment of a stringent upper bound. Furthermore, we validate our findings through an extensive set of practical experiments encompassing a wide range of tasks in the fields of computer vision and natural language processing. These experiments serve to demonstrate the capabilities of CO2 in terms of convergence, generalization, and scalability when deployed across configurations comprising up to 128 A100 GPUs. The outcomes emphasize the outstanding capacity of CO2 to hugely improve scalability, no matter on clusters with 800Gbps RDMA or 80Gbps TCP/IP inter-node connections.

Published as a conference paper at ICLR 2024
CO2: EFFICIENT DISTRIBUTED TRAINING WITH FULL COMMUNICATIO--OOMUUTATUUU OVERLAP
Weigao Sun, Zhen Qin, Weixuan Sun, Shidi Li, Dong Li, Xuyang Shen,
Yu Qiao, Yiran Zhong'
OpenNI Shanghai AI Laboratory (sunweigas gao, zhongyi ran @pjlab. org. en https: ttttb/bbgttttttt////ggt
thub. com//ppenNtLPLab//C..
ABSTRACT
The fundamental success of large language models hinges upon the efficacious im- plementation of large-scale distributed training techniques. Nevertheless, building a vast, high-perfen perfoomeffrm cluster featuring high-speed communication intercon- nectivity is prohibitively costly, and accessible only to prominent entities. In this work, we aim to lower this barrier and democratize large-scale training with limited bandwidth clusters. We propose new approach called CO2 that introduces local- updating and asynchronous communication to the distributed data-pan mata-parallel training, thereby facilitating the full overlap of COmmunication with COmputation. CO2 is able to attain a high scalability even on extensive multi-node clusters constrained by very limited communication bandwidth. We further propose the staleness gap penalty and outer momentum clipping techniques together with CO2 to bolster its convergence and training stability. Besides, CO2 exhibits seamless integration with well-established ZeRO-series optimizers which mitigate memory consumption of model states with large model training. We also provide mathematical proof of convergence, accompanied by the establishment of a stringent upper bound. Fur- thermore, we validate our findings through an extensive set of practical experiments encompassing wide range of tasks in the fields of computer vision and natural language processing. These experiments serve to demonstrate the capabilities CO2 in terms of convergence, generalization, and scalability when deployed across configurations comprising up to 128 A100 GPUs. The outcomes emphasize the outstanding capacity of CO2 to hugely improve scalability, matter clusters with 800Gbps RDMA or 80Gbps TCP/IP inter-node connections.
INTRODUCTION
Distributed optimization is crucial for the efficient training of aarge-scale deep neural networks. Mini- maich parallel optimization methods (Goyal et al. 2017 Li et al., 2014) like stochastic gradient decent (SGD) with distributed data parallel (DDP) paradigm are commonly used, but communication overhead can pose significant challenges when scaling out to larger GPU clusters. Existing techniques leverage gradient bucketing to partially overlap communication with backward computation to enhance training efficiency, but residual overhead remains challenge in scenarios with large model sizes and limited inter-node communication bandwidth.
Various strategies have been proposed to address the communication--elated issues. These strategies can be classified into three following categories: 1) Communication Compression in Single Iteration,
which comprises techniques such gradient sparsification (Shi et al., 2019 Li & Hoefler 2022
Barnes et af. 2020) and gradient quantization (Alistarh et al.) 2016 Tang et al. 2021 Li et al. 2021
Deutmers al. 2)2T). These methods accelerate distributed training by minimizing the volume of
communication traffic during each iteration. 2) Communication Frequency Reduction is represented
by work such as (Lin et al., 2018 Wang al. 2019 2020, Wang & Joshi 202T), which maintains the communication volume per iteration but lessens the frequency of communication events. 3) Communication and Computation Overlapping aims to overlap communication with computation,
"Corresooniigg author.

---

Published as a conference paper at ICLR 2024
either in single step across multiple steps. For instance, the distributed module in PyTorch (Li (et a 2020) leverages gradient bucketing to partially overlap gradient all-reducing with the backward
pass. Additionally, asynchronous distributed training methods (Dutta et al. 2021 Cohen et al. 2021
Mishchenko et al. 2022 Su et al. 2022; Koloskova et al. 2022) synchronize the transmission of stale gradients with the latest local computation on each worker.
These communication--fficiee approaches improve convergence speed but often yield sub-optimal final optimization outcomes. None of these methods have achieved complete overlap of communica- tion with computation, i.e. 100% scalability throughout the entirety of the training process on a large cluster, especially under varying communication conditions while preserving the performance.
In this paper, we propose CO2 which enables complete overlap of Comunication with Computation. This is made possible through the use of the local updating strategy and asynchronous communication, where each worker node independently optimizes its model parameters without requiring synchro- nization with other peers at the local updating stage. As shown in Fig. T to guarantee parameter consistency across all workers, model parameter synchronization is performed asynchronously after completing every r local updates. By selecting appropriate local updating step accordance with the communication environment, we can achieve full overlap of model parameter synchronization with multiple local computation steps, thereby attaining 100% scalability.
Maintaining training stability well performance with asynchronous parameter updates is crucial for the success of our method. We introduce two additional techniques: staleness gap penalty and outer momentum clipping. The staleness gap penalty mechanism effectively quantifies and penalizes the discrepancy between different parameter versions during the update of outer momentum. This approach significantly contributes the overall training stability and performance by addressing the inconsistencies in parameter updates. Similarly, outer momentum clipping serves as an essential tool to mitigate the emergence of anomalous values the outer momentum. This helps prevent extreme values, thus further bolsters the training stability. In addition to our empirical findings, we establish robust theoretical convergence bounds. These bounds provide compelling evidence that our proposed framework is capable of achieving convergence rate that is on par with that of baseline optimizers. It is worth noting that CO2 can seamlessly integrate with widely adopted ZeRO- optimizers (Rajbhandari et al. 2020), which are well-established for their capacity to reduce memory consumption within paralparllll distributed training scenarios.
We evaluate CO2 by testing it on various tasks such image classification, semantic segmentation, point cloud processing, autoregressive language modeling, and bidirectional language modeling. Our thorough assessment demonstrates that CO2 performs equally well in terms of convergence speed and generalization capabilities when compared to other optimization methods.
Our primary contributions can be summarized as follows:
Outstanding scalability, CO2 enables outstanding scalability by fully overlapping com- munication with computation. This approach excels in large-scale multi-node clusters with limited communication bandwidth.
Good convergence and generalization performance. We empirically demonstrate that CO2 achieves convergence and generalization performance that is close to well-established baseline optimizers, such SGD and Adamw. Our results underscore the robustness of CO2 across diverse range of tasks and datasets.
Theoretical convergence analysis. We provide a rigorous theoretical analysis that guaran- tees the convergence of CO2. This theoretical foundation enhances our understanding of the practical effectiveness of our approach.
Compatibility with ZeRO-series optimizers. CO2 can integrate with ZeRO-series opti- mizers to reduce memory usage for large-scale model training.
2 RELATED WORK
Local Updating, To our known, local updating in distributed deep learning can date back to Zhang et al. (016). Subsequently, extensive research has delved into local updating methods. Local- SGD Stich 2018) provided rigorous theoretical convergence proof for local updating methods. Dost-Local-Sooe (Lin et al. 2018) and Overlap-Local-S (Wang et al., 2020) reduced commu- nication costs but raised concerns about convergence. SlowMo (Wang al. 2019) improved
2

---

Published as a conference paper at ICLR 2024
FP
BP
LU
FP
BP
BP
SGD
Grad. AR
Grad. AR
Comp. Stream
FP
BP
FP
OU
FP
BP
FP
CO2
Param. AAR
Comm. Stream
Wall Time
FP: Forward Pass
BP: Backward Pass
LU: Local Update
OU: Outer Update AAR: Asynchronous All- -lldee
Figure 1: Visualization of CO2 and SGD. We exemplify the mechanism of CO2 with local step count 7 = 2. This configuration dictates that the outer update starts after every two local steps, concurrently launching AAR operation model parameters. This strategy is made to make the full overlap of AAR communication with local computation possible. CO2 can effectively reduce the wall time required for training compared to the conventional SGD in DDP paradigm.
convergence stability and Cooperative SGD (Wang & Joshi 2021) presented unified framework of communication effficient SGD algorithms. However, communication overhead remains chal- lenge. Our proposed CO2 introduces asynchronism to overlap parameter synchronization and local computation, achieving 100% scalability with good convergence performance.
Asynchronism. There is longstanding history of asynchronous distributed training techniques (Chen et al. 20T6). Recently, Zhang et al. (2015): Barkai al. (2019) have introduced methods that penalize gradient staleness in asynchronous SGD employing well-defined gap metrics, Dutta et al. (2021) have delved into the trade-off's between staleness and convergence within the context of asynchronous training. Moreover, work like [Chhnn et al. (202I); Mishchenko et al. (2022). through rigorous theoretical proofs and empirical experiments, have established that asynchronous SGD can achieve convergence comparable to that of mini-batch SGD, even in the presence of arbitrary gradient delays. In a related vein, Su et al. (2022) proposed GBA as an approach to seamlessly transition between synchronous and asynchronous training, for recommendation models. Koloskova et al. (2022) provided enhanced convergence guarantees for asynchronous SGD in the context of distributed federated learning. These contributions collectively advance our understanding of asynchronous training techniques in distributed settings. Our method, as comparison, only introduces one-step asynchronous staleness the outer loop, thus obtaining more robust convergence and performance.
Efficient Communication. To improve communication efficiency in synchronizing gradients in distributed training, previous methods often utilized quantization techniques reduce the commu- nication overhead, such as 1-bit Adam (Tang et al. (02T). 1-bit LAMB (Li et al.) (021), and 0/1 Adam (Lu et al., 2022). In a distinct vein, Assran et al.aa222000 (2019) proposed stochastic gradient push to perform approximate distributed averaging at each iteration, which offers an alternative dimension for accelerating communication. These methods significantly enhance communication efficiency by mitigating communication overhead within single iteration, which is orthogonal to proposed method. More substantial benefits can be realized when integrating with these methods.
vein.|Assran
METHOD
In this section, we start with an overview of our proposed CO2 and share two novel techniques that enhance its convergence and training stability. Additionally, provide mathematical proof of convergence, which proves that CO2 is capable achieving a convergence rate that comparable to those of previously widely used optimizers.
3.1 THE OVERVIEW
For the distributed data-parallel training a large-scale deep neural network, the target can be formulated as:
min 1 BG, 1 (x) Ct),
(I)
where the objective is to minimize function with respect to the parameters x E R The optimization problem involves summation over workers, each indexed by i. For each worker there are loss function 7,69) and data samples Ç; drawn from the distribution Di-
3

---

Published as a conference paper at ICLR 2024
In the routine DDP training, model parameters replicated across all workers within the communi- cation group. Each worker is then provided with distinct batches of data samples and independently performs forward and backward passes to compute distinct losses T,(0) and gradients g g())_ To ensure uniformity of model parameters across all workers within the same DDP communication group, the gradients on all workers are exactly synchronized through an all-reduce operation, resulting in the aggregated gradient g. i.e., 9 δ Ei= g() This aggregated gradient is then utilized to update the model parameters across all workers, ensuring consistency of parameter values across each worker.
Local-updating methods, such as Local-SGD or other variants, struc- ture the overall training process into two discernible phases: the inner loop and the outer loop. The inner loop unfolds within each worker and en- compasses localized updates. Within this phase, every worker maintains its dedicated set of model parameters, fa- cilitating autonomous updates without the need for gradient synchronization at every step. Upon the completion of local updates, the training pro- cess transitions to the outer loop. In
Algorithm 1: CO2 Algorithm
Input: Data samples ((,) on worker i; Inner learning rate "it Inner loop steps т; Outer learning rate a; Outer momentum factor 3; Outer loop steps T;
Initial outer momentum mo =
for t E 10. , 1} do
3
for k E {0, 1, ... -- 1} on worker i do
FP & BP: 91.k
(8) (T L St.k
(1)
LU: a (0)
(6)
1,41 11.6 7
end
Launch AAR for Li.T AAR(Z! if is jis_completed(
(1)
this phase, an all-reduce operation on model parameters is executed to ex- actly synchronize the model param- eters across all computational work- 10
}) is not True then
01-1, = wait (AAR(x)
end
ers, i.e., x ? E = ensur-
(i)
Update staleness gap:
ing convergence alignment. Recent work (Wang et af., 2019) introduces additional outer iterates to enhance
A . 1.1 0..0- + 1"
12
Update one-step stale outer momentum:
the convergence of communication- efficient base optimizers, which has 13 been proven effective.
m. = 3m,-1 + X. (xt-1.0 - 11-1,))
Update outer iterates:
T1+1.0 2r,0 - a Clip(m Φ)
CO2 introduces asynchronous all- end
reduce on model parameters within the outer loop to achieve full overlap between communication and computation throughout the entire training procedure. As delineated in Algorithm [ during the outer iteration f, all workers independently perform local updates involving the FP, BP, and LU on their respective worker nodes. The inner loop iterates steps, from k = to k = T - 1. Notably, during this process, no gradient synchronization takes place among the workers, resulting in distinct model parameters (i) on each worker i at the conclusion of the inner loop. This necessitates a synchronization operation to harmonize the model parameters across all workers.
Consequently, upon the completion of inner loop, CO2 launches all-reduce operation on ! for each worker i to synchronize them, i.e., (=R( ) In contrast to conventional methods, this all- reduce is asynchronous with subsequent computations, which means the subsequent computations are not dependent upon the results obtained by ( Instead, it leverages the outdated model parameters x " which are subjected t 1. It is worth noting that, at this moment, AAR(x may not have finished. Therefore, a waiting operation wait(-) should be conducted if it is checked not completed, this guarantees the execution of (z) (4) (-.-). The resulting æt-1,7 is subsequently employed for updating the stale outer momentum, which is then applied to perform an outer iterate to obtain the initial parameters 1.+1.0 oo for next outer loop. For the special case of the first outer iteration, lines 8-13 in Algorithm T are skipped, since there is no outdated parameters to be used for updating outer iterates. For each respective outer loop, the stale outer momentum lags by only one step compared to using the new version of Ft,r. This strategy introduces minimal one-step staleness in the outer loop and has subtle impact on network convergence.
(t))
previous AAR operation in the preceding outer iteration
4

---

Published as a conference paper at ICLR 2024
The AAR operation runs in parallel with the subsequent outer updates well as local computations in the next outer loop. Given the practical considerations of cluster sizes and variable communication conditions, configuring appropriate number of local steps enables the seamless overlap of the entire asynchronous communication with computations.
3.2 STALENESS GAP PENALTY
CO2 achieves full overlap between communication and computation through one-step asynchronous communication, enhancing scalability up to 100% with an appropriate number of local steps. However, the asynchronous nature adds noise to the optimization process, potentially causing discrepancies in training loss and generalization accuracy. To compensate, we propose a staleness gap penalty mechanism, which requires an accurate metric quantify the staleness gap of outer momentum.
Recall the formulation in (1), where fx) Be L (a) is the local objective function at worker Before constructing the staleness gap, we first assume that each fi(x) is L-Lipschitz as stated in Assumption 3.1
Assumption 3.1 For some existed L there is [VVx(()) V S L|| yll, for all inputs y R" and E (1,2,-... ..}.
The Lipschitz assumption stated above is intuitive, it suggests that viable alternative metric for precisely quantifying gradient differences is to consider the distinctions in model parameters. This insight serves motivation for us employ the model parameters from different versions as indicators of the staleness gap for outer momentum.
To quantify the staleness gap in outer momentum, straightforward approach involves computing the difference in outer momentum between different iterations, such as iterations t + and f, and normalizing this difference assess its magnitude. Considering m 3m, + (æx, Xt.r)
and m, = 3m;-1 + (#t-1.0 2t x;-1,7), the staleness gap of m; in comparison to m4+1 can be straightforwardly represented as:
m m = []3 - m4-1) + (xi.o xt,r) - (.t- - 11-11
8|| (m m.-1) ||a., 1, ||,,,r 20 2,-1,
essential to emphasize that at the initiation of iteration I, the ongoing asynchronous all-reduce operation involving øt,- is generally in progress and has may not reached completion. This situation implies that a globally averaged value for Xt,7, accessible for all workers, is unavailable. Considering the outcomes presented in (2), suitable and readily accessible metric for measuring the staleness gap, denoted as A, can be defined as:
Definition 3.1 The staleness gap of the outer momentun A R at step t is quantified as the ratio of the model displacement within the outer loop to the maximum distance the parameters can traverse within the inner loop, in iteration 1 - 1. This can be formally articulated as follows:
||I,,0 21-1,01
A: T T - 1101
+1"
In (3), the term ||00,0 2.ol denotes the magnitude of the model parameter changes within which occurred with the presence of stale outer momentum.
the outer loop at iteration T|| 2;-1,oll represents the maximum distance that the model parameters can traverse in inner steps without utilizing stale outer momentum. Typically, during the training process, both the gradient and the learning rate within the inner loop experience decay, leading to reduction in the model pa- rameter distance within a single inner step. This behavior elucidates why || Æt-1.p signifies the maximum distance achievable in a single inner step. The notation 1 signifies the computation of is parameter-wise where n corresponds to the number of parameters. All computations within conducted element-wise. At each iteration, recalculate the staleness gap by performing
and then penalize it when updating the outer momentum via: m, = fm,-1 A. (IL-1.0 -000--1-1 -1-1.77
3.3 OUTER MOMENTUM CLIPPING
Ensuring training stability is of paramount importance, particularly for large language models. We empirically find that the asynchronism introduced by CO2 occasionally have adverse effects
5

---

Published as a conference paper at ICLR 2024
training stability. To mitigate the potential issues related unusual values in outer momentum and enhance overall training stability, we have implemented a coordinate-wis outer momentum clipping CXip(X, Φ) max{-0, xin((x,)) 6}}, which is used to update the outer iterates as 1,+ 2r,0 a Clip(m:, Φ), where ø denotes the clipping threshold, which satisfies Φ > 0.
3.4 CONVERGENCE ANALYSIS
Considering the staleness introduced by approach, it is imperative to establish its convergence guarantee from a theoretical standpoint. To do so, we begin by presenting the following assumptions, which align with standard practices in this domain.
Assumption 3.2 For all E (1.2....,Cr there exists finite positive constant 0² such that Eç-D. VL 9 (x: vf(()) s² i.e. the variance of fe(æ)) is bounded.
Assumption 3.3 There exists a finite positive constant V such that E g 13 [g V.
Under these assumptions, our convergence results are given below, where the detailed proof be found in Appendix A.4
Theorem
If we take x 1/Ar, Yt y and ignore the clip operation, such that X at,
and TT 2 GL3 [1 1+ v max { {-((1_--)) 4-8 1)) then under Assumptions 3.1 3.2 3.3 and
TT
DG Vf(x) - (x)) S 6², where is positive finite finte constant, we have:
+ EE (x,,))) - (vCiT)
+ + Cz
(GTT))
where f = Ei f()). The theorem indicates that, when the total steps T'r is sufficiently large, ie, T' >> - the RHS dominated by (verr) So, when the number of workers is G times more, we only need G times less total steps to achieve the same error.
EXPERIMENTS
We conducted thorough assessment of CO2 in range of deep learning applications spanning puter vision (CV) and natural language processing (NLP). Our testing encompassed an array model architectures, including convolution networks, transformer networks, and linear transformer networks. For more specific information regarding the tasks, models, their categorizations, parameter counts, and datasets used, please refer to Table in Appendix[A] It is worth noting that the ToM (7B) (Qin et al., 2023b 2024a b) experiment mainly aims to evaluate CO2's scalability, so we only pre-train it relatively small dataset WikiText-100 assess the convergence.
4.1 EXPERIMENTAL SETUP
We compared CO2 against other state-of toefthhharttt optimization techniques, including SGD, Adamw, Local-SGD/Adame
and SlowMo. Hyperparameters for each approach were meticulously tuned for maximum performance. For CO2, we conducted an extensive hyperpa- rameter search for 7, ranging from (1,3, 6, 12, 24, 48, 96, 192) to find the optimal balance between efficiency and performance. All experiments were executed five times with distinct random seeds to ensure robust results. For hardware, software, and hyperparameter details, see Appendix A. A.d A.2
Boerkap-Lcaa-SGDDAAaaae
Image Classification (IC). We train ResNet-50 (ee et al. 2016 Zhou et al. 2020 VIT Dooovviskiiy et al. 2020) Tang et al., (024), and VVT (Sun et al. 2023) on Image IK dataset, using eight total mini-bated size of 8192 and train for 90 epochs with a cosine learning rate schedule. For the training VIT and VVT, they closely adhere analogous hyperparameters. We uniform total batch size of 2048 for both and train for 300 epochs.
A100 nodes with 64 GPUs. For the training of ResNet-50,
Semantic Segmentation (SS). We leverage the pre-trained VVT on ImageNet- K as the backbone model, adopt semantic feature pyramid network (Semantic FPN) (Lin et al. 2016) as the decoder and finetune on the ADE20K dataset, using four 3090 nodes with 32 GPUs.

---

Published as a conference paper at ICLR 2024
Table 1: Convergence performance on CV tasks. CO2 performs better than other local-updating
methods with clear margin and is comparable standard optimizers such SGD//damw.
SGD
Local- SGD(Adamw)
Overtap-Local- SGD(Adamw)
SlowMo
CO2
Task
Model
(Adamw)
ResNet--00
76.92 (== 0.05)
75.57 (#: 0.76)
76.28 (# 0.18)
71.12 (#: 0.11)
77.14 (#: 0.09)
VIT (Base)
81.33 (1: 0.04)
78.43 0.22)
78.04 0.35)
79.83 0.16)
(1: 0.08)
VVT (Large)
83.64 (#: 0.06)
81.09 (#: L,15)
80.33 (#: 0.49)
82.75 (4: 0.27)
83.38 (:: 0.06)
SS
VVT (Lagge)
47.82 (#= 0.05)
44.25 2:24)
45.21 (# 1.36)
47.51 0.12)
47.80 (+: 0.11)
PC
MAF
68.56 (t: 0.08)
64.25 (d: 2.11)
63.78 1.92)
68.69 0.32)
68.89 (::: 0.39)
Red Meni 50
VIT
WT
60
60-
15
Val
18
20
13
Localscip
20-
Alatim
20
Adamw LocalAsare
30
Spochs
EDUcTs
Epoctss
Figure 2: Validation curves for image classification tasks. Three models ResNet-50, ViT, and VVT are trained ImageNet-1K for 90, 300 and 300 epochs, respectively. Our CO2 exhibits robust convergence and good generalization performance when compared to other existing methods.
3D Point Cloud (PC) Reconstruction. We pre-train the Point- Moint-MAE Pang al. 2022 Li et al. 2023) via reconstruction task using ShapeNet dataset (Chang et al. 2013). We use a point number of 1024, a total batch size of 256, epoch number of 300, train on four 3090 nodes with 32 GPUs.
Autoregressive Language Modeling (ALM). For ALM tasks, we train autoregressive GPT-2 mod- els with 125M (Small), 355M (Medium), and 770M (Large) parameters the OpenWeb'Te Text dataset (Radford et al. 2019). All experiments use a context length of 1024 and a total batch size of 480. training on eight A 100 nodes with 64 GPUs Additionally, we train the large linear trans- former (Qin et al., 2022a bj 2023a|c) model TransNormer-LLA (7B) on WikiText-100 103 for 100K iterations on sixteen A100 nodes with 128 GPUs, in order to verify the performance of CO2 on large language models. The ZeRO-2 optimizer is integrated with CO2 to lessen memory redundancy during training L.. (7B).
Bidirectional Language Modeling (BLM). For BLM tasks, we train a RoBERTa (Large) (Liu et al. 2019) model on a combined dataset consisting of Wikipedia and BookCorpus (Wettig et al. (0022) for approximately 3 epochs, using a total batch size of 4096 and sequence length of 128, train on four 3090 nodes with 32 GPUs
4.2 CONVERGENCE AND GENERALIZATION
Computer vision results. Table D shows the results of convergence and generalization in computer vision tasks. CO2 optimizer consistently outperforms other efficient optimization techniques in terms of Top-1 validation accuracy for ResNet-50, ViT, and VVT architectures. Although CO2 has slightly lower accuracy than the Adamw for both ViT and VVT, the reduction is marginal and falls within the expected range of random fluctuations. In addition to its good performance, CO2 achieves more efficient communication, including the potential for zero communication delay compared to the baseline methods. A similar trend is observed for VVT in semantic segmentation tasks. CO2 demonstrates a notable performance advantage over (Overlap-)Local-S slight performance gap of 0.02% when compared to Adamw. For 3D point cloud reconstruction, our proposed CO2 outperforms all other counterparts, exceeds the bascline Adamw by 0.33%.
and SlowMo, with only a
We also present validation curves that track training progress of these models on IC tasks. As shown in Fig. 2 CO2 features a more stable validation curve for ResNet-50 and maintains consistently higher accuracy curves compared to (Overlap-) Local-AAmw and SlowMo for ViT and VVT.
Natural language processing results. Table 2 and Fig. presents the convergence and generalization outcomes for natural language processing tasks. The table shows that CO2 consistently outperforms alternative methods across GPT-2 variants (Small, Medium, and Large) for ALM tasks, resulting in lower validation perplexity values as the model size scales. Furthermore, CO2 achieves the best.
7

---

Published as a conference paper at ICLR 2024
Table 2: Convergence performance on NLP tasks. Quantitative perplexity (PPL) results for GPT-2, TransNormer-LLMM and RoBERTa are presented. CO2 shows lower perplexity (lower is better) than the baseline Adamw in all these experiments.
Overlap- Local-Adamw
Task
Model
Adamw
Local-Adamw
SlowMo
CO2
(Small)
(1 0.36)
7.95 2.04)
1.03)
7.34 (c: 0.89)
0.73)
GPT-2 (Medium)
6.61 (#: 0.53)
7.49 1.87)
7.26 (#: 1,44)
6.41 (#= 0.65)
6.36 (# 0.66)
ALM
GPT-2 (Large)
6,26 (d:: 0.58)
7.00 (d: 1.91)
7.18 (d:: 0.98)
6.29 (d: 0.61)
6.13 0.52)
TN-LLM (7B)
16,82 (::: 0,86)
18.63 (#: 3.13)
17.83 (= 2.95)
16.95 (# 101))
16.78 0.95)
BLM
RoBERTa (Large)
3.96 (#: 0.37)
4.38 (#: 0.83)
4.52 (#: 1.42)
3,98 (#: 0.85)
3.95 (#: 0.96)
0P1/2
GP1-2
GP1-2 (args))
Addame SlevaMf
Locai-tsimm
Lccalasame
Sowelo
CO22
002
CO2
10-
06
20x
-005
601
FBX
100K
06
20c
10k
CBC
10k
100K
26
20k
4.8
2.K
120K
Stego
Stops
Figure 3: Validation curves for autoregressive language tasks. We train GPT-2 on OpenWebTer. for 100K steps in three sizes: 125M (Small), 355M (Medium), and 770M (Large). CO2 exhibits robust convergence and the best generalization performance when compared other existing methods.
training perplexity on the large TransNormer-LLLM (7B) model, when trained on a relatively small dataset. We also give validation curves with respect relative time (in seconds) for GPT-2 models of varying sizes as illustrated in Fig. 7 Appendix A.3
For BLM tasks, CO2 demonstrates comparable performance the standard Adamw and Mo by a significant margin. These results emphasize the capacity of CO2 to optimize encoder-only transformer models without compromising accuracy.
algorithms, surpassing
4.3 SCALABILITY
We tested our proposed CO2 using up to 16 DGX-A100 servers to train TransNormer-LLLMM with 7B parameters. Each server has 8 A100 GPUs interconnected via NVSwitch for an inter-GPU bandwidth of 600GBps. We evaluated CO2's scalability under different cluster communication conditions using two kinds of inter-node network configurations. The first uses RoCE solution with 8 RDMA adapters per server for inter-server communication at 800Gbps. The second uses TCP/IP Ethernet with significantly lower bandwidth around 80Gbps. Note that, mitigate fluctuations, the presented throughput results are averaged values extracted from iterations 100 to 200 for each experiment.
We evaluated methods, Adamw and CO2, on RoCE RDMA and TCP/IP networks. On the RoCE RDMA network, CO2's communication advantage became prominent in larger clusters (more than 64 GPUs) resulting in higher throughput than Adamw. On the TCP/PP network, CO2 performed similarly to its performance on the RoCE RDMA network when T was set to 48. However, Adamw performed poorly due pronounced communication latency. We also test the influence of T on communication efficiency, with varying levels of overlap. Results are shown in Fig. 4a).
Using CO2 with TCP/IP and 12 caused moderate reduction in speed and scalability, but it still outperformed Adamm on TCP/IP due to overlap between communication and computation. Transitioning from GPUs to 16 GPUs caused a notable drop in throughput values for both CO2 and Adamw due to the diminished inter-node connectivity.
Table 3 shows throughput and scalability of CO2 and Adamw. Scalability ratios were calculated for different communication scenarios, specifically when transitioning from 16 GPUs to 128 GPUs. CO2 exhibits scalability ratios exceeding 1, mainly due to inherent measurement fluctuations.
4.4 ABLATION STUDY
We conducted ablation studies on CO2 using a small-scale GPT-2 (Small) model across four servers, each with eight 3090 GPUs. The experiments lasted for 100K steps.
8

---

Published as a conference paper at ICLR 2024
1nt
TransNormar--LLM (7B)
(Small)
3.0
CD2 (RoCE RDMA, 12)
1.30
Adamw (RoCE RDMAI
2.5
(TCP)P,T== 48)
per
CO2 (TCPAP.T= 12)
Adamm (TCP/I))
Throughout (words
100
GPT2 (Smalt)
8.50
1.0
8.00
Valttelion
20
40
60
60
100
120
140
160
150
200
Num. of GPUs
(a) Scalability of CO2.
(b) Effects
Figure 4: (a): Scalability of CO2. Throughput (words/sec) results on distinctive inter-node network configurations are presented. CO2 exhibits pecfect 100% scalability on both configuration (b): Effects of Training speed and generalization performance results w.r.t. 7 are presented. A larger value of leads higher communication efficiency but worse generalization behaviors.
Table 3: Quantitative scalability performance of CO2. Throughput (words/sec) results CO2 and Adamw from 1 to 16 DGX-A 100 servers are presented. Scalability ratio records the scalability from 2 to 16 nodes to take into account the inter-node connections. With suitable configuration of T, CO2 can reach the scalability ratio of I, which outperforms the baseline Adamw.
Throughput
Scalability
Ethernet
Method
Ratio
GPUs
16 GPUs
32
64 GPUs
128 GPUs
(16 128)
CO2
17980 48)
35692 (+: 126)
72491 183)
150073 (#: 362)
307557 (#: 617)
1.08
RDMA
(r=12)
Adamw
19276 (= 59)
38888 (#: 118)
77782 (# 93)
151554 (#: 209)
289106 (# 423)
0.93
(m=48)
18090 95)
35969 (#: 193)
71249 179)
147507 (:: 315)
304736 (#: 729)
1.06
TCP/P
C02
(r=12)
17995
(#: 108)
36095
373)
564)
Adamw
18444
3488 (#: 115)
7077 (+ 127) 12855 (# 308)
22810 (::: 526)
0.82
Performance and scalability w.r.t. We conducted experiments examine time per iteration and validation perplexity across different values of т. Increasing 7 generally leads to faster training but higher perplexity. However, at 12, computation and communication overlap, resulting in significant acceleration. Larger values benefit from reduced communication frequency. = 12 also yields a slightly lower perplexity, which can attributed to the regularization effect.
Staleness gap penalty and outer momentum clipping. Table 4 shows that incorporating the staleness gap penalty has more significant positive impact than outer momentum clipping. The latter mainly targets training stability rather than performance improvement.
Table 4: Ablation results staleness gap penalty and outer momentum clipping. We pre-train tre-train GPT-2 (Small) with 100K steps for ablation. Both train and validation perplexity results indicate that staleness gap penalty has noteworthy improvement for convergence performance.
Model
Steps
Metric
CO2
CO2 w/o Penalty
CO2 w/o Clipping
Train PPL
7.36
7.52
7.39
Validation PPL
7.39
7.56
7.42
CONCLUSION
We proposed CO2, a method for large-scale distributed training on clusters with low-cost and limited- speed interconnections. CO2 enables exceptional scalability through local-updating and one-step asynchronous communication. Mathematical analysis is provided on convergence rate. Our experiments in computer vision and natural language processing demonstrate CO2's ability to achieve 100% scalability even on clusters with very limited communication bandwidth. We have enhanced CO2's convergence and stability through the introduction of two techniques: the staleness gap penalty and outer momentum clipping. CO2 is highly compatible with established optimizers and can be integrated with ZeRO-series optimizers reduce memory consumption for large model training.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This work is partially supported by the National Key R&D Program of China (N0.20222D0110100))
ETHICS STATEMENT
Given the widespread adoption of large language models trained on extensive corpora, the significance of distributed training across expansive clusters has grown substantially. Inefficient training practices result in substantial computational resource wastage. leading to elevated economic costs and an increased emission of CO2, a significant contributor to the global greenhouse effect. In the context of distributed training for neural networks on GPU clusters, efficient communication plays pivotal role. However, existing methods still grapple with low communication efficiency, particularly on modern, large-scale GPU clusters characterized by limited inter-node connections. Our proposed approach addresses this challenge by advocating for the full overlap of communication and computation, thereby substantially expediting model training even in scenarios with severely restricted communication capabilities. This not only reduces training time but also translates into cost savings and environmental benefits, as contributes to mitigating the impact on the environment.
In addition, our proposed method is relatively general and not limited specific models tasks. can be well extended to other tasks not included in this paper, such as stable diffusion for image generation, object detection for autonomous driving, etc.
REPRODUCIBILITY STATEMENT
For the sake of facilitating reproducibility in our experiments, we provide a comprehensive listing of the hardware and software components employed in our study. This information is thoughtfully presented in Appendix A and encompasses details such as the specifications of GPUs, inter-node communication configuration CUDA, cuDNN, and NCCL. Additionally, we give links to the open-source code repositories that were instrumental in the execution of experiments across spectrum of tasks in the paper. These concerted efforts reinforce the transparency and replicability of our research.
and the versions of critical software components including PyTorch,
REFERENCES
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD: Communication- efficient sgd via gradient quantization and encoding. arXiv preprint arXiv:1610.02111 October 2016. doi: 1048550/ARXXV.161000211
Mahmoud Assran, Nicolas Loizou, Nicolas Ballas, and Mike Rabbat. Stochastic gradient push for distributed deep learning. In International Conference on Machine Learning, pp. 344 353. PMIR, 2019.
Saar Barkai, Ido Hakimi, and Assaf Schuster. Gap aware mitigation of gradient staleness. arXiv preprint arXiv: /909. 1009, September 2019. doi: 10.48555/ARXIV.1900 10802.
Leighton Pate Barnes, Huseyin A. Inan, Berivan Isik, and Ayfer Ozgur. rTTp-k: A statistical estimation approach distributed sgd. arXiv preprint w071: May 2020. doi: 10.48550/ARXIV. 1076..0761.
Heinz H Bauschke, Jérôme Bolte, and Marc Teboulle. A descent lemma beyond lipschitz gradient continuity: first-order methods revisited and applications. Mathematics of Operations Research, 42(2):300-348, 2017.
Angel X. Chang. Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, Jianxiong Xiao, Li Yi, and Fisher Yu. ShapeNet: An Information-Ris Rich 3D Model Repository. Technical Report arXiv: 15Xiv:1512.005 151.00312 [cs.GR], Stanford University - Princeton University - Toyota Technological Institute at Chicago, 2015.
10

---

Published as a conference paper at ICLR 2024
Jianmin Chen, Xinghao Pan, Rajat Monga, Samy Bengio, and Rafal Jozefowicz. Revisiting distributed synchronous sgd. arXiv preprint 1604.06604009000 April 2016.
Alon Cohen, Amit Daniely, Yoel Drori, Tomer Koren, and Mariano Schain. Asynchronous stochastic optimization robust to arbitrary delays. arXiv preprint arXiv:2106.11879 June 2021. doi: 10. 48550/ARXIV:2006. 11879.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint arXiv:2110.02811, October 2021. doi: 10.4550/ARXI/221000226..
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 11929, 2020.
Sanghamitra Dutta, Jianyu Wang, and Gauri Joshi. Slow and stale gradients can win the race. IEEE Journal on Selected Areas in Information Theory, 103):1122-1024, sep 2021. doi: 10. 10.1109jssai. 2021.3 1103770
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training imagenet hour, arXiv preprint arXiv:1706.00 02677, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Qirong Ho, James Cipar, Henggang Cui, Seunghak Lee, Jin Kyu Kim, Phillip B Gibbons, Garth A Gibson, Greg Ganger, and Eric P Xing. More effective distributed ml via stale synchronous parallel parameter server. Advances in neural information processing systems, 26, 2013.
Anastasia Koloskova, Sebastian U. Stich, and Martin Jaggi. Sharper convergence guarantees for asynchronous sgd for distributed and federated learning, arXiv preprint arXiv:2206.08307, June 2022. doi: (0.48550/ARXIV222066
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 convergence speed. arXiv preprini arXiv:2104.0660 April 2021. doi: 10.48550/RRXV.2104..0
Mu Li, David G Andersen, Jun Woo Park, Alexander J Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene Shekita, and Bor- Yiing Su. Scaling distributed machine learning with the parameter server. In 11th USENIX Symposium operating systems design and implementation (OSDI 14), pp. 583-598, 2014.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. Pytorch Distributed: Experiences on accelerating data parallel training, arXiv preprint 1rXiv:2006.15744, June 2020. doi: 10.48550/ ARXIV.2006. 15704.
Shidi Li, Christian Walder, Alexander Soen, Lexing Xie, and Miaomiao Liu. Sampled transformer for point sets. arXiv preprint vrXiv:2302.14344. 2023.
Shigang Li and Torsten Hoefler. Near-optimal sparse allreduce for distributed deep learning. arXiv preprint arXiv:2201.077 075598, mar 2022. doi: 10.1145/5532221.55083
Tao Lin. Sebastian U. Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't It use large mini-batcher use local SGD. arXiv preprint 1xxiv:1808. 07217, August 2018. doi: 10.4555ARKIV.100.0000
Tsung- Yi Lin, Piotr Dollár, Ross Girshick, Kaiming He, Bharath Hariharan. and Serge Belongie. Feature pyramid networks for object detection. arXiv preprint 1rXiv.1612.03144 December 2016. doi: 085550/RRIVV.1122.0
11

---

Published as a conference paper at ICLR 2024
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized bert pretraining approach. arXiv preprint 1tXh::107..11992 July 2019. doi: 1044850/ARXIV.199.116999
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Reprexentations, 2018.
Yucheng Lu, Conglong Li, Minjia Zhang, Christopher De Sa, and Yuxiong He. Maximizing communication efficiency for large-scale training via 0/1 adam. arXiv preprint arXiv:2202.06009 February 2022. doi: 0.44850/ARXII.2202.000
Konstantin Mishchenko, Francis Bach, Mathieu Even, and Blake Woodworth. Asynchronous SGD beats minibatch SGD under arbitrary delays. arXiv preprint arXiv:2206.07638 June 2022. doi:
10.48550ARRXIV.220.0766
Yatian Pang, Wenxiao Wang, Francis EH Tay, Wei Liu, Yonghong Tian, and Li Yuan. Masked autoencoders for point cloud supervsu learning. In European conference on computer vision, pp. 604-621. Springer, 2022.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025-7041, Association for Computational Linguistics, December 2022a.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosFormer: Rethinking softmax in attention. In ICLR, 2022b.
Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling. In The Eleventh International Conference on Learning Representations (ICLR), 2023a.
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Fei Yuan, Xiao Luo, Yu Qiao, and Yiran Zhong. Scaling TransNormer to 175 billion parameters. 2023b.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. In NeurlPS, 2023c.
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Xiao Luo, Yu Qiao, and Yiran Zhong. LMan A faster and better large language model with improved transnormer, 2024a.
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightning Attention- 2: A free lunch for handling unlimited sequence lengths in large language models, 2024b.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Shaohuai Shi, Xiaowen Chu, Ka Chun Cheung, and Simon See. Understanding Top-k sparsification in distributed deep learning. arXiv preprini 1rXiv:1911.0877 November 2019. doi: 10.48550/
Sebastian U Stich. Local SGD converges fast and communicates little. 1805.09705 2018
arXiv preprint
Wenbo Su, Yuanxing Zhang, Yufeng Cai, Kaixu Ren, Pengjie Wang, Huimin Yi, Yue Song, Jing Chen, Hongbo Deng, Jian Xu, Lin Qu, and Bo zheng. GBA: A tuning free approach to switch between synchronous and asynchronous training for recommendation model. arXiv preprint arXiv:2205. / 1048, May 2022. doi: 10.48500ARXX1.220.11052
12

---

Published as a conference paper at ICLR 2024
Weixuan Sun, Zhen Qin, Hui Deng, Jianyuan Wang. Yi Zhang. Kaihao Zhang, Nick Barnes, Stan Birchfield, Lingpeng Kong, and Yiran Zhong. Vicinity vision transformer. IEEE Transactions on Pattem Analysis and Machine Intelligence, 2023.
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 convergence speed. arXiv preprint arXiv:2102.02888 February 2021. doi: 10,4850/ARXXV.22022228888
Xiaqiang Tang. Weigao Sun, Siyuan Hu, Yiyang Sun, and Yafeng Guo. MS-Net: A multi-path sparse model for motion prediction in multi-scenes. IEEE Robotics and Automation Letters, 9(1) 891-898, 2024. doi: 10.
3333443...
Zhenheng Tang, Shaohuai Shi, Xiaowen Chu, Wei Wang, and Bo Li. Communication--fficcnn distributed deep learning: A comprehensive survey. arXiv preprint arXiv:2003 2020.
Jianyu Wang and Gauri Joshi. Cooperative SGD: A unified framework for the design and analysis of local-update sgd algorithms. The Journal of Machine Learning Research, 22(1):9709 9958,, 2021.
Jianyu Wang. Vinayak Tantia, Nicolas Ballas, and Michael Rabbat. SlowMo: Improving communication-efficenn distributed sgd with slow momentum. arXiv preprint 111.0:1910 October 2019. doi:
Jianyu Wang, Hao Liang, and Gauri Joshi. Overlap local-SGD: An algorithmic approach to hide communication delays in distributed SGD. In ICASSP 2020 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, may 2020. doi: 10. 10.1 10.11099ias 4020.9053834.
Alexander Wettig, Tianyu Gao, Zexuan Zhong, and Danqi Chen. Should you mask 15% masked language modeling? arXiv preprint arXiv:2202.0800 2022.
Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3D Shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1912-1920, 2015.
Jian Zhang, Christopher De Sa, Ioannis Mitliagkas, and Christopher Ré. Parallel SGD: When does averaging help? arXiv preprint 1rXiv:1606.07365 June 2016. doi: 1..48550AAXXIV.166..0777
Wei Zhang, Suyog Gupta, Xiangru Lian, and Ji Liu. Staleness-aware Async-SGD for distributed deep learning. arXiv preprint 15Xv:1511.00954 November 2015. doi: 10445550ARXXI.1551.0033
Beitong Zhou, Jun Liu, Weigao Sun, Ruijuan Chen, Claire J Tomlin, and Ye Yuan. pbSGD: Powered stochastic gradient descent methods for accelerated non-convex optimization. In I/CAI, pp. 3258- 3266, 2020.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In The IEEE International Conference on Computer Vision (ICCV), December 2015.
13

---

Published as a conference paper at ICLR 2024
Appendix
A EXPERIMENT DETAILS
We show specific experiment plans regarding the tasks, models, their categorizations, parameter counts, and datasets used in Table
Table Tasks, Models, and Datasets implemented in the experiments. IC: image classification, SS: semantic segmentation, PC: Point Cloud, ALM: autoregressive language modeling, BLM: bidirectional language model. TN-LLM: TransNormer-LLMM.
Field
Task
Model
Parameters
Model Type
Dataset
ResNet-50
25.6M
ConvNet
ViT (Base)
86.6M
Transformer
ImageNet-Il Ik
VVT (Large)
61.8M
Linear Transformer
CV
VVT (Large)
65.5M
Linear Transformer
ADE20K
PC
Point-MAE
10.2M
Transformer
ShapeNet
GPT-2 (Small)
125M
Transformer (Decoder)
GPT-2 (Medium)
355M
Transformer (Decoder)
OpenWebText
ALM
NLP
GPT-2 (Large)
770M
Transformer (Decoder)
TN-LLM (7B)
6.7B
Linear Transformer
WikiText- 103
BLM
RoBERTa (Large)
355M
Transformer (Eneoder)
Bikippdia+BoooCrruus
A.1 HARDWARE AND SOFTWARE
Hardware. Our experimental setup comprises up to 16 DGX-A100 servers, with each server featuring 8 A100 GPUs These GPUs are interconnected via NVSwitch, providing an inter-GPU bandwidth 600GBps. Two distinct inter-node network configurations are used. The first config- uration leverages RoCE (RDMA over Converged Ethernet) solution, employing 8 RoCE RDMA adapters in each server to facilitate inter-server communication with a bandwidth capacity of 800Gips. The second configuration employs the TCP/IP protocol which operates at substantially reduced bandwidth compared to RoCE. Specifically, offers approximately 1/10th the bandwidth, totaling around 80Gbps. Part of experiments are conducted on 3090 GPU cluser with total 10 servers. Each server is equipmented by eight 3090 GPUs.
Software. Experiments implemented in PyTorch 1.13.0 with CUDA 11.7, cuDNN 8.0, and NCCL 2.1..3. Our algorithm is developed upon FairScale 0.4.13. The image classification ex- periments build on
image-models and https://github.com OpenNLPLabe DDenNLLLaa/Viciiitt Vision- Transformer The semantic segmentation experiments build on The 3D point cloud reconstruction experiments build on hitps: /ggthub.com/Pang Yatian/Point- MAE] The autoregressive language modeling experiments build The
https://github.cc httpl/giiiuuccom/haaiioiee wwwm/waa362//VV
on https://github.cc https:/gittub..om/karppaa wwt../aaaaaatynannnniiiii
and http::///uubbbooommccookeessrare bidirectional language modeling experiments build on http://ggtibbbooomm/mrretttotnnn Train
A.2 HYPERPARAMETER DETAILS
We detail hyperparameter settings for all tasks in the following:
Image Classification. The image classification experiments are conducted on ImageNet1 K dataset, which contains 1.28M training images and 50K validation images from 1, 000 categories. We adopt ResNet-50 (He et al., 2216; et al. 2020), VIT (Dosovitskiy et al., (020). and VVT (Sun et al.) 2023), which respectively correspond to convolutional network, transformer network, and linear transformer network. For the training of ResNet-50, the total mini-batch size is 8192, and the training runs 90 epochs. We sourced hyperparameters from DeepLeamiinExxmmpls The momentum value
14

---

Published as a conference paper at ICLR 2024
is set to 0.875 and the learning rate is initialized to 0.256 for a batch size of 256, the value is linearly scaled for batch sizes distinct from this reference. The learning rate schedule adheres to a cosine schedule. while a linear warm-up mechanism for the learning rate is introduced for the initial 5 epochs. We set the weight decay to 1/32768 and we do not apply weight decay on batch normalization trainable parameters. A label smoothing coefficient of 0.1 is adopted to enhance model robustness.
For the training of ViT and VVT, they closely adhere to analogous hyperparameters. A uniform total batch size of 2048 is employed, sustained across 300 epochs. A cosine learning rate schedule supplemented by linear warm-up during the initial 5 epochs is incorporated. A label smoothing coefficient of 0.1 is applied, while weight decay set 0.05. We use an initial learning rate of 5x 10 this value diminishes with cosine schedule, with 5 epochs dedicated to the warm-up phase. The augmentation follows previous literature, which includes practices such as random cropping and random horizontal flipping. Throughout the training, all models undergo 300 epochs of training on the training set, adopting a crop size of 224 x 224. The evaluation metric is the top-1 accuracy.
Semantic Segmentation. For the semantic segmentation task on the ADE20K dataset, the pre- trained ViT and VVT models undergo a finetuning process. The ADE20K dataset is composed of 150 distinet classes, distributed across 20210, 2000, and 3352 images for training, validation, and testing, respectively. We leverage the ViT and VVT models as the foundational backbone. already pre-trained on ImageNetIK. K. we adopt the Semantic Feature Pyramid Network (Semantic FPN) (Lin) et al. 2016) as the decoder. Hyperparameters governing the training of both ViT and VVT models are determined with reference to VVT.
3D Point Cloud Reconstruction. We trained the Point-MAE Pang al. 2022 Li et al. 2023 via
the reconstruction task using ShapeNet dataset (Chang et al. 2015). We further adopt the Chamfer distance as reconstruction loss, as well as a point number of 1024, a batch size of 256, epoch number of 300, cosine learning rate schedule, Adamw optimizer (oohchhlov & Hutter 2018) with learning rate of 0.0005, and weight decay of 0.05. The distributed training performance is evaluated not only by the training loss but also by the accuracy in the validation dataset of ModelNet40 (Wu et al. 2015). Detailed reconstruction pre-training and validation settings may be found in Pang et al. (2022):
ct al. (2023).
Autoregressive Language Modeling, For ALM tasks, we trained autoregressive GPT-2 mod- els (Radford et al. 2019) on Open WebTee dataset (Radford al. 2019]. Three GPT-2 models with 125M (Small), 355M (Medium), and 770M (Large) parameters are tested to prove the scalability of CO2. All experiments use the context length of 1024 and use a total batch size of 480. In addition, to further validate the scalability and efficiency of CO2, we conducted experiments large efficient language model, namely TransNormer-LLM, (7B) for 100k steps on small dataset WikiText-103 ZeRO-series optimizers are integrated with CO2 to lessen memory redundancy during training TaansNomer--LM. (7B).
Bidirectional Language Modeling. For BLM tasks. we trained a RoBERTa (Large) model on a combined dataset with Wikipedia We utilize a total batch size of 4096 and sequence length of 128. A polynomial learning rate decay scheduler was employed, with a peak learning rate of 2e-3 after 6% of total updates as warm-up.
Fnnd and BookCorpus (Zhu et al. 2015) for approximately 3 epochs.
A.3 ADDITIONAL EXPERIMENTAL RESULTS
A.3.1 TRAINING CONVERGENCE CURVES
In main paper, we show validation accuracy curves on classification task of three models, i.e., ResNet- 50, ViT and VVT, in Fig. 2 we show their respective training loss curves in Fig. [ Similarly, Fig 3 demonstrates the validation perplexity curves of GPT-2 with different sizes (Small, Medium and Large) on autoregressive language tasks. We show their respective curves of training perplexity in Fig. 6
These curves further validate the effectiveness of the proposed CO2.
hhtp::/dummps.wikimmdi...rg
15

---

Published as a conference paper at ICLR 2024
Resheb-30
StawMfe-SPP SSP
Stanki
Stowfo
C02
CO2
DO2
50
1.00
250
300
150
200
250
300
Epocte
Epacts
Epocht
Figure 5: Training curves for image classification tasks. Three models ResNet-50. VIT and VVT are trained on ImageNet- IK for 90, 300 and 300 epochs, respectively. Our CO2 exhibits robust convergence compared to other existing methods.
G773 (5mal)
GPT2 (Meckum)
GPT-
Asarme
Adamn Local Asurm Slanttc 003
Adamma Local edamm Skontfo
Looal Adammi
5 ceetto
Toah
50C
Figure 6: Training curves for autoregressive language tasks. Three variants of GPT-2 with 125M (Small), 355M (Medium), and 770M (Large) parameters are pre-trained Open WebText for 100K steps. Our CO2 exhibits the fastest convergence rate compared to other existing methods.
A.3.2 TIME TO ACCURACY CURVES
In order highlight the benefits of CO2 in accelerating the training process, give validation curves with respect to relative time (in seconds) for GPT-2 models of varying sizes as illustrated in Fig. 7 The outcomes demonstrate that CO2 excels only in achieving superior generalization performance but also in minimizing the duration of the training process, thereby affirming its practical utility in distributed training.
30.
G2T22 15ma8)
13
(FF12 (Meckam)
13
GPP12
13-
Asame
12
Adart
12
Local-daar
Addmix
VW
3.20
Neldvve Nim (5)
Aduliic Time
Figure 7: Validation curves w.r.t. relative time for autoregressive language tasks. Three variants of GPT-2 with 125M (Small), 355M (Medium), and 770M (Large) parameters are pre-trained on OpenWeTText WebText for 100K steps. Our CO2 exhibits the fastest convergence speed on relative time compared to other existing methods.
A.3.3 CONVERGENCE RESULTS WITH MORE DETAILS
To show advantages of CO2 on the trade-off of training efficiency and convergence performance, we list the corresponding throughput and settings of r for both CV and NLP tasks in Table 6 and Ta The presented throughput results represent averaged values extracted from iterations 100 to 200 for each experiment. Given the extensive nature of our experiments, convergence outcomes in the paper originate from two distinct clusters. Specifically, the IC and ALM tasks were executed
16

---

Published as a conference paper at ICLR 2024
Table 6: Convergence performance on CV tasks. CO2 performs better than other local methods with clear margin and is comparable to standard optimizers such as SGD//damw. For IC and SS tasks, we present the throughput results in images/sec, the image resolution 224 x224; For PC task, we present the throughput results in point clouds/sec, each point cloud has 1024 points.
SGD (Adamw)
Local- SGD(Adamw)
Overlap- Local-SGD(Adamv)
SlowMo
C02
Task
Model
Acc/ /Tpt Thpt
Acc/ Thpt/r
Acc/ Acc/Tbpt/r
Acc/ Acc/Tbpt/r
Acc / Thppt/
ResNet-50
76.92/ 108739
75.57 108758
76.28 108765 /1
77.12/ 108741 /1
77.14/ 108753
IC
VIT (Base)
81.33 39422
/ 39512 /8.43/39512// /
78.04/ 39511/3
79,83/ /9509
80:95 / 39533 /
VVT (Large)
83.64/ 44375
81.09 / 44390/ 1
80.33 / 44392/1
82.75 / 44376/
83.38 / 44387
S8
VVT (Large)
47.82/ 5384
44.25 / 5528 /6
45.21 5545 /6
47.51 /5521 /6
47.80/ 5562/
PC
Point-MALE
68.56 / 5859
64.25 5931/3
63.78/5950/ /
09.69/5904/
68.89/ 5956
Table 7: Convergence performance on NLP tasks. Quantitative perplexity (PPL) results for GPT-2, LoMmer-L and RoBERTa are presented. CO2 shows lower perplexity (lower is better) than the baseline Adamw in all these experiments. For ALM and BLM tasks, we present the throughput results in words/sec. TN-LLM represents TransNormer-L.M.
Overlap- Local-Adamw
Adamw
Local-Adamn
SlowMo
C02
Task
Model
Ace/ Tbpp
Ace / Thpt/
Ace/ Thp//pp//r
Ace/ Thp//hpt/r
Ace/ Thpt/
GPT-2 (Small)
7.44 6,543c6
7,95 6.556c6 3
8.11 6.55666 3
7.34 6,55446 /3
7.37 6,55666 /
GPT-2 (Medium) GPT-2 (Large)
6.61 2,08426
7.49/ / 2.09946 /
7,26 2.09266 /
6.41 2.099166 /
6.36 / 2,09226 /
ALM
6.26/ 1.052e6
7.00 1.05966 7
7.18 / 6
6.29 / 1.05366 6
6.13 1,05666 /6
TN-LLM (7B)
16.82/ 0.28
18.63 / 0.30366 / 12
17.83 / 0.306e6 / 12
16.95 0,30166 / 12
16.78 0,30866 12
BLM
RoBERTa (Large)
3,96 / 2262
4.38 4.38/ /815/815/6
4.52/ 2877 /6
3.98 /
3.95/ 2892/892/6
an A100 cluster equipped with RoCE RDMA high-speed inter-node connections, featuring 8 nodes and total of 64 GPUs. Conversely, the SS, PC, and BLM tasks were conducted on 3090 cluster with a standard TCP/IP Ethernet inter-node connection, comprising 4 nodes and 32 GPUs. Upon observing tasks trained on the A100 platform, it is evident that the throughput of CO2 surpasses other counterparts, although with slight advantage. This advantage becomes more pronounced on tasks trained on the 3090 platform. It reasonable infer that this advantage will be amplified on larger clusters, especially those with slower nnder-node connections.
Prior to commencing large-scale convergence experiments, we meticulously tune the values of r for each task using simple grid search strategy within the range of {1, 6, 12, 24, 48, 96, 192) to reach a balance of accuracy and throughput. Hyperparameter tuning experiments are conducted on smaller versions of models when the model to be trained is large, owing to the associated high time and resource costs. We start the tuning of T from to larger candidate values given the consideration of high-speed communication on the corresponding platform. It is worth noting that, in addition to CO2, other methods such as Local-SGD/Adam and SlowMo also require tuning of T. In practice, we exclusively tune for CO2 and employ the same tuned values for Local-SGD/AAama and SlowMo. This is because of the similar role played by r in these local-updating methods. Employing identical T values for different methods on the same task ensures a fair comparison of their convergence and training throughput.
A.3.4 COMMUNICATION COMPUTAATON OVERLAP RATIO
We conducted quantitative measurements on the time allocated to local computations and asyn- chronous communication. This assessment specifically targeted CO2 with different configurations on an A 100 cluster equipped with 128 GPUs, employing slower TCP/PP inter-node network. Adher- ing to the settings outlined in the paper, utilizing -rans- (7B), we maintained consistency in our experimental conditions. The measured duration for a single-step local computation was approximately 0.109s, while an all-reduce communication incurred cost of 1.566s. Subsequently, we computed and calculated the conmunication/comptttiin overlap ratio for various т values, presenting the results in the Table g
17

---

Published as a conference paper at ICLR 2024
Table 8: Communicatonn-compttttoo overlap ratio results. The varying values of and the corresponding overlap ratios on TransNormer-LLLM (7B).
6
12
48
Overlap Ratio
6.52%
20.39%
41.81 81%
83.28%
100%
100%
A.3.5 SCALABILITY RESULTS ON GPT-3 (13B)
To further effectively demonstrate the scalability advantages of CO2 on large models, conducted experiments specifically GPT-3 (13B). a well-known language model with autoregressive trans- former architecture. The results are shown in the Table 9 includes throughput comparisons with SlowMo. The results reveal that GPT-3 (13B). CO2 consistently achieves higher throughput platforms with different communication conditions.
Table 9: Throughput results on GPT-3 (13B). Throughput (words/sec) results of CO2, SlowMo and Adamw from 16 DGX-A100 servers are presented. With suitable configuration of r, CO2 can reach the highest throughput which outperforms SlowMo and the baseline Adamw.
Throughput
Ethernet
Method
GPUs
GPUs
32 GPUs
64 GPUs
GPUs
CO2
9468
18587
37905
75682
151846
RDMA
SlowMo
12
9463
18620
37467
75503
148987
Adamw
9519
19131
38824
75326
147723
CO2
48
9508
18672
38151
75573
151678
SlowMo
48
9424
18511
38092
75209
146024
TCP/IP
CO2
12
9324
9694
18792
36968
78638
SlowMo
9359
9628
18085
33864
66908
Adamm
9482
2283
4192
7984
13832
A.3.6 PENALTY PERFORMANCE ON NON-DECAY LEARNING RATE
To evaluate the performance of staleness gap penalty on non-decay learning rate, we conducted a comparative experiment on GPT-2 (Small) using the CyclicLR schedule provided in PyTorch with the triangular2 policy [ training for 100K steps. The test results and their comparisons with CosineAnnealingLLR are summarized in the table I0
Table 10: Performance of staleness gap penalty on CyclicLR. We pre-train GPT-2 (Small) with 100K steps for ablation. Both train and validation perplexity results indicate that staleness gap penalty is effective on CyclicLR schedule.
Method
Staleness Gap Penalty
Train PPL
Validation PPL
CO2 with CosineAnnealingL..
No
7.52
7.56
CO2 with Amnellnnne
Yes
7.36
7.39
CO2 with CyclicLR
No
7.58
7.63
CO2 with CyclicLR
Yes
7.45
7.51
The experimental outcomes consistently affirm the favorable impact of the staleness gap penalty tech- nique on training performance, regardless of the choice between Cosine and CyclicLR. While the enhancements PPL may not be substantial, they exhibit reliable positive trend. is noteworthy that PPL values using CyclieLR surpass those using CosineAnnealingLLR. potentially owing to sub-optimal learning rate tuning. Despite the cyclic nature of the learning rate CyclicLR, the staleness gap penalty continues to exert its positive influence. This could be attributed to the peri- odic decrease the learning rate during CyclicI_R's R's training iterations, facilitating the efficacy of the
18

---

Published as a conference paper at ICLR 2024
staleness gap penalty. It is pertinent to acknowledge that both tested schedules incorporate learning rate warm-up up during the initial 4000 iterations, which also represents form increasing learning rate. From technical standpoint, posit that there may exist specifie learning rate schedules where the staleness gap penalty technique may not exhibit optimal performance. In such instances, we advocate exploring alternative decaying learning rate schedules considering the option to disable the staleness gap penalty.
A.4 CONVERGENCE PROOF
The convergence proof is conducted for constant learning rate, i.c., "t = y and constant staleness gap, i.e., À = 1/A, and we ignore the outer momentum clipping operation make the analysis easy follow. Besides, we follows the standard assumptions (Wang et al. 2019) as below:
Assumption A.1 For some existed O, there [[V V s L||= ull, for all inputs 2, Rn and i E (1,2, ,G}.
Assumption A.2 For all i E (1.2,...,C G}, there exists a finite positive constant o² such that
Ex-D. | < () [()) S 0², ie., the variance of ft(x)) bounded.
Assumption A.3 There exists a finite positive constant V such that E l E [ge, S V.
Under these assumptions, the momentum update rule in Algorithm becomes:
m, -m,-1 A(x:-1.0 1.7)
(5)
Let us define:
1
gt.k
G
9ikk
(6)
Note that the local update rule is:
(6)
(1)
(8)
1,4+1
= 1,6
-9i.
(7)
So we have:
G
Σ
(
01,++1 =
G
x
1tt+1
f=1
Σ
=
G
x t.k
19th
(8)
Ki,k
19t.k
Øt,k-1
(gt.k 1- + 9).k)
k
Z1,0 -7 > gr.-
1=0
The momentum update rules then becomes:
m, = 8m + X(xx 1,0 - x.-),,) = -8m--1 + Xy
Σ
9
(9)
k=0
The outer update rule becomes:
1. F1,0 cTT4
T-1
F1,0 a,34-1 axy Σ g1-1k
(10)
1-1
= -1,0 - axy Σ
g + 3(1,,0 x1-1,0).
k=0
19

---

Published as a conference paper at ICLR 2024
We define:
Yt.0 = 24,0 +
(11.0 - 100
(11)
- B
So have:
B
B
y,+1,0 - Уt.0 = 0++1,0 - -,,00 +
1- (I++1.0 - -xt,0) - 1
(21.0
xc-
(I1+1.0
I1,0 (x1,0 xx11,00)
1 B
(12)
ady
gr. 1,k-
1
k=0
then define yo,r 1,+1.0 and
ad?
dn 3 9
Yt.k-
M1.k
91
-1
(13)
Since fü is L-Smooth, f Σ. ft is also L-Smooth. then according to L-Smooth of f and the Descent Lemma in Bauschke et al. (2117) we have
f(y..+1)= =f (m.
yykk
1 - B 91-1,k
(14)
Xy
X2
S f((r)))))
(ftck).9:-111. +
19. 19:-1.
- B
2(1 - .32
Taking expectation, then have
Ay
X2g2L
E[f(ree41) < f(y.x)
+
E [119
(15)
1-
2(1 B32²
20

---

Published as a conference paper at ICLR 2024
According to
15
and we define fin inf f(x)
(Wang et al. 2019
We have:
T
E f f [21,
Tr
1= 0=
2 (f (z0,0) fint) + mVL
E
VV (x), - Be, 1,
VGTT
TT
t=0 k=0
4G - 1)
(1
8GVL2T 132
-x-1
+
TT 2f (I0,a)
X
Tt (1- 32)
fint
+ mVV
VGTT
1 T-1r-1
E IIV (x
Ee, [gcck] + Be, [gc.k] - Be, [9c-1,1]]]
TT
t=0 K=
4GV L2
1)
1 1 3
-1)
L2V g²
+
Tt
X
TT (1 - 82)
(ʃ (20.0)
fint)
+ mVL
(16)
VGT
[V Ee xl1
TT
t= k=
Effect of Local SGD
2 T-1--
E (V (mI.k) B. (e) (Ee,, gg..k Ee l
TT
t=0 k=0
Effect of cross term
E E,,, 1 Ee 1.c
t=0 K=
Effect of asynchronous updates
4GV L2 1)
1 - 3
X -
L2T g2
+
Tt
Tt
(1 (2)
Efffect outer momentum
Note that Ee, [gc.k] results:
Ey, 191, the cross-term and asynchronous term eliminate, which
E (11, 5 2 (J (2 fint) + mVL
TT
VGTT
T-1T-1
E (( - Er, []l
L=0 k=0
(17)
Elfect of Local SGD
4GV IGVL?(T- 1) 1 3
8GVL2T B
Tt
Tt
(1 - 32
Effect of outer momentum
Then, finally have:
ΣΣ E IV (=)
(
(CTTT
+0
GT
(18)
TT 0=
T
21

---

Published as a conference paper at ICLR 2024
B DISCUSSIONS
B.1 DISTINCTIONS WITH EXISTED METHODS
The stale-synchronous parallel (SSP) framework (Ho et al. 2013 Tang et al. 2020) aims to mitigate the straggler problem with relaxed synchronization. Specifically, SSP allows faster workers to perform more updates than slower ones, reducing the waiting time of faster workers. However, to maintain model consistency and ensure convergence, SSP imposes staleness bounded barrier that limits the iteration gap between the fastest and slowest workers. Comparing with CO2, their differences mainly lie in these three aspects: 1) The SSP framework is built upon the Parameter-Serwee architecture. While CO2 as well as the Local-SGD and SlowMo algorithms investigated in the paper all based on the All-Reduce architecture. 2) The SSP-like algorithms indeed introduces asynchronization into its training procedure, however, it does not allow the communication to be overlapped with computations. And its staleness is not punished, which will leads unsatisfactory convergence performances. As comparison, the asynchronous communication in CO2 makes the communication can be overlapped with multiple-step local updates, which hugely improves the training efficiency, And the staleness gap penalty technique allows to mitigate the negative effects on convergence results. These imply the scaling efficiency and convergence performance of SSP-like algorithms will be worse than CO2 in practice. 3) The convergence analysis of SSP framework indicates that the convergence rate of SSP in normal cases ((+)) Comparing with CO2, when the total steps Tr sufficiently large, i.e,
So theoretically, the convergence rates of T G the convergence rate of CO2 is O VGTTT these two algorithms are similar.
The Cooperative SGD (Wang& & Joshi 2021 aims to unify variety of local-update SGD algorithms, such as local SGD, elastic averaging SGD, and decentralized parallel SGD, under the proposed unified Cooperative SGD framework. It also provides a unified convergence analysis for the methods under this framework. As one of the by-products, it proposes to add an auxiliary variable on the elastic averaging SGD, which serves as local copy of the model. It can be seen as preliminary approach that enables the asynchronous communication of model averaging overlap with computation. However, our method leverages more sophisticated asynchronous communication design to improve the convergence. The differences between our method and the Cooperative SGD are threcfold: I) The CO2 algorithm is built upon SlowMo, which involves slow momentum and outer updates to improve the convergence; 2) CO2 leverages more sophisticated asynchronous communication design together with the outer momentum and outer updates to hide more communications; 3) CO2 also introduces staleness gap penalty and outer momentum clipping to improve the convergence well as training stability.
The Overlap-Local SGD (Wang al., 2020) method is also able overlap communication by local computations via adding an anchor model on each node. The distinctions between CO2 and Overlap-Local-SGii are threefold: 1) Overlap-Local-SSOn
achieves communication-com computationn overlap on the top of naive -GD, whose convergence performance is normally not good enough. On the other hand, the CO2 is built on SlowMo, which involves outer loop includes outer momentum and outer updates to improve the convergence of native Local-SGD. 2) As mentioned by the reviewer, CO2 introduced staleness gap penalty and outer momentum clipping to prevent divergence and improve training stability, which has been verified effective via extensive experiments. is only tested on the CIFAR-10 10 image classification task using ResNet-
3) The Overlap-Local-SSNn 18, while CO2 has been widely verified multiple tasks both CV and NLP fields, including image classification, semantic segmentation, 3D point cloud reconstruction, autoregressive language modeling and bidirectional language modeling.
B.2 MEMORY OVERHEADS OF CO2
As an asynchronous method, CO2 leverages the advantages of a high overlap ratio facilitated by asynchronous communication, while contending with an augmented memory footprint due to the one-step asynchronous delay. Here we present a preliminary comparative analysis of the memory footprint of the methods discussed in the paper, using the widely adopted Adamw optimizer as the baseline and excluding considerations of mixed precision. Local-Adamw, a variant of Adamw, conducts local forward and backward passes and updates parameters without synchronizing gradients. It shares the same memory footprint Adamw but exhibits lower communication frequency and
22

---

Published as a conference paper at ICLR 2024
diminished accuracy performance. SlowMo, can be seen variant of Local-Adamw, introduces slow momentum calculation and outer model updates to enhance accuracy. Despite maintaining the same communication frequency Local-Adame SlowMo incurs additional memory footprint (twice that of the model itself) due to the introduced slow momentum and outer updates. In comparison, CO2 introduces one-step-delay outer momentum and outer model updates to overlap communication with multiple steps of local-updating computation, enhancing cnaununncationccopuaaaiiii overlap. This improvement comes at the cost of an additional memory footprint (twice that of the model itself) compared to SlowMo, attributed to asynchronous operations. In summary. So requires additional memory footprint which corresponds to 2 model copies than Local-Adamw, Then on the top of SlowMo, CO2 involves another additional memory footprint which corresponds to 2 model copies, so totally, CO2 needs additional memory overheads of 4 model copies comparing with the vanilla Adamw or Local-Adamw.
Despite the heightened memory overhead introduced by CO2, it enables full overlap (in the best case) of communication and computation, particularly advantageous for large clusters with limited inter-node connections. Techniques such as ZeRO-1, ZeRO-2, and ZeRO-3 have been incorporated to alleviate redundant memory footprint issues associated with optimizer states, gradients, and model parameters within the inner loop training of CO2. These measures could mitigate the adverse effects of CO2's augmented memory overheads.
B.3 COMPATIBILITY WITH PARALLELISM TECHNOLOGIES
Given the additional memory footprints associated with asynchronous updates, the compatibility of CO2 with various parallelism techniques becomes an indispensable consideration, especially the context of large-scale model training. During the implementation of the algorithm, we systematically investigated its compatibility with ZeRO-series optimizers. Currently, ZeRO-1, ZeRO-2, and ZeRO-3 integrated into the inner loop training of CO2, with the exclusion of the outer loop. We conducted thorough experiments to validate the effectiveness and efficiency of these integrations. In particular, with respect to ZeRO-3, which entails more sophisticated communication during both forward and backward passes, we had considered its compatibility with CO2. Given CO2's s local- local-updating behavior in the inner loop, continuous gradient synchronization is deemed unnecessary each step. As result, communication operations for CO2 with ZeRO-3 involve only two all all-gather collective communication primitives at the beginning of the forward and backward passes, respectively. which differs from the original implementation of ZeRO-3.
Concerning tensor parallelism, sequence parallelism, and their integration with CO2, while we have not implemented them at this moment, their mechanisms appear non-conficting with CO2's inner loop training. Tensor parallelism divides large matrix multiplications across multiple devices, and sequence parallelism divides the input sequence into chunks, processing them simultaneously on different devices for efficient training of long sequences. We will try to integrate them with CO2 in the future work.
23


Reviews:
[{"review_id": "NnaXV8SS9N", "reviewer": "Reviewer_8oQm", "summary": "This paper proposed a novel distributed training method: CO2, which can overlap communication and computation in distributed training. This technique is particularly useful when there are a large number of GPU nodes and the inter-connection between nodes are very slow. Compared to previous works, this paper introduces (1) penalty on stale momentum (2) momentum clipping. Empirical ablations show these two techniques are crucial to improve the training convergence performance. The authors also conducted extensive empirical studies, including experiments on image classification, large language model training, to demonstrate the effectiveness of the proposed method.", "questions": "See the above section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper has extensive empirical studies across different learning tasks as well as different network environments.\n- The authors also provided a convergence analysis for the proposed method.", "weaknesses": "- The idea of overlapping communication and computation is not new, as mentioned in the paper. The key contribution of this paper would be introducing the staleness penalty and momentum clipping mechanisms. They also present solid experimental resutls.\n- The comparison with previous works are not enough. For example, totally overlapping communication and computation has already been achieved. like Wang et al, 2020. Overlap-Local SGD. The authors should include more discussions on the differences. or even include this method as a baseline.\n- It is not very clear the convergence analysis was performed on which algorithm. Does the analysis consider staleness penalty and momentum clipping? Also, the convergence analysis looks like following previous works. It'd be better to cite few at the very beginning of the analyses.", "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": "JLc9xtFQ3p", "reviewer": "Reviewer_cSfz", "summary": "This work proposes a new distributed training algorithm called CO2, which aims to improve the communication efficiency of data-parallel training by overlapping local training iterations with parameter averaging from the previous global step. The proposed method is tested across multiple machine learning tasks and achieves better scalability than the baseline approaches while maintaining comparable convergence properties.\n\n---\nPost-rebuttal update: after reading authors' responses and other reviews, I decided to keep my weakly positive score unchanged and increase the confidence of my review. I think that the contributions of the study are solid and I am in favor of accepting the submission, but I am not fully sure that the work will have significant impact on the field in light of prior closely related publications.", "questions": "* What were the values of $\\tau$ for each experiment?\n* Which stage of ZeRO have you used for the TransNormer experiment?\n* In Table 2, it is somewhat surprising to see that CO2 (an asynchronous method) obtains consistently lower perplexity than a non-asynchronous adaptive method (AdamW). Do you have any explanations of that phenomenon?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Overall, the proposed method is conceptually simple yet shows promising results.\n* The paper has a broad range of experiments, covering 5 setups with models that are widely used in practice.\n* Authors conduct a detailed ablation study for the components of CO2, as well as measure its scalability in different environments.", "weaknesses": "* While I am not an expert in distributed optimization, to my understanding, similar methods allowing full overlap of communication and computation have been proposed previously. See, for example, [1] from the related work section: on page 17, they state that \"as long as the number of local updates τ is large enough, the communication can be fully overlapped with the local computation.\" This appears to be quite close to the primary contribution of this work, therefore I believe that the submission needs to describe the key distinctions from prior work in more detail.\n* I think that the experimental setup description could benefit from more details. For example, while the authors mention that their hyperparameters were tuned \"to find the optimal balance between efficiency and performance\", we do not see neither the exact values of $\\tau$ for each experiment nor the exact description of the tuning procedure. Also, authors mention that they leverage ZeRO for TransNormer experiments, but do not state the exact type of the optimizer within that family.\n* Lastly, the majority of model sizes used in this work have quite small parameter counts (fewer than 1B), and therefore it is a bit surprising to see communication as the bottleneck for training even on 80Gbps networks. I think that it would be beneficial to provide more detailed breakdowns of computation and communication times (for example, the time to process 1 microbatch and 1 batch of data, as well as the time to exchange parameters) in each setting to demonstrate the necessity of large $\\tau$.\n\n[1] Cooperative SGD: A unified Framework for the Design and Analysis of Communication-Efficient SGD Algorithms. Jianyu Wang, Gauri Joshi. JMLR 2021", "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": "mRne7UMitP", "reviewer": "Reviewer_vfwG", "summary": "The paper introduces CO2, a framework for improved communication/computation overlap for distributed deep learning, especially in the case of limited network bandwidth. CO2 leverages local SGD, performing a fixed (tunable) number of local iterations while allreduces perform synchronization in the background, allowing communication to almost always be hidden. To ensure good convergence, CO2 computes a staleness gap metric and uses this to scale updates, as well as a clipping mechanism to limit the variance of updates. A convergence bound is proven and experiments on a variety of network architectures and datasets show convergence matches that of standard SGD and other communication-avoiding algorithms; in the low-bandwidth regime, CO2 additionally offers significantly improved performance and scalability.", "questions": "1. I think the paper would be stronger if the claims of \"perfect 100% scalability\" were toned down and better contextualized. (See above for some details.)\n2. How should $\\tau$ be selected? Is hyperparameter tuning the only way to do so?\n3. How does the paper improve upon prior works which tune the communication frequency (see above for some references)? Could these approaches be used to tune $\\tau$ automatically?\n4. Please add implementation details and a discussion of memory overheads. I think memory may be especially relevant for larger models such as LLMs.\n5. Please add the observed variance to the accuracy results. It would also be good to include error bars in the scaling performance results.\n6. How do other models considered in the paper (e.g., ResNets or ViTs) scale?\n7. How do other communication-efficient (e.g., SlowMo) methods scale on the fast and slow network?\n8. How much communication/computation overlap is actually achieved by CO2, particularly at scale?\n9. A more minor point: The paper refers to gradient bucketing as a way to overlap communication and computation (e.g., in Section 1). I think this is not quite correct; rather, gradient bucketing is a latency/bandwidth tradeoff (performing fewer allreduces on larger buffers). While this can be more efficient, and consequently improve communication/computation overlap, it does not itself enable overlap.\n\n-----\n\nIn light of the authors' response and promised updates, I have raised my score. They have addressed a number of points above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. This paper is addressing an important situation: communication-bound training workloads. This can occur due to both large models and slower interconnects. I appreciate that the paper specifically and clearly calls out lower-bandwidth networks as an area it is focused on. While the idea is relatively straightforward, it includes some details to get it to work well in practice.\n2. The paper adequately specifies its proposed algorithm and includes some theoretical justification to support its claims.\n3. There are extensive experiments on a variety of models, including relatively large ones, demonstrating roughly equivalent convergence curves, indicating that the method does not compromise learning.\n4. Scalability studies are also conducted, showing slightly improved performance on high-bandwidth networks and significantly improved performance on low-bandwidth networks relative to a standard allreduce implementation.", "weaknesses": "1. I think the claims of \"perfect 100% scalability\" are a bit oversold. This relies on appropriately selecting $\\tau$, the number of local steps between global communications; it seems clear that if you can arbitrarily set the amount of computation done to hide communication, you can easily hide it. (Though I wish to be clear that the paper is clear that you can't make $\\tau$ arbitrarily high and still achieve good convergence.) This also neglects other aspects of training which may limit scalability (e.g., I/O for data ingestion).\n2. The paper does not provide guidance on selecting an appropriate $\\tau$, and in its experiments searches over a small set of potential values. This seems like a challenging parameter to tune in practice, as it could significantly increase hyperparameter tuning costs.\n3. It is not clear to me how the paper improves upon existing communication-efficient works which try to tune the communication frequency to achieve both good learning and runtime performance. In particular, works like Wang & Joshi, \"Adaptive Communication Strategies to Achieve the Best Error-Runtime Trade-off in Local-Update SGD\", or Haddadpour et al., \"Local SGD with Periodic Averaging: Tighter Analysis and Adaptive Synchronization\", seem like relevant points of comparison.\n4. The paper lacks implementation details. Specifically, it does not specify how the asynchronous allreduce is implemented (e.g., is it using a NCCL allreduce on a separate CUDA stream?). It is also not clear whether the asynchronous allreduce is operating on a separate weight/gradient buffer from the one being used for computation; or what the memory overheads of the method are.\n5. While I appreciate that the experiments were run multiple times (Section 4.1), the results do not include any measure of variance. This makes it hard to understand whether CO2 amplifies the variance between runs and how much methods actually differ.\n6. Scalability is only evaluated on one model. I would be interested to see how models other than the TransNormer-LLM scale; in my experience, smaller models tend to benefit less from communication optimizations as they are already often able to hide most communication.\n7. The scaling study in Section 4.3 does not include any comparisons with other communication-efficient methods. Given that SlowMo demonstrates very similar convergence curves, it seems prudent to see whether CO2 offers better scalabiltiy.\n8. From a performance perspective, the paper is missing a detailed analysis substantiating its claims. In particular, the communication/computation overlap achieved is never actually measured.", "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": "cvWcLhf3eK", "reviewer": "Reviewer_m91N", "summary": "To address the communication problem in large-scale distributed training of deep neural networks, the paper proposes a combination of local-SGD and asynchronous communication to derive a new distributed training algorithm named CO2. In CO2, two novel approaches are developed to ensure that CO2 aligns the convergence performance with conventional distributed data-parallel algorithms. Experiments are conducted on a 64-GPU testbed, showing that CO2 outperforms existing methods significantly. The studied problem is timely and important. The paper is also well-written.", "questions": "- How about comparing with SSP-like algorithms in terms of theoretical convergence bound and empirical scaling efficiency? \n- How $\\tau$ is set in Table 1?\n- How about the end-to-end training performance (i.e., time to accuracy)?\n- How to choose $\\tau$ in a new distributed GPU cluster?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Propose a new distributed training algorithm, CO2, using local updates and asynchronous communication to alleviate the communication problem in conventional synchronous data-parallel distributed training. \n- New tricks to address the convergence problem in stale gradients.\n- Comphesive experiments to show the effectiveness of CO2.", "weaknesses": "- Some stale parallel algorithms (e.g., SSP [ref1]), whose key ideas are quite similar to CO2, were not included in the discussion and comparison. The survey paper [ref2] may help find SSP-like methods for comparison.\n- It seems that 100% scaling efficiency is over-claimed. The scaling efficiency highly depends on $\\tau$. Higher $\\tau$ has better scaling efficiency but has worse convergence performance. Thus, achieving 100% scaling efficiency with a $\\tau>1$ while sacrificing the convergence performance cannot conclude the algorithm has true 100% scaling efficiency.\n\n[ref1] More Effective Distributed ML via a Stale Synchronous Parallel Parameter Server, NeurIPS 2013.\n[ref2] Communication-efficient distributed deep learning: A comprehensive survey, arXiv 2020.", "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": "tQfITz6G7v", "reviewer": "Reviewer_rgZH", "summary": "The paper proposes an approach called CO2 to improve throughput of distributed model training by overlapping computation and communication. Building on prior work that perform multiple training iterations with local updates before each global model synchronization, CO2 enables further throughput improvement by making the global synchronization asynchronous and overlapped with the next round of local updates. CO2 proposes two techniques for addressing the convergence issues of the asynchrony: (i) staleness penalty gap, and (ii) outer momentum clipping. The paper presents theoretical analyses of the convergence guarantees of these two techniques. The evaluation results show that CO2 can achieve convergence results comparable to baselines that are fully synchronous (e.g, Adam) and better than those using local updates (e.g, LocalSGD). The experimental results also show the throughput and scalability of CO2 are better than Adam.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper is tackling an important problem since communication is a major bottleneck for scaling model sizes and training hardware, and so approaches for reducing communication overheads are very relevant to the community. \n\nThe idea of overlapping communication with computation is reasonable given the cost-effectiveness. I also liked the fact that the paper attempts to quantify and fix the resulting staleness update problem. \n\nThe evaluation considers a diverse and important set of workloads and hardware environments, which helps to understand the generality of CO2.", "weaknesses": "I observe some critical problems in the draft that raise the question of whether CO2 can simultaneously achieve good convergence and high throughput. \n\n1. The convergence and throughput trade-off of inner loop step count ($\\tau$) is not clearly reported in evaluation. In particular, the convergence results in Tables 1 & 2 should include the corresponding $\\tau$ and throughput. I was unable to determine whether the good convergence results are achieved with $\\tau$ that also provides throughput benefits. \n\n2. The paper is silent on the memory overheads of CO2 relative to baselines, even though Algorithm 2 suggests that multiple copies of the model is required to support asynchronous communication. \n\n3. Equation 3 assumes learning rate decay in the inner loop which is not true for learning rate schedules, such as cyclic, which involve learning rate increases. \n\n4. It is unclear to me whether CO2 can achieve expected throughput benefits in scenarios with parallelism techniques (e.g., tensor slicing, sequence parallelism, and zero stage 3) that introduce communication to forward/backward passes. It seems these (synchronous) communication operations would interfere with the overlapped communication and hurt overall throughput. Evaluating such scenarios could help to better understand the generality of CO2.", "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": "MBHcGwzFcV", "reviewer": "Reviewer_Qwcz", "summary": "The paper proposes CO2, a new approach that enables efficient distributed training of large language models on clusters with limited bandwidth. CO2 introduces local-updating and asynchronous communication to the distributed data-parallel training, allowing for full overlap of communication with computation. The approach achieves 100% scalability even on clusters with limited communication bandwidth. The paper also introduces staleness gap penalty and outer momentum clipping techniques to improve convergence and training stability. The proposed approach is validated through extensive experiments on computer vision and natural language processing tasks as well.", "questions": "How many runs for each task? I understand that training a Language Learning Model from scratch can be quite costly. However, conducting the experiment only once may not yield persuasive results.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "+ The paper is well-written and comprehensible.\n+ The code is available in this work.\n+ The utilization of local updating and asynchronous communication makes a full overlap of computation and communication. \n+ The paper provides enough theoretical explainability and empirical validation. \n+ The experimental results are sound and promising.", "weaknesses": "I do not have much to comment on the weakness, as this work goes beyond my acceptance threshold.", "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 (spotlight)