1-bit LAMB: Communication Efficient Large-Scale Large-Batch Training with LAMB's Convergence Speed
Authors: Conglong Li
Ammar Ahmad Awan
Hanlin Tang
Samyam Rajbhandari
Yuxiong He
Venue: ICLR 2022 Submitted (2022)
Keywords: optimization
communication compression
natural language processing
language model pre-training

Abstract: To train large models (like BERT and GPT-3) on hundreds of GPUs, communication has become a major bottleneck, especially on commodity systems with limited-bandwidth TCP network. On one side large batch-size optimization such as LAMB algorithm was proposed to reduce the frequency of communication. On the other side, communication compression algorithms such as 1-bit Adam help to reduce the volume of each communication. However, we find that simply using one of the techniques is not sufficient to solve the communication challenge, especially under low network bandwidth. Motivated by this we aim to combine the power of large-batch optimization and communication compression, but we find that existing compression strategies cannot be directly applied to LAMB due to its unique adaptive layerwise learning rates. To this end, we design a new communication-efficient algorithm, 1-bit LAMB, which introduces a novel way to support adaptive layerwise learning rates under compression. In addition, we introduce a new system implementation for compressed communication using the NCCL backend of PyTorch distributed, which improves both usability and performance. For BERT-Large pre-training task with batch sizes from 8K to 64K, our evaluations on up to 256 GPUs demonstrate that 1-bit LAMB with NCCL-based backend is able to achieve up to 4.6x communication volume reduction, up to 2.8x end-to-end time-wise speedup, and the same sample-wise convergence speed (and same fine-tuning task accuracy) compared to uncompressed LAMB.

Under review as a conference paper at ICLR 2022
1-BIT LAMB: COMMUNICATION EFFICIENT LARGE- SCALE LARGE-BATCH TRAINING WITH LAMB'S CONVERGENCE SPEED
Anonymous authors Paper under double-blind review
ABSTRACT
To train large models (like BERT and GPT-3) on hundreds of GPUs, communication has become major bottleneck, especially commodity systems with limited- bandwidth TCP network. On one side large batch-size optimization such as LAMB algorithm was proposed to reduce the frequency of communication. On the other side, communication compression algorithms such as 1-bit Adam help to reduce the volume of each communication. However, we find that simply using of the techniques is not sufficient to solve the communication challenge, especially under low network bandwidth. Motivated by this aim to combine the power of large-batch optimization and communication compression, but we find that existing compression strategies cannot be directly applied to LAMB due to its unique adaptive layerwise learning rates. To this end, we design a new communication- efficient algorithm, 1-bit LAMB, which introduces novel way to support adaptive layerwise learning rates under compression. In addition, we introduce new system implementation for compressed communication using the NCCL backend of PyTorch distributed, which improves both usability and performance. For BERT- Large pre-training task with batch sizes from 8K to 64K, our evaluations on up 256 GPUs demonstrate that 1-bit LAMB with NCCL-based backend is able achieve up to 4.6x communication volume reduction, up to 2.8x end-to-end time- wise speedup, and the same sample-wise convergence speed (and same fine-tuning task accuracy) compared to uncompressed LAMB.
INTRODUCTION
Training large-scale deep learning models in distributed fashion is computation-heavy and ex- pensive (Brown et al. 2020). In addition to computation, communication overhead becomes a serious system challenge for such large models. A recent study of BERT pre-training with Adam demonstrates that the allreduce communication can take to 94% and 75% of total training time per step on clusters with Ethernet and InfiniBand inter-node connections, respectively (Tang et al. 2021).
To achieve communication efficient distributed training, there are two promising directions: large batch optimization and communication compression. LAMB optimizer, which can be viewed as Adam with adaptive layerwise learning rates, is an example of large batch optimization (You et al. 2020). LAMB can scale the batch size of BERT pre-training to 64K without losing accuracy, thereby greatly reducing the total training time as larger batch sizes leads to less frequent communication, On the other hand, recent works on communication compression such as 1-bit Adam demonstrate that it is possible to combine 1-bit compression with Adam's 's convergence speed, thereby reduce BERT pre-training communication volume by 5x (Tang et al. 202T).
Both LAMB and 1-bit Adam demonstrate great benefit for distributed training. Unfortunately. our studies show that simply using one of them is not sufficient to fully address the communication issue, especially under limited network bandwidth and large number of GPUs/machines (Section 3). We find that communication is still a non-trivial overhead when running large-scale scale distributed training with LAMB, convergence speed compared to LAMB at batch sizes 16K or larger for BERT pre-training (You) et al. 2020). Using the same methodology, BERT experiments show that 1-bit Adam, similar
with the larger batch sizes. Previous study shows that Adam provides slower

---

Under review as a conference paper at ICLR 2022
to Adam, also has slower convergence speed compared to LAMB at batch size 16K. Even with the communication compression, this batch size limitation would hurt the communication efficiency when the number of GPUs/machines is large.
LAMB and 1-bit Adam are two unique optimizers. However, the techniques behind them are complementary: large batch optimization reduces the frequency of communication, and compression reduces the volume of communication. Motivated by this we aim to combine LAMB's large batch optimization algorithm with compression strategies behind 1-bit Adam. However, we find that they are not directly compatible due to LAMB's unique layerwise learning rate update strategy, which requires information that are missing when communication and optimizer states are compressed (Section
The studies and challenges above motivate us to design new algorithm called 1-bit LAMB (Sec- tion 4). Learning from the insights behind 1-bit Adam, 1-bit LAMB is a 2-stage algorithm which uses LAMB (warmup stage) to "pre-condition" a communication compressed momentum SGD algoirthm (compression stage). At compression stage where original LAMB algorithm cannot be used to update the layerwise learning rates, 1-bit LAMB employs a novel way to adaptively scale layerwise learning rates based on information from both warmup and compression stages. As result, 1-bit LAMB is able to achieve large batch optimization (LAMB)'s convergence speed under compressed communication, which is impossible using existing approaches.
In addition to the 1-bit LAMB algorithm, we propose a new NCCL-based compressed communication backend which provides better usability and performance than previous work (Section (5) This back- end can be applied to 1-bit LAMB, 1-bit Adam, and other communication compression algorithms. We evaluate I-bit LAMB using BERT pre-training and GLUE/SQuAD fine-tuning tasks (Section (6) Results show that under different batch sizes from 8K to 64K and with up to 256 GPUs, 1-bit LAMB with NCCL-based backend is able achieve up to 4.6x communication volume reduction and up to 2.8x end-to-end time-wise speedup for BERT pre-training compared to uncompressed LAMB, together with the same sample-wise convergence speed and same GLUE/SQuAD fine-tuning task accuracy. The 1 1-bit LAMB optimizer as well as the NCCL-based communication backend has been open sourced in deep learning optimization library (name hidden to maintain anonymity).
2 RELATED WORK AND BACKGROUND
To achieve communication efficient distributed training, techniques include decentralization (Lian et al. 2017; Koloskova* et al. 2020 Li et al.) 2018), asynchronous communication (Zheng et al. 2016 Chaturapruek et 2015), and gradient compression/quanizztion which we focus on in this paper. Before communication, we could compress the original gradient g into Cg]g). where [. is the compress operator As a result the communication volume could be greatly reduced.
Compression can be achieved by quantization, sparsification, sketching, etc. (Ye & Abbe 2018
Alistarh et al.] 2017 AAgaral et al. 2018: Yu et al.] 2019) Spring et al., 2019 Ivkin et al. 2019
Shi et al. 1021). Several works focus unbiased compression methods (original and compressed tensors have the same expectation), such as centralized compressed parallel SGD AAlistarh et al.
2017) and many others (Wangni al.. 2018: [Shen et al.) 2018 Zhang et al., 2017: Wen al. 2017: Jiang& & Agrawal (2018). On the other hand, recent works about biased compression methods demonstrate better compression rate and the same convergence rate by using an error compensation
technique (Seide et al.
2014; Bernstein et al. 2019 Stich et al. 2018: Zheng et al. 2019 Phuong &
Phong 2020 Yu et al.
2019 Shi al. 2019: Ivkin et al. 2019 Sun et al. 2019
Basu al. 2019
Vogels et al. 2019 Tang et al. 2021).
The correscompenstted compression is proposed in the 1-bit SGD work (Seide al. 2014): 4): instead of compressing the gradient at each iteration directly, they compress the sum of the gradient and the last step's s compression error. By using error compensation the training can achieve promising convergence speed even with 1-bit compression (representing the gradient by +1 signs and scale). Recent works provide theoretical guarantee of this method (Bernstein et al., 2019), and also demonstrate that admits the same asymptotic convergence rate as the uncompressed one (Sich et al. (2018). In addition, error compensation method enables almost any compression methods (Stich et al. 2018), either biased or unbiased, to converge as fast as the uncompressed case.
[+) could also include randomness.
2

---

Under review as a conference paper at ICLR 2022
Adam (Kingma & Ba 2013) can be viewed SGD with momentum and adaptive learning rate scaling on each coordinate of the gradient. has demonstrated promising convergence speed and hyperparameter robustness
many deep learning tasks. Recently, Tang et al. (202I) proposed 1-bit Adam which combines the efficiency of error-compensated 1-bit compression with Adam's convergence speed. They show that emoo-coompensated compression does not work for Adam directly, because Adam is non-linearly dependent on the gradient (the variance term). On the other hand, they find that Adam's variance becomes stable an early stage of training. To this end, they design new 2-stage algorithm, 1-bit Adam: At warmup stage, vanilla Adam is used. At compression stage, they stop updating the variance and use it fixed precondition, and communicate based on the momentum applied with ertor-compensated 1 1-bit compression. Their experiments on up to 256 GPUs show that I-bit Adam achieve the same convergence behaviour and final accuracy as Adam, together with up to 5x less communication volume and 3.3x faster end-to-end throughput.
To further improve training efficiency large scale, being able to support large minibatches while keeping the convergence speed is critical factor. Recently You et al. (2020) find that it is difficult to keep Adam's convergence speed at batch sizes 16K or larger for BERT pre-training. To this end they proposed LAMB which can be viewed Adam with adaptive layerwise learning rates. By using LAMB, they are able to scale the batch size BERT pre-traie pre-training to 64K without losing accuracy, thereby, reducing the BERT training time from 3 days to around 76 minutes. The major idea of LAMB is that it utilizes a layerwise scaling coefficient to regulate the update of each layer, and the updating rule can be summarized as
m
(1)
= mm + (1 B1)9. .u
11)
=B20 + -1-1111 (9:
(1.
u
m
.g = -clip
|
Cmins Craz
(1)
[10
1 +
(1)
!1)
(1) ( 76 ui
Here 9: FF Ex), mi ve x denote the stochastic gradient, momentum, second moment (i.e., the variance), and the model parameters at the model's (-th layer at step f; 31 and 83 are the decaying factor; y is the learning rate; 17 an additive constant avoid division by 0; clip(z, a, b) := min(max a), b} is the clipping operation
(1)
is a layer-wise scaling factor that regulates the update of x into certain range. One thing note is that within each layer, each tensor (e.g., weight and bias) will have its own scaling coefficient ct The underlying intuition of LAMB's scaling coefficient is that when the update relatively large compared the parameter, we should apply lower learning rate to that layer (and vice versa).
MOTIVATION AND INSIGHTS
3.1 1-BIT ADAM IS NOT SUFFICIENT FOR LARGE-BATCH DISTRIBUTED TRAINING
1-bit Adam demonstrates the same convergence speed as Adam for BERT pre-training task with batch size 4K (Tang al., 202T). On the other hand, the LAMB work shows that it is difficult to keep Adam's 's convergence speed at batch sizes 16K or larger for BERT pre-training (You et al.) 2020). To find out whether 1-bit Adam is sufficient for largh-batch distributed training, we perform similar experiment using BERT pre-training task at batch size 16K. Using You al. (2020)'s 's training parameters and tuning procedure (details in Appendix A. for LAMB and Adam, we perform BERT pre-training with LAMB and 1-bit Adam, respectively. Then use the two pre-trained BERT model to perform SQuAD 1.1 fine-tuning (details in Section (6. Results in Table T show that similar to Adam, 1-bit Adam has slower convergence speed compared LAMB at larger batch size.
Here (a=),, v and z all denote element-wise operations. For simplicity weight decay is omitted. In the LAMB paper the clip function is only applicd to || || without mentioning the exact clipping function configurations, and our experiments show that ["x varies lot among different layers. Thus we apply the clipping function to the whole ratio, which is more stable among different layers. With this clipping
function able to achieve similar SQuAD accuracy compared to the original LAMB.
3

---

Under review as a conference paper at ICLR 2022
Table 1: BERT-Large pre-training (batch size 16K) final validation loss, and SQuAD average/max
dev set Fl scores over 32 runs using the pre-trained BERT models. The first two columns are from the original LAMB work. The last two columns are our experiments using the same training parameters.
BERT pre-training optimizer
LAMB You et al. 2020) Adam ( You et al. 2020)
LAMB
I-bit Adam
BERT validation loss
1.362
1,504
SQuAD Avg. FI
90.716
89.225
SQuAD Max F1
91.345
88,540
91.119
89.494
Preelto Leperierm
3K4b
0.4
Aln axtort
20
0.2
10-
LANB
18 25
35
66
16
3
3K
4K
3X
Steps
scaling coefficients
(h) variance norms
Figure I: LAMB's scaling coefficients (e)" in 1 and variance norms for different layers in the first BertLayer during BERT-Large pre- training seqlen 128 (5993 steps in total). We set the lower/upper bound of the scaling coefficient (Cmin and Craz in at 0.01 and 0.3.
3.2 COMMUNICATION IS STILL AN OVERHEAD FOR LARGE- BATCH DISTRIBUTED TRAINING
Tang et al. (202)) demonstrate that when pre-training BERT with Adam at batch sizes 64 to 4K, the communication may take up to 94% and 75% of total training time per step clusters with Ethernet and InfiniBand networks. To investigate whether the communication overhead still exists, we conduct similar profiling experiments but with LAMB optimizer and batch sizes 8K to 64K. We evaluate two clusters: one with 4 NVIDIA V100 GPUs per node and 40 Gigabit Ethernet inter-node network (4.1 Gbps effective bandwidth); the other with 8 V100 GPUs per node and 100 Gigabit InfiniBand EDR inder-node network (close to theoretical peak effective bandwidth). Results show that even with larger batch sizes the communication still contributes up to 91 % and 52% of the training time on two clusters (details in Appendix A.2). indicating the opportunities to improve large-batch distributed training efficiency by communication compression.
3.3 INVESTIGATING BERT PRE-TRAINING WITH BASELINE LAMB
Tang et al. 1022 find that Adam's 's variance term becomes stable an early stage of training (after around 15% of total training for BERT), which is why 1-bit Adam can "freeze" the variance after the warmup stage and use it as a fixed precondition during compression. LAMB also has the variance
term as a non-linearly gradient dependency, and LAMB's scaling coefficient (c, in (T) depends on the variance. Thus we investigate how LAMB's scaling coefficient and variance change during training using BERT pre-training task.
Figure D presents LAMB's scaling coefficients and variance norms for different layers in the first BertLayer (other layers in the model have similar patterns, details in Appendix A.3). Results demonstrate that the scaling coefficients keep increasing until reaching plateaus, because the update tends to become smaller compared to the parameter during training. In addition, many scaling coefficients become stable at an early stage. Results also demonstrate that LAMB provides adaptive learning rate in two folds: 1) different layers may reach scaling coefficient plateaus at different time; 2) different layers may reach different scaling coefficient plateaus. We believe that this is one of the reasons why LAMB can provide better convergence (or reach the same convergence with less hyperparameter tuning) at larger batch sizes compared with Adam. On the other hand, LAMB's s varaince terms are less stable compared with Adam: many layers have their variance norms constantly decreasing during the whole training, up to two orders of magnitude difference. As we see in next section this makes b-bit Adam's 's strategy affect convergence when directly applied to LAMB.
4

---

Under review as a conference paper at ICLR 2022
Table 2: BERT-Large pre-training (batch size 64K/32K seqlen 128/512) final validation loss, and
SQuAD average/max dev set Fl scores over 32 runs using the pre-trained models. "LAMB + basic
1-bit" is the experimental algorithm described in Section 3.4
1-bit LAMB" the proposed work
described in Section
BERT pre-training optimizer
LAMB
You et al. 2020
LAMB (ours)
LAMB + basic 1-bit
1-bit LAMB
BERT validation loss
1,451
1.494
1.443
SQuAD Avg. F1
90.265
90.069
90.524
SQuAD Max FI
90.584
90.555
90.409
90.788
3.4 EXISTING COMPRESSION STRATEGY AFFECTS LAMB'S CONVERGENCE
To combine 1 1-bit compression and large-batch training's communication efficiency, first we attempt directly apply 1-bit Adam's 's strategy LAMB. We design and evaluate an experimental two-stage algorithm "LAMB + basic 1-bit": At the warmup stage we use vanilla LAMB. At the compression stage we stop updating the variance and LAMB's s scaling coefficient and use them as precondition, and communicate based on 1-bit compressed momentum (details about this algorithm design in Appendix A.4). For simplicity we only apply this experimental algorithm to BERT pre-training seqlen 128 phase, and still only use vanilla LAMB in seqlen 512 phase.
Table presents the BERT pre-training final validation loss when using vanilla LAMB, the experi- mental algorithm described in this section, and the proposed 1-bit LAMB algorithm (Section 4]. We also use the pre-trained models to fine-tune tune SQuAD 1.1 and present the F1 scores. Results show that simply freezing both variance and LAMB's scaling coefficient would affect the convergence speed and lead to lower SQuAD scores (and lower GLUE scores in Section (6) This is because LAMB's variance term is less stable demonstrated in our study. On the other hand, the proposed 1-bit
LAMB is able to provide the same convergence speed as vanilla LAMB, and next we will describe its algorithm design.
4 I-BIT LAMB ALGORITHM
The proposed 1-bit LAMB optimizer introduces a novel way to update the adaptive layerwise learning rate during the compression stage. There are two major differences between 1-bit LAMB and the original LAMB: 1) During compression stage, 1-bit LAMB updates the layerwise learning rate based novel "reconstructed gradient" based on the compressed momentum. This makes 1-bit LAMB compatible with error compensation and be able to keep track of the training dynamic under compression. 2) 1-bit LAMB also introduces extra stabilized soft thresholds when updating layerwise learning rate at compression stage, which makes training more stable under compression.
Problem setting
In this paper, we focus on the following optimization task:
min
f(x)
E Bew F(x: (()
ERR
=ff((x))
where d the dimension of the input model æ, D; is the data distribution of individual data sample g(9) on the i-th worker, F(x; is the loss function.
Notations and definitions Throughout this paper, we use the following notations:
Vf(-) denotes the gradient of a function f.
fz()) == ((;
|| llp denotes the n,-norm for vectors and matrices. Notice p means the infinity norm.
(-) denotes the randomized compressing operator, where denotes the random variable. One example is the randomized quantization operator, for example, (o.7).) with probability 0.7 and (0.7) 0 with probability 0.3.
V- denotes the square root of the argument. In this paper if the argument is vector, then it returns vector taking the element- square root.
(z)) denotes the element- wise square operation if a is vector.
5

---

Under review as a conference paper at ICLR 2022
a/b denotes the element- wise division operation if both and b are vectors and their dimension matches.
Algorithm 1-bit LAMB
(1) (9) 1: Initialize: Tg m O, 0, (2) = for each layer. Learning rate initial error å = O, numher of total iterations T, warm-up steps Tw,, three decaying factor B1, B3, ,83 for LAMB's momentum, variance, and scaling coefficient. ,()) = 1. Fmin. Tmaz. Fthreshold for 1-bit LAMB. We use 1-bit compression as the compressing operator, which means Cu [ign(w]]] keeps the sign of each element in the original tensor
|=| Sign(æ) where Sign(æ) is the tensor that only
2: Running the original LAMB in for Tw steps, and each step Goug
(1) Paraeg + Sa)e)
3: At the end of step Tw, for each layer store the variance term (defined as v, (!) in (T) !! while still keep
updating 2i in the future steps. Also stop updating caug.
4: for t = TT.. .......? do
5: (On i-th node)
6: Randomly sample 5,² and compute local stochastic gradient g: = V (m: .((" and update the
local momentum m, according to m, (c) = f. + (1 B1)g
7: Compress the fused momentum m into m. C. m, _ and update the compression
(4)
error by σ! m, 6,2 m
8: Send the ric. to the server.
9:
(On server)
Take the average over all m, (+) it receives and compress it into C. [ - mi 8 and
10:
update the compression error accordingly by 8. Ej= m, (4) 8, me. Send m. to all the workers.
11:
12:
(On j-th node)
13:
Set m. 7/ for the 1-th layer
8.m (.)(111(1- -2). Reconstruct global gradient g. =mm!"
(1)
82)
91 (1)
17:
(1) Tw
/02
18:
r (1) clip
rr (1 (ehreshold) x xr r + Tihreshola) x r((2))
(1)
19:
Jt
(1) clip Pt .Tmin: maar
(1)
20:
G (t) (1) (i)
Cavg- (
21:
Update model of the I-th layer at!" 20 - 704
(0) m
"Tw
22:
end for end for
23:
24:
Output: ac
We summarize 1-bit LAMB in Algorithm I During the compression stage, we freeze the variance in order to apply the eиror compensation mechanism correctly (Tang et al. 2021). However, this brings
two challenges for LAMB's s case: 1) We cannot update LAMB's 's scaling coefficients (c in ([) during compression stage based LAMB algorithm, because it requires uncompressed momentum and up-to-date variance: 2) In vanilla LAMB the scaling coefficients become stable during training. However, because LAMB's 's variance term is stable for some layers but we freeze them during compression, we need to adjust scaling coefficients during compression compensate this.
To this end, 1-bit LAMB uses a novel way to adaptively update LAMB scaling coefficients during the compression stage to compensate the difference between frozen and actual variance. During the warmup stage, we use vallia LAMB and keep track the moving average of each layer's scaling coefficient (used during the compression stage because the scaling coefficient is not stable at beginning). At the end of warmup, we stop updating the moving average, and store the frozen variance be used during compression stage. On the other hand, still keep updating another "fresh" " variance by reconstructing the global gradient based on this and last step's s compressed momentum.

---

Under review as a conference paper at ICLR 2022
To update LAMB scaling coefficient during compression stage, we compute 1-bit LAMB scaling ratio which is the max element among the (frozen varianceffresh variance). We use the max element as the scaling ratio because it is effective and cheap to compute based our experiments. To avoid extreme ratios and dramatic change between ratios, we use two kinds of clipping configurable by the user. (The clipping from vanilla LAMB algorithm are not used during the compression stage.) Then we compute this step's LAMB scaling coefficient using the 1-bit LAMB scaling ratio and the moving average at the end of warmup, and use it to update the model. Our study shows that those layers with less stable variance tend to have more dynamic 1-bit LAMB scaling ratios, which demonstrate the desired adaptiveness (details in Appendix A.5)) Appendix provides a theoretical analysis for 1-bit LAMB.
5 PROPOSED SYSTEM DESIGN
To realize the compressed communication system level, (Tagg et al. (2021) designed custom collective primitive, called "compressed allreduce", using Message Passing Interface (MPI). In addition implementing 1-bit LAMB using this MPI-based backend, we introduce new system implementation for compressed communication using the NCCL backend of PyTorch distributed, following the 3-phase design in the MPI-based backend: 1) The gather step, where each worker sends its i-th chunk to worker i, is implemented using NCCL's Alltoall and AllGather. 2) The average step, where each worker averages all chunks it receives. Here each worker acts as the server defined in Algorithm[] [ 3) The scatter step, where each worker receives the average of all i-th chunks from worker i, is implemented using NCCL's AllGather. Compared to the MPI-based, this new NCCL- based implementation significantly improves the usability since NCCL is integrated with Py PyTorch distributed. In addition, evaluations show that the performance of the NCCCL-bsee implementation is better than the MPI-based for Ethernet-based systems and on-par for InfiniBand Band-based systems. Thus we will mainly present the results with NCCL-bbase backend, but we also include a comparison between MPI and NCCL-based implementations in Appendix A..6
EVALUATION
Dataset and models We evaluate the convergence and performance of 1-bit LAMB and uncom-
pressed LAMB for BERT-Large (L = 24, H = 1024, A = 16, 340M params) treitraining task. We use the same dataset as Devlin et al. (2019). which is a concatenation of Wikipedia and CooksCorpus with 2.5B and 800M words respectively. Compared to the original BERT model, one notable change that we applied PreLN instead of PostLN for better training stability (Zhang Hel 2020 Xiong ct al. 2020). We use the GLUE fine-tuning benchmark (Wang al.
2018) and SQuAD 1.1 fine-tuning task to evaluate the convergence of the BERT models trained by LAMB and 1-bit LAMB.
Hardware We use the two clusters described in Section 3.2 We use 8 256 GPUs for BERT
pre-training tasks to measure 1-bit LAMB's s performance gain. For fine-tuning tasks we use 4 GPUs. One thing to note is that because 1-bit LAMB introduces additional memory overhead (one persistent copy of the fresh varaince and one temporary copy of last step's momentum) and because the V100 GPUs on the Ethernet cluster have 16GB memory instead of 32GB, we were not able to fit large batch sizes for BERT pre-training when the number of GPUs is small. For seqlen 128 and 512, we need to use at least and 16 GPUs on the Ethernet cluster, respectively.
Training parameters For BERT pre-training, we set the parameters in ①1 as 81 0.9, B2 0.999. Cmin 0.01 and Cmaz = 0.3 for LAMB and 1-bit LAMB. For 1-bit LAMB, we set 33 = 0.9. Tmin 0.5, Tmaz 4.0, and Tthreshold 0.1 in Algorithm [ For convergence analysis, we set total batch size as 64K for seqlen 128 and 32K for seqlen 512. For performance analysis, we test different batch sizes from 8K to 64K.
For BERT pre-training seqlen 128, the learning rate starts from 1 x 10 exponentially increases to 12 x 10 3 as a warmup in the first 450 steps, then decays into 0.9 of the original after every 250 steps. The total number of steps is 5993. For 1-bit LAMB we use the first 1000 steps (16.7%) as the warmup stage. For BERT pre-training seqlen 512, the learning rate starts from 0, exponentially increases to 2 x 10 as warmup in the first 150 steps, then decays into 0.9 of the original after
7

---

Under review as a conference paper at ICLR 2022
12
Training loss
JANB sealen 126
LAMB seglen 532
sedjon 128
3-t1 EAMB sndlen
100M
200M
300
4 400M
Samples
Figure 2: Sample-wise convergence speed for BERT-Large pre-training with LAMB and 1-bit LAMB.
Table 3: GLUE development set results using the pre-trained BERT-Large models. "Original" results are from Devlin et al. (2019) using BertAdam. "LAMB" results use the uncompressed LAMB for BERT praitrai "LAMB basic 1-bit" is the experimental algorithm described in Section 3.4
1-bit LAMB" is the proposed work. The latter 3 cases use the same shared training parameters during
pre-training and fine-tuning. Spearman correlations are reported for STS-B, and accuracy scores
reported for the other tasks. Each task's 's scores are the median scores over 32 runs.
MNLI-(m/mm)
QQP
QNLI
SST-2
CoL.A
STS-B
MRPC
RTE
Average
Original LAMB
86. 86.7/85.9
89.3
92.7
94.9
60.5
86.5
85.4
70.1
83.6
85.4/85.5
91.4
91.9
92.8
60.9
90.1
86.9
70.4
83.9
LAMB basic 1-bit
84.8/84.7
91.2
91.4
92.4
54.0
89.6
84.2
66.6
82.1
1-bit LAMB
85.5/85.6
91.3
92.3
93.1
59.2
90.0
86.5
71.5
83.9
every 150 steps. The total number of steps is 555. For 1-bit LAMB we use the first 107 steps (19.3%)
as the warmup stage.
For 1-bit LAMB's hyperparameter tuning, we find that the only parameter that requires nontrivial tuning is the number of warmup steps, which depends when do the LAMB's scaling coefficients become stable as discussed in Section 3.3 thus could potentially be automated. For the other parameters (83, T'ins I'maz and Tthreshold in Algorithm I find tuning them does not lead to very different convergence since they are just used to avoid extreme value/dramatic change of LAMB's B scaling coefficients. Similarly the clipping parameters Cynin and Craz in for LAMB and 1-bit LAMB do not require non-trivial tuning.
For GLUE benchmarks
Adam optimizer and perform single-task training on the dev set. Following the setup in the BERT paper (Devlin et al. 2019), we use a batch size of 32 and fine- tune for 3 epochs for all GLUE tasks. For each task, we select the best learning rate among {2 10-5, 3 x 10-5, x 10 10-$, 5 x 10-5], For SQuAD fine-tuning we use Adam optimizer and the dropout
same parameters as published by HuggingFace (batch size 24, learning rate = 10
0.1, 2 epochs).
Convergence analysis Figure 2 presents the BERT pre-training sample-win convergence results. For both seqlen 128 and 512. 1-bit LAMB provides the same convergence speed as LAMB, while takes much less time process each batch due to communication compression. We already presented the BERT pre-training validation loss and SQuAD fine-tuning results Section 3.4 Table And Table presents the GLUE results. For both SQuAD and GLUE, results show that 1-bit LAMB provides the same fine-tuning task accuracy as LAMB, while simply freezing both the varaince and LAMB's scaling coefficients would hurt the accuracy.
Performance analysis Computed 1/(warmup_rand ratio + (1 warmupp_atioo//// for FP16 training, 1-bit LAMB offers 4.6x and 4.1x end-to-end communication volume reduction for BERT-Large seqlen 128 and 512, respectively. To measure the actual end-to-end performance gain, first we perform throughput analysis where we run the warmup (i.c., bascline LAMB's performance) and compression stage of 1-bit LAMB for 200 steps each, and measure the average throughput of the two stages. Figure 3 presents the results with NCCL-based compressed communication backend under different batch sizes and number of GPUs. For seqlen 128, 1-bit LAMB provides up to 4.5x speedup during the compression stage, which is equivalent to 2.8x end-to-end speedup (computed as 1/(wamupp ratio + (1 warmup romioo/oomprsssonn _taggessttee__ provides up 3.0x speedup during the compression stage, which is equivalent to 2.2x end-to-end speedup. This demonstrates 1-bit LAMB's 's better scalability compared to LAMB. is also worth mentioning that 1-bit LAMB Ethernet (4.1 Ghps effective bandwidth, GPUs per node) is able
stage_speedup), speedup)). For seqlen 512, 1-bit LAMB
8

---

Under review as a conference paper at ICLR 2022
Table 4: Total runtime of BERT-Large pre-training with LAMB and 1-bit LAMB (256 GPUs with Ethernet connections. Batch sizes 64K/32K for seqlen 128/512.)
Seqlen 128
Seqlen 512
Total
LAMB
657 min
74 min
731 min
1-bit LAMB
301 min (2.2x) 50 min (1.5x)
351 min (2.1x)
12400
16020
30000
Sumplesseann
10900
14000
Samplecserand
25:00
Samplzn
10030
23000
Sumpleg
20000
8000
22000
15000
4300
4030
5000
129
16
12s 25s
15 22 65 135 256
16 22 126:256
Number of GPUs
Number GPUs
Number GPUs
Number GPUS
(a) seqlen 128, batch size 8K
(b) seglen 128, batch size 16K
(c) seqlen 128, batch size (d) seqlen 128, batch size 32K
64K
5000
LANE
6000
LANE InTEBaed
4000
2000
3000
2000
2000
1000
1000
16 32 64 128
16 32 54 129 256
16 32 64 128 256
Number
GPUS
(e) seqlen 512, batch size 8K
(f) seqlen 512, batch (g) seqlen 512, batch size 32K
16K
Figure 3: Scalability of 1-bit LAMB with NCCL-based backend for BERT-Large pre-training on V100 GPUs. LAMB lines represent the throughput at 1-bit LAMB's warmup stage (i.e., baseline LAMB). 1-bit LAMB lines represent the throughput at compression stage. Annotations represent highest speedup achieved each figure. Note that this is the speedup between warmup and compression stage. The end-to-end speedup also depends on the percentage of warmup. All figures share the same legend as 3(g)
to achieve comparable throughput as LAMB on InfiniBand (near 100 Gbps effective bandwidth, 8 GPUs per node), which demonstrates 1-bit LAMB's ' efficiency considering the hardware differences.
In addition to throughput analysis, also measure the total runtime of pre-training at batch size 64/32K for both 1-bit LAMB and LAMB. As shown in Table overall 1-bit LAMB is able to provide 2.2x and 1.5x speedup for seqlen 128 and 512. These numbers consistent with the end- to-end speedup calculated in the throughput analysis (2.0x and 1.5x based on results in Figure ((d) (g)). For seqlen 128, the end-to-end speedup based on runtime is slightly higher than the speedup based on throughput. We find that it is because uncompressed LAMB's s larger communication volume makes more sensitive the occasional fluctuation of the actual network bandwidth.
CONCLUSION
To reduce both the frequency and volume of communications for large-scale training, we propose an eroo--comennaaed LAMB preconditioned momentum SGD algorithm, 1-bit LAMB, which combines the power of large batch optimization and communication compression by introducing an novel way to support adaptive layerwise learning rates during communication compression. We also introduce an easier-to-use and more efficient compressed communication backend system based on NCCL. Evaluations show that 1-bit LAMB with NCCL-based backend is able to achieve up to 4.6x communication volume reduction and up to 2.8x end-to-end time-wise speedup for BERT pre-training compared to uncompressed LAMB, together with the same sample-wise convergence speed and fine-tuning task accuracy.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Naman Agarwal, Ananda Theertha Suresh, Felix Xinnan X Yu, Sanjiv Kumar, and Brendan McMa- han. cpSGD: Communication-efficini and tufferentially-privaal distributed SGD. In Bengio, and R Garnett (eds.), Advances in Neural
H Wallach, H Larochelle, K Grauman, N Information Processing Systems 31, pp. 7564-7575. Curran Associates, Inc., 2018.
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD: Communication Efficient SGD via gradient quantization and encoding. In I Guyon, U V Luxburg, S Bengio, H Wallach, R Fergus, S Vishwanathan, and Garnett (eds.); Advances in Neural Information Processing Systems 30, pp. 1709-17200 Curran Associates, Inc., 2017.
Debraj Basu, Deepesh Data, Can Karakus, and Suhas Diggavi. Osparse-local-spe Distributed sgd with quantization, sparsification and local computations. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 14695-14706. Curran Associates, Inc., 2019.
Jeremy Bernstein, Jiawei Zhao, Kamyar Azizzadeneshelii and Anima Anandkumar, signSGD with majority vote is communication efficient and fault tolerant. In International Conference Learning Representations, 2019. URL
net hwt.:/eeee//ernneeee Poforumiid-BBie JACY7
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Ben- jamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL httpt: htttl:::///pooceeddeee neuri ps .
Sorathan Chaturapruek, John C Duchi, and Christopher Ré. Asynchronous stochastic convex optimization: the noise is in the noise and sgd don care. In C. Cortes, N. D. Lawrence, D. D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Neural Information Processing Systems 28, pp. 1531-1539 Curran Associates, Inc., 2015.
Jacob Devlin, Ming-Was Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT, 2019.
Nikita Ivkin, Daniel Rothchild, Enayat Ullah, Vladimir braverman, Ion Stoica, and Raman Arora. Communication-efficiaa distributed sgd with sketching. In H. Wallach, H. Larochelle, A. Beygelz- imer, F. d Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Curran Associates, Inc., 2019.
Systems 32, pp. 13144-13154
Peng Jiang and Gagan Agrawal. A linear speedup analysis of distributed deep learning with sparse and quantized communication. In S. Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi and R. Gareett (eds.), Advances in Neural Information Processing Systems 31, pp. 2541.. Curran Associates, Inc., 2018.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/1412.6990, 2015.
Anastasia Koloskova Tao Lin Sebastian U Stich, and Martin Jaggi. Decentralized deep learning with arbitrary communication compression. In International Conference Learning Representa- tions, 2020. URL
net www.eeee////uuneeeeeeeemee Lteee//oooouuumdsssscccckCr
Youjie Li, Mingchao Yu, Songze Li, Salman Avestimehr, Nam Sung Kim, and Alexander Schwing. Pipe-sgd: A decentralized pipelined sgd framework for distributed deep net training. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Besa-Bianchi, and R. Garnett (eds.), Advances in Curran Associates, Inc., 2018.
Neural Information Processing Systems 31, pp. 8056-8067
10

---

Under review as a conference paper at ICLR 2022
Xiangru Lian, Ce Zhang, Huan Zhang, Cho-Jui Hsieh, Wei Zhang, and Ji Liu. Can decentralized algorithms outperform centralized algorithms? a case study for decentralized parallel stochastic gradient descent. In Guyon, U V Luxburg, S Bengio, H Wallach, R Fergus, S Vishwanathan, and R Garnett (eds.), Advances in Neural Information Processing Systems 30, pp. 5330-5340. Curran Associates, Inc., 2017.
T. T. Phuong and L. T. Phong. Distributed sgd with flexible gradient compression. IEEE Access, 8: 64707-64717, 2020.
Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient descent and application to data-parallel distributed training of speech dnns. In Interspeech 2014, September 2014.
Zebang Shen, Aryan Mokhtari, Tengfei Zhou, Peilin Zhao, and Hui Qian. Towards more efficient stochastic decentralized learning: Faster convergence and sparse communication. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Leurning, volume 80 of Proceedings of Machine Learning Research, pp. 4624 4633, Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018. PMLR.
S. Shi, Q. Wang, K. Zhao, Z. Tang, Y. Wang, X. Huang, and X. Chu. A distributed synchronous sgd algorithm with global top-k sparsification for low bandwidth networks. In 2019 JEEE 39th International Conference on Distributed Computing Systems (ICDCS), pp. 2238-2247, 2019.
Shaohuai Shi, Xianhao Zhou, Shutao Song, Xingyao Wang. Zilin Zhu, Xue Huang, Xinan Jiang, Feihu Zhou, Zhenyu Guo. Liqiang Xie, Rui Lan, Xianbin Ouyang, Yan Zhang. Jieqian Wei, Jing Gong, Weiliang Lin, Ping Gao, Peng Meng. Xiaomin Xu, Chenyang Guo, Bo Yang, Zhibo Chen, Yongjian Wu, and Xiaowen Chu. Towards scalable distributed training of deep learning on public cloud clusters. In Proceedings of Machine Learning and Systems, 2021.
Ryan Spring, Anastasios Kyrillidis, Vijai Mohan, and Anshumali Shrivastava. Compressing gradient optimizers via Count-Sketches. Proceedings of the 36th International Conference on Machine Learning, 97:5946 5955, 2019.
Sebastian U Stich, Jean- Baptiste Cordonnier, and Martin Jaggi. Sparsified sgd with memory. In S. Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 4447-4458 Curran Associates, Inc., 2018.
Jun Sun, Tianyi Chen, Georgios Giannakis, and Zaiyue Yang. Communication--fficienn distributed learning via lazily aggregated quantized gradients. In H. Wallach, H. Larochelle, A. Beygelzimer, F. Alché-Buus Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 3370-3380. Curran Associates, Inc., 2019.
Hanlin Tang, Chen Yu, Xiangru Lian, Tong Zhang, and Ji Liu. Double Parallel stochastic gradient descent with double-pass compr-comeenstted compression. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 6155-6165, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
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 Lcrge-Scale Training with Adam's Convergence Speed. In Proceedings of the 38th International Conference on Machine Learning, pp. 10118-10129, 2021.
Thijs Vogels, Sai Praneeth Karimireddy, and Martin Jaggi. Powersgd: Practical low-rank gradient compression for distributed optimization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d Alché- Buc. E. Fox, and R. Garnett (eds.). Advances in Neural Information Processing Systems 32, pp. 14259-14268. Curran Associates, Inc., 2019.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi:
1..18663/v1/118-555 httpst www. ac web. .oceboorg wwwpgg//ahllllloooooooooooooo ogy / 5444

---

Under review as a conference paper at ICLR 2022
Jiangiao Wangni, Jialei Wang. Ji Liu, and Tong Zhang. Gradient sparsification for Communication- Efficient distributed optimization. In S Bengio, H Wallach, H Larochelle, K Grauman, N Cesa- Bianchi, and R Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 1299- 1309. Curran Associates, Inc., 2018
Wei Wen, Cong Xu, Feng Yan, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Terngrad: Ternary gradients reduce communication in distributed deep learning. I. Guyon, U. V. Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems 30, pp. 1509-1519. Curran Associates, Inc., 2017.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 10524-10533 PMLR, 13-18 Jul 2020. Ut.http::/pprooo
ngs. mml...mlr.. prwss/ 9/ wwt.vl11//////////1 2b.
Min Ye and Emmanuel Abbe. Communication Conoonnttotmttoiinn efficient gradient coding. Proceedings of the 35th International Conference on Machine Learning, 80:5610-5619, 2018.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations. 2020. URL wttp///pee.e.....ooo//p/ http:/iwwwww.eeeeeeeeeeeeeeeeeee http://ww....t.ttaaaaaaaaaaaaaaa
Yue Yu, Jiaxiang Wu, and Longbo Huang. Double quantization for communication-efiiiennn dis- tributed optimization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 4438 4449. Curran Associates, Inc., 2019.
Hantian Zhang, Jerry Li, Kaan Kara, Dan Alistarh, Ji Liu, and Ce Zhang. ZipML: Training linear models with end-to-end low precision, and a little bit of deep learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 4035-4043, International Convention Centre, Sydney, Australia, 06-11 11 Aug 2017. PMLR.
Minjia Zhang and Yuxiong He. Accelerating training of mansformer-based language models with progressive layer dropping. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 14011- Cur- ran Associates, Inc., 2020. URL / neurips........p.pp............. cc CC wttp:pps.ccppppp///0/////
Shuai Zheng, Ziyue Huang, and James Kwok. Communication-effiienn distributed blockwise momentum sgd with error-feedback. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d Alché- Buc. E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 11450-11460. Curran Associates, Inc., 2019.
Shuxin Zheng, Qi Meng, Taifeng Wang, Wei Chen, Nenghai Yu, Zhiming Ma, and Tie-Yan Yan Liu. Asynchronous stochastic gradient descent with delay compensation for distributed deep learning. CoRR, abs/1609.08330 2016
12

---

Under review as a conference paper at ICLR 2022
Table 5: BERT-Large pre-training seqlen 128 profiling results.
Cluster
Num. node
Num. GPU
Batch
Batch size
Grad
Forward
Backward allreduce
Backward everything else (ms)
Step (ms)
allreduce
Network
size
accum.
(ms)
Type
GPU
step
(ms)
Ethernet
256
16
8K
55
3579
117
191
91%
Ethernet
64
256
16
16K
111
3533
227
195
87%
Ethernet
64
256
16
32K
8
224
3599
462
233
80%
Ethernet
64
256
16
64K
16
445
3674
919
215
70%
Ethernet
128
16
8K
112
3759
233
121
89%
Ethernet
16
64
16
8K
223
3433
464
109
81%
Ethernet
8
32
16
8K
16
445
3528
923
38
72%
Ethernet
4
16
16
8K
32
881
3436
1827
33
56%
Ethernet
8
16
8K
1773
2087
3696
31
28%
Ethernet
16
8K
128
3532
234
7329
30
InfiniBand
16
128
64
8K
96
335
179
36
52%
InfiniBand
16
128
64
16K
2
192
346
356
37
37%
InfiniBand
16
128
64
32K
381
377
714
37
25%
InfiniBand
16
64
422
1422
32
16%
InfiniBand
64
64
SK
2
192
332
352
34
36%
IndiniBand
4
32
64
8K
4
384
339
711
31
23%
InfiniBand
16
64
8K
768
270
1436
31
11%
InfiniBand
64
8K
1534
167
2869
31
4%
A APPENDIX
A.1 TRAINING PARAMETERS FOR 1-BIT ADAM EXPERIMENTS IN SECTION 3.1
For experiments in Section 3.1 for both LAMB and 1-bit Adam use batch size 16K, 28125/3125
steps for seqlen 128/512, weight decay = 0.01, linear LR warmup and decay. For LAMB, we use learning rate 3.54 10 10% LR warmup, clipping configs (Cmin and Cmaz in ( 0.1 and 1. For 1-bit Adam, we use learning rates E (1 x 10-1 1 x 10-4,3 x 10-*}, LR warmup (5%, 10%, 20%). ). All of these training parameters (except LAMB clipping configs)
from the LAMB paper. For 1-bit Adam, following the original work's 's strategy we set the number of warmup steps as 4000 (out of total 28 1252 steps) for seqlen 128 and 475 (out of 3125) for seqlen 512.
A.2 DETAILED RESULTS FOR PROFILING EXPERIMENTS SECTION 3.2
Table 5 presents the detailed profiling results. Results show that even with larger batch sizes the allreduce communication still contributes to a great portion of the training time per step, up to 91 % and 52% on two different kinds of clusters. And this overhead is larger when the number of nodes is larger, when the batch size smaller, when the network bandwidth is lower.
A.3 DETAILED RESULTS FOR LAMB EXPERIMENTS IN SECTION
Figure 4 presents LAMB's 's scaling coefficients for different layers during BERT pre-training sequence length 128 (sequence length 512 has similar patterns). Only the c s see_reelleeeeeeeeeeeeeeeeeeeeeee www.relatiii ielshinsship bi has very unstable scaling coefficient. This is because this bias only has elements, representing the two states of whether the two sequences are next to each other.
Figure presents LAMB's variance norms for different layers during BERT pre-training sequence length 128 (sequence length 512 has similar patterns). We believe that there are two reasons why LAMB's 's varaince terms are less stable: LAMB has larger batch size, smaller number of steps, and layerwise adaptive learning rates compared to Adam. 2) Because LAMB requires the calculation of the scaling coefficient for each layer, we cannot fuse all the variance together as Adam. And each separate variance could have less stable norm compared to single fused variance.
A.4 DESIGN OF THE EXPERIMENTAL. ALGORITHM IN SECTION 3.4
Because LAMB's ' scaling coefficient is essentially part of the learning rate (and it's just a scalar), up- dating this coefficient would not affect communication compression' 's error compensation mechanism
13

---

Under review as a conference paper at ICLR 2022
Werdentedding hovlenentertepp
Lagerieem-e Legttlermb
Preetin Layefisena Layefterms
ORY
Posvetin Layeriem Layellrm
0/4
0.4
Asn vocst-w
LAMB scaling
5:2
0.2
0.2
26
45
1K
48.
Steps
2xpps
Steps
(a) Berddmbedsings
(b) First BertLayer, part
(c) First BertLayer, part
dente, arta dants arta
LayeTeem
Mac
Loyothome.
0.2
0.1
0.0g 1.C
0%
36
Staps
Staps
(d) BendLMPradiction
(e) cls.seq relationship
Figure 4: LAMB's 's scaling coefficients (cf²) in (T) for different layers during BERT-Large pre-training
seqlen 128 (5993 steps in total). Since all 24 Bertl Layer have similar patterns, we just present the first. one. We set the lower/upper bound of the scaling coefficient (Cmin and Gmaz in (T) at 0.01 and 0.3.
LAWBS vetience surm
varance
10-4
10-1
10-1
20
30
22-1
Fasidon embocings
Tocan 1yos bgd
Lepotteriz
10-0
10-1
Preati Uneetern
Alle
20-1
10-4
Fadl altn Lyeeteess
napurw
16
S2o00
Stops
Steps
(a) BentEmbeddings
(b) First Berth ayer, part
(c) First Bertl BertLayer, Layer, part
10-
Leno
dente 3Ct
Layodiom
Livellam
(d) Mrrdiittrcdit
(c) cls.seq relatiosii
Figure 5: LAMB's 's variance norms for different layers during BERT-Large pre-training seqlen 128. Since all 24 BertLayer have similar patterns, we just present the first one. The y-axis is in log scale.
as long as the change small enough between two steps However, we find that it's challenging to update this coefficient during compression stage because in LAMB's algorithm, updating this scaling coefficient requires both momentum and variance term. However, the error compensation mechanism requires freezing the variance term the beginning of the compression stage due to its nonlinear dependency to the gradient. In addition, we find that due to the error compensation mechanism, the norm of some layers' momentum term could become larger/smaller compared to the uncompressed case. As a result, we find that updating LAMB's s scaling coefficients during compression stage based
In fact described in Section 4 the proposed 1-bit LAMB algorithm does require updating LAMB's scaling coefficient during compression stage, but in way different from original LAMB algorithm.
14

---

Under review as a conference paper at ICLR 2022
Word errbeedings
Pre Pr-am Layerfome Layerformb
oK/b
fot Layeriarma lag LayerMornaz
Intenedateb
loatton umbedd tye eneotiine
Am supuat
Gupuea
1-00t LAMS scalling
1on LiB
36
36
5K
15
2%
35
06
16
25
56.
6K
Steps
Steps
BortEmbeddings
(b) First BertLayer, part
(c) First BertLayer, part
Leyenlemm
wa jelatsnntss 6n miMiinces
1-401 LAMB scalng ratas
tener arth
Nas
Lapa
1-08 LANB
16
30
3%
56
68
36
45
3k 60
Stepss
(d) Berd
(c) dis.segrrelationsig
Figure 6: 1-bit LAMB scaling ratios (r!") in Algorithm 1 for different layers during BERT pre- training sequence length 128. Since all 24 Bertl have similar patterns, we just present the first one. The number of warmup steps 1K out of 5993 total steps. We set the clipping configurations of the scaling ratio (Tmin, Tmozi l'threhhodd in Algorithm D at 0.5, 4.0, 0.1.
on original LAMB algorithm would produce suboptimal results and slower convergence speed. Thus for Section 3.4 experimental algorithm, stop updating LAMB's 's scaling coefficients during compression stage, and just use a calculated scaling coefficient moving average at the end of the warmup stage. For compression stage also tried communicating based on 1-bit compressed gradient, but leads to much slower convergence speed. We believe it's because gradients are less stable, which could lead higher compression error and slower convergence.
A.S 1-BIT LAMB ALGORITHM DESIGN AND IMPLEMENTATION DETAILS
Figure presents 1-bit LAMB scaling ratios (r2 in Algorithm 1 for different layers during BERT pre-training sequence length 128 (sequence length 512 has similar patterns). When comparing with Figure 5 we find that for those layers with less stable varaince (e.g.. 3 kinds of embeddings, weights in BertLayer), the corresponding 1-bit LAMB scaling ratios are also larger. As result 1-bit LAMB able to adaptively update LAMB scaling coefficients during compression according to the difference between the frozen and fresh variance.
A.5.1 REDUCING NUMBER OF COMMUNICATIONS BY MOMENTUM FUSION
Different from Adam, LAMB has the scaling coefficients that need to be updated separately for each layer. For the communication of the compressed momentum during 1-bit LAMB's compression stage, if we also communicate separately for each layer, the number of communications (which is 302 for BERT-Large model) will greatly affect the overall performance. Thus fuse all layers' momentum into a contiguous 1D buffer and just do one communication over this fused momen- tum. This fused momentum not duplication, but different "view" of all layers' momentum. We implement this momentum fusion by Norchh_utile ftt.aatllllllllleeeeeeeeeeeeee tensors and torch. warchh_uuille uuili.................... unf wweeteeendeeeeeeeeeeeeeee tensors..
A.5.2 REDUCING COMPRESSION ERROR BY MOMENTUM SCALING
We find that after momentum fusion, the compression emoor increase lot for some layers which greatly increase the chance of divergence. This is because: 1) When communicating base on the fused momentum, all layers' momentum will be compressed the same scale due to the nature of
15

---

Under review as a conference paper at ICLR 2022
LANB
LANS
LABLLAS (NOCJ
1-111 LAMO (SCCL)
Samplestter
20000
15000
15300
10000
10000
5000
5000
1K
32%
64K
D6
16k
32K
BAK
Botch
Batch stae
(a) 256 GPUs on Ethernet cluster
(b) 128 GPUs on InfiniBand cluster
Figure 7: Comparing MPI and NCCL-based compressed communication backend implementations based on performance of BERT pre-training seqlen 128.
1-bit compression (representing the tensor by =1 signs and a single scale). Thus the layers with very small/large momentum scale will have larger compression error; 2) Due to LAMB's 's layerwise adaptive learning rate, each layer is learning different speed, which could further increase the variance of momentum scale among layers. We solve this issue by computing an average momentum scale at the end of warmup stage and using it to compute momentum scale coefficient for each layer. During the compression stage, we multiply each layer's local momentum by its scale coefficient, then do the compressed communication, then divide each layer's global momentum by the same scale coefficient. By performing this momentum scaling, all layers' momentum will have similar scales when passed to the compressed communication. Thus we are able to greatly reduce the compression error and chance of divergence. Since the momentum scale coefficients are fixed during the compression stage, won't 't affect 1-bit compression's 's error compensation mechanism.
A.6 COMPARING MPI AND NCCL- COMMUNICATION BACKEND
All of the results in main paper are evaluated with the NCCL-based compressed communication backend implementation. Figure 7 presents the performance comparison between MPI and NCCL- based implementations. Compared to the MPI-based implementation introduced in the 1-bit Adam work, our NCCL-based implementation is able to provide better performance on the Ethernet cluster (where OpenMPI library is used for MPI backend) and on par performance on the InfiniBand cluster (where MVAPICH2-GDR is used for MPI backend).
A.7 ADDITIONAL CONVERGENCE ANALYSIS ABOUT NUMBER OF WORKERS AND BATCH SIZE
This section provides additional convergence analysis on BERT pre-training and SQuAD fine-tuning. First we investigate whether the number of workers/GPUs affect the convergence 1-bit LAMB. Table 6 presents the SQuAD average/max F1 scores using the BERT models pre-trained with LAMB and 1-bit LAMB optimizers, with different number of GPUs, and with the same training parameters described in S Results show that both LAMB and 1-bit LAMB achieve similar SQuAD Fl scores under different number of workers, and 1-bit LAMB always provides better results.
Next, we investigate whether the batch size affects the convergence of 1-bit LAMB. In main paper we compared LAMB and 1-bit LAMB when pre-training BERT with batch size 64K/32K at seqlen 128/512. Here we compare them batch size 16K and 4K. At batch size 16K, for LAMB we use the same training parameters described in Appendix A. and for 1-bit LAMB we set the number of warmup steps as 4000 (14%) for seqlen 128 and 475 (15%) for seqlen 512. At batch size 4K, for LAMB we follow the original work and use 1 /12000/
steps for seglen 128/512, weight decay 0.01, linear LR warmup and decay, learning rate = 1.77 x 10 2.5% LR warmup, clipping configs (Cmin and Cmax in (I) as 0.1 and 1. For 1-bit LAMB we set the number of warmup steps as 16000 (14%) for seqlen 128 and 1900 (15%) for seqlen 512. For 1-bit LAMB's 's other hyperparameters in
Algorithm[ [ we keep them the same under all batch sizes: 33 0.9, Tiin 0.5, I'az = 4.0, and
Tekreshold 0.1. As shown in Table 7 results show that 1-bit LAMB provides better SQuAD FI
scores than LAMB under all batch sizes.
16

---

Under review as a conference paper at ICLR 2022
Table 6: SQuAD) average/max dev set Fl scores over 32 using the BERT models pre- trained (batch size 64K/32K at seqlen 128/512) with different optimizers and different number of GPUs/workers. The first column from the original LAMB work. The last two columns are our experiments using the same training parameters.
SQuAD Avg./Max Fl
LAMB (You 2020
LAMB
bit LAMB
128 GPUs
-/90,584
90.257/90.557 90.265/90.555
90.494/90. 740
64 GPUs (from Table
590,584
90,524/90. 7..84/990.88 90,469/90.811
32 GPUs
-/90.584
90.263/90.58%
Table 7: SQuAD average/max dev set Fl scores over 32 runs using the BERT models pre-trained with different optimizers and different pre-training batch sizes. The first column is from the original
LAMB work. The last two columns are our experiments using the same training parameters.
SQuAD Avg.MMax F1
LAMB (You et al. 2020)
LAMB 90.265/90.555
1-bit LAMB
Bsz 64K/32K for seqlen 128/512
190.584
90.524/90.788
(from 2
Bsz 16K/16K for seqlen 128/512
-191.345
90.716/91.11 90.917/91.086
90.906/91.14
Bsz 4K/4K for seqlen 128/512
1/91.137
90,942/91.111
A.8 ADDITIONAL CONVERGENCE ANALYSIS: RESNET-50 ON CIFARI00
This section provides additional analysis of training ResNet-50 CIFARI00. Because ResNet-50 much smaller than BERT-Large and because we use GPUs on a single node for the training, 1-bit LAMB and LAMB have very similar throughput. Thus this experiment focus whether 1-bit LAMB can achieve the same convergence as LAMB under the same training hyperparameters. For both LAMB and 1-bit LAMB, we use 250 epochs (about 12K steps), batch size 1K, learning rates E {0.005, 0.01, 0.02), ). linear LR warmup in the first K steps, no LR decay, weight decay = 0.01. For LAMB set the clipping configs (Cmin and Cmaa in (T) 0.01 and 0.3. For 1-bit LAMB we set the number of warmup steps as 1K (8%), and set other hy hyperparameters in Algorithm aas: B3 = 0.9, Tmin 0.5, Tmaz 2.0. and Fthreshold 0.1. Figure presents the sample-wise training loss and testing accuracy, and Table
S presents the final testing accuracy. Results show that under all three learning rates, 1-bit LAMB provides on par or better convergence and final accuracy than LAMB.
A.9 THEORETICAL ANALYSIS
Notice that for 1-bit LAMB, only use original LAMB at warm-up, and then essentially run error-compensated momentum SGD with coordinate-dependeet learning rate v"Tw" Therefore here
we consider the LAMB-based warm-up phase as a way to find a good precondition variance term Vn,, be used in the compression phase. Below we are going to introduce the convergence rate for the compression phase after warm-up. up. We first introduce some necessary assumptions, then we present the theoretical guarantee of the convergence rate for 1-bit LAMB.
Assumption 1. We make the following assumptions:
1. Lipschitzian gradient: f(-) assumed to be with L-Lipschitzian gradients, which means
[Vv - vf(y)) - L| - -vll. Vx, vy,
2. Bounded variance: The variance of the stochastic gradient is bounded
(VF c( [ < o², Va, Vi.
3. Bounded magnitude of error for Cu-l-]: The magnitude of worker's local errors 8 and the server's global error 8, are assumed to be bounded by constant
n E.
n
E 18
5
Vi
The most important reason we use the third assumption above is that for 1-bit compression. the signal-to-noise assumption does not hold if there is only non-zero element in the tensor. In (Tang
17

---

Under review as a conference paper at ICLR 2022
0.800
109
0.775
LAMB
LAMB LR-O A1
1-H LAMB LRt0.01
1-511 LAMII
Training loss
10-2
0.750
LAMB LRwOb.02
Eit LAMB LR-0000
0.725
0.700
10"6.
Testing
0.675
LAMB LR-0.005
1-bit LAMB LR=0.005
10-4
LAMB LR-0.01
1-bit LAMB 18-0..0
LI--0.02
0.600
3M
6M
9M
12M
3M
6M
9M
12M
Samples
Samples
(a) Training loss
(b) Testing accuracy
Figure 8: Sample- wise training loss and testing accuracy for ResNet-50 on CIFARION
Table 8: Final testing accuracy for ResNet-50 on CIFAR
Final testing accuracy
LAMB
1-bit LAMB
LR=0.005
0.68
0.68
LR=0.01
0.68
0,73
IR=0.02
0.66
0.69
et al. 2019), authors proved that our assumption is equivalent to the signal-to-noiss assumption when the gradient magnitude is upper bounded.
Next we present the main theorem for 1-bit LAMB.
Theorem 1. Under Assumption 7 for 1-bit LAMB, we have the following convergence rate
8-.3 Tmmxx I.2
T
Ymax Vmax
Σ B())(1
(1 3)
2E (x))( - 2Ef(y/)
+ 24 / L
LV Vmax o
TT & L2V2 O
max
(3)
Ymin
Wimin
n(1 B)2 Ymin
where V = diag (1/0f/1 1/07 ... 1/// (d)
(1)
(2)
is diagonal matrix spanned by Ur
Here 1/v in is a moving sum over all history (V (V)()))² not || [a)) [|- Therefore it will not approach to O for the following two reasons: We stop updating v after the warmup stage, which means V f( at this time will not approach 0 since it has not converged yet; ii) Since B2 = 0.999, this means is heavily dependent the history gradients, therefore even though V f=(x) approaches 0, v itself will still taking time to decay to 0. Therefore it is reasonable for us to assume an upper bound for Vinax
Given the generic result in Theorem D we obtain the convergence rate for 1-bit LAMB with appropri- ately chosen learning rate "y-
Corollary 1. Under Assumption 7 for 1-bit LAMB, choosing ?
7max
7min
16LV mas + we have the following convergence rate
1-9
V
T-1
max
S
VnT
T3
T
where we treat f(x1) - f-. 3 and as constants.
This result suggests that: 1-bit LAMB essentially admits the same convergence rate distributed SGD in the sense that both of them admit the asymptotical convergence rate O(1/vnT), which means we can still achieve linear speedup w.r.t. the number of workers n.
For theoretical analysis, follow the framework of previous work (Tang et al. 2021). Below first summarize the updating rule of 1-bit LAMB, followed by the proof of Theorem then the proof of Corollary [
18

---

Under review as a conference paper at ICLR 2022
A.9.1 PROOF TO THE UPDATING FORM
Since our algorithm is equivalent running parameter-server prototype communication on each chunk of the gradient, so below we will assume parameter-server model (which means the tensor not required be divided into n chunks) for simplicity. Although we use a layer-wise updating strategy in 1-bit LAMB, for the simplicity of theoretical amalysis, we consider the case where we use the same updating rule for all the layer, but our results can be easily applied to this layer-wise case.
According to the algorithm description in Algorithm at iteration + 1, the updating step of the momentum term M4+1 can be divided into two steps:
1. Local Update and Compress: each worker locally update my and use the error-ompensate strategy for compressing.
m, =Bm; + (1 - -)) (8)
(e) -Cu
(2) Ct- (1) mi
1+ 7
Ct-1
+ ((0)
1+1
m,
m b43
2. All workers send its m (1) compress it again using corpensompen...onn
to the server. The server takes the average over them and
e+
Σ
(0) +
m
m
n
1=1
Ct-2
m141 =Cw
m
6
841 + + 8
m
3. The server broadcast m,+1 all workers, and all workers update the local model according to
T++1 2t 7/--1 m++1
Notice that here we do not include the updating rule for {ct} because for the following part will only view it as scaling factor.
So actually the updating rule above can summarized as
2-2 6,
m =Cr,
m 1 +
4-1
-m++ + C-2 8, 8Mt1 (from the definition of 8141)
6-1
Cw
mi
(i)
C-2
((0)
+
Ct2
6 8
C-1
C--1
1=1
m
4-2
0,
0,41 -8(1) (i)
+
C-2
6₁ 8.+1
(from the definition of 6!2)
n
C-1
Ct-1
i=1
3
2-
(1) 0 +
=8m+
+
9t
+
+ +
8c+1
C-1
71 i=1
17
i=1
i=1
Denote
n
(1)
9: N1
gt
8
+ 8
1=1
19

---

Under review as a conference paper at ICLR 2022
the update rule of m, can be summarized as
= 3mm-1 + (1 (1-3) +
2t-2 8,-1 8,.
C-1
and
21+1 = 21 - -7Va--mm
where V = diag(1/ 1/2,, 1//vea)
the diagonal matrix that spanned with vr... In order to simlify the notation, we define >t := 7 as time-varying learning rate, which reduces the updating rule above into
m. 55m (1 5)g. Tt-1 8, 8,,
7t
1t+1 -c mm.
A.9.2 PROOF TO THEOREM I
Notice that in for 1 1-bit LAMB, the learning rate for each coordinate is different. In order to simplify our analysis, we instead consider another function that is defined as
H(z) F(V
also
= = f(Vt2),
where V a diagonal matrix. In this case have
v Vh(z).
Therefore the updating rule of 1-bit LAMB in the view of h(-) is
V z Vizze 7V3 (våm:)
It can be easily verified that
m: (11 -
81-8 9: +
Σ
g (6-=(00_1 -d,)
sm0
sm0
=(1 3)
g
1
FF(V zz (1)) Σ
(0. (0.1) - -8,)
T1
==0
i=1
==0
which means
n
V mim; =(1 -
g1-s
V V +v)))
- -3,)
g 8 (6,1
n1
0=0
5m0
n
=(1 - ()
g
VH(V 2() Σ gr-o (8.-111
0=0
i=1
2=0
-(1
3)
+
g (6-1-1 3z),
==0
==0
where (z) is the corresponding averaged stochastic gradient computed in the view of loss function h(-). Then, if we define mz(z) V m,, the updating rule of (2) admits
m+(2) 8m,-1 (z) (1 (z) + V - V.3.,
20

---

Under review as a conference paper at ICLR 2022
and
V 2141 =V - -VV2mm 2++1 =2+ (2).
(5)
From (4) and ( we shall see that using different learning rate for each coordinate is equivalent to optimizing a new loss function defined on scaling the original coordinate and using a uniform learning for all coordinates. Therefore below we first study the behavior of the eror-compensated momentum SGD using coordinate- independent time-varying learning rate.
coordinate-ineeeedeen
Below are some critical lemmas for the proof of Theorem I
Lemma 1. Given two non-negative sequences {at and {be that satisfying
bbs
O2
with E [0, 1). we have
1 D&
p)
Proof. From the definition, we have
k
kーs
b
S₁
b
b
S
Σ
p
Dg
p br
b.b,
12 b2
(due to 7
which completes the proof.
Lemma 2. Under Assumption for any sequence that follows the updating rule of
Øt+1 =0, mi
m: =3 + (1 9)g, +
7 5,
5,,
7
if
Eg. Vf(x.). Blg, vf(2))11 2²
Vt.
[V) v < L| vll. Vz,Yy.
21

---

Under review as a conference paper at ICLR 2022
then guarantee that
1- (1-7L- (1 B32
22
T
B f(,))
(=0
2Es(211) - -2E
Loo 2721,-02
7
(1 B)²
n(1 - B)²
Proof. Instead of investigating æ, directly, we introduce the following sequence
Y1 -1 Y: + 7/110
1- B
The updating rule of y, admits
1 1- (m++1 - -m²) + 7:0/41 - 94-10=
Ytt
Yo+1 Yt 1x++1
2t-
1 1
78 - - 6
(8m= + - 1-)g+++1 .ct1 1- 8.) +
=m
3
7
1 B
--
Since f(-) is with L-Lipschitzian, we have
-
E (Vf())) 3e+1 -vi) + EE 1 yill
L7² E gl
E. (V.( 9.+1)
Li
BBE (V)))) v[(
+ E 11
2
2 [
Ttt E (x
(2()))1 + L2 EE g 11
s )x
Ly
BBl
[r+1
1:112
L
111
2
" "
T.2
(2
3)m,
T- 1,-1
Int Elgg
1
L2
47 L
Li E6 g++1
V ((x 1]
EE m:
10
11,1112
(Tthreshold <
(1-
3)2
(1 - 3)2
L2
L
S " E V f111+1)11
E m,
gl1g++111
(1-
(1 3)2
4-, 1,2
Inf
Lnfo² 2n
s 7 E V (1++))11
EE m,
E
V
(x(+1)))
(1 - 3)²
(1 3)2
Summing up the equation above from t to = T we get
(1
7:L
T
4 L2
4I2 DT
7
Lo² Σt 72
E/(yr=1) - Ef((n) S
m ||
+
2
3)
(1
3)2
2n
t=0
Therefore if we 74 <
and it's within certain range % E Mmin 7.ma]. the equation above leads to
2r.
T
(1 - (1-Tmax L)
E
{(
1=0
2Ef(Uo) - 2E/((T+1)
8 I
8 1.22T I.22T
L I 2T
E mm 1 +
(8)
Vwin
(1 2min
(1 2min
7Ymia
t=0
22

---

Under review as a conference paper at ICLR 2022
Notice that we have
my -(1 - (3
8t
9. +
8 - 5,))
==0
==0
which by using Lemma
we have
2
T²I + B)2"
2
m 1
|
| 1
[[((x))
(9)
3322
t=
1=0
t=0
Combing
and
together we get
8
I
1 Vmax L (1
Imax
B (x)))
1=0
21E (yo) 22E (ET(()++1)
247
2
L o
8 L2202T
Ymmi
(1
m/mmi
n(1 3) ymin
21E ( 2Ef(y*)
24 L
8 L²
Ymii
(1 3)2
77min
n11 - 3)
Proof to Theorem Since using per-coordinate learning rate for loss function f(-) is equivalent to use constant learning for all coordinates but for loss function fa(),, the only two thing that change are
Different L-Lipschitzian coefficient: the L-Lipschitzian coefficient for h(-)) is
[[Vh(=) V
v v (V V
VV(V x)) vs(v _
<L² |vtx v v+yll ||
=L y
< < 12x
| y 1
Therefore the effective L-Lipschitzian coefficient h(x) is Vmax
Different delinition of d,; from (4) we shall see that actually the compression error in the view of h(-) is V so in this case we have
BllV + 8el1 S Vmaxx2
Proof. From Lemma 2 we have
1 7max LVmax
8 122
T
E0
fxx))1
(1- 2
2E 2Ef(2z) 2E
3 mmm V3
V 2T
LV 2 i 2T
8 L2V2 max 7min
TT
Timin
(1 7min
77min
n11
A.9.3 PROOF TO COROLLARY
Proof. By choosing Tmax = 7min
1-8
we can guarantee that
16L VoL / -T+
V2 max
1 1
(1- 19)2
23

---

Under review as a conference paper at ICLR 2022
So 3
leads to
16LV max + +ov - + I {
T
(1-13
(11 (11-3)IV3 -BBIVTTT + +
+32L c T V3
max
vn
max
ElVf(> (Bf((o) - -f(D"))
16LV max
$
+ T
(1-3)
T
VnT
3)- (1-+))+
2
32L22
e 3 v3
VT
VnT
max
Treating f(y1) - f 3 and L as constants, from the inequality above we get
T
E (f(x))
T 10
VIT
T T
completes the proof.
24


Reviews:
[{"review_id": "zsaoiuQcysD", "reviewer": "Reviewer_5SMi", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strengths:\n* In general, the results demonstrated in the work show measurable improvements in efficiency over the baselines: the gains of 2-3x might be quite significant for practical scenarios of large-scale learning. Authors also compare their method with several reasonably chosen baselines to demonstrate the importance of their contributions: while there definitely could be more pretraining tasks and models to be considered, given the space limits I believe the current scope of experiments is sufficient to justify the claims made in the paper.\n* Moreover, the theoretical analysis of the algorithm allows to validate its scalability properties and to establish convergence rates mostly equivalent to those seen in previous works on distributed optimization.\n\nWeaknesses:\n* In the text, I could not find any information about the actual compression algorithm used for experiments. I think this is a quite important part of the algorithm which should be described in detail, since both the empirical results and the analysis (Assumption 1.3) rely on it.\n* Although the simplicity of the method is definitely a plus, some parts of the contribution of this work appear to lack any significant novelty. For instance, it was not clear to me how the primitive described in Section 5 differs from regular All-Reduce and whether this contribution is solely a software engineering one (that is, porting the existing protocol from MPI to NCCL). The method also seems to generally follow from 1-bit Adam, with the exception of a gradient recovery mechanism (although authors indeed demonstrate that direct application of the techniques from prior work yields worse results).\n* The presentation of the method can be slightly improved: for instance, is it necessary to mention a server as a separate entity if in practice all communication happens by means of All-Reduce?\n\nQuestions:\n* Do I understand correctly that Algorithm 1 mostly reduces to 1-bit Adam if we remove all lines required to update $c_t^{(l)}$ (in particular, lines 15–20) and change LAMB to ADAM for the preconditioning phase?", "summary_of_the_paper": "This work studies the problem of distributed training with large batches in a communication-bottlenecked setup, where regular versions of algorithms such as LAMB become a constraint. Authors propose a new algorithm, which compresses the gradient momentum before aggregation and then reconstructs the gradients to recover the scaling coefficients for LAMB. This idea allows authors to achieve faster convergence compared to naive compression for LAMB while maintaining higher communication efficiency than a non-compressed version.", "summary_of_the_review": "Although this paper proposes a straightforward yet definitely useful method for speeding up distributed training, in my opinion the clarity of presentation both for method and results can be improved. I give a score of a weak accept for now, with possible improvements if authors address the issues outlined in the review.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "DdN1jOXaqml", "reviewer": "Reviewer_otsu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "The paper is well-written. The idea of 1-bit lamb is well-supported by the convergence analysis and good empirical results. In overall, I think this is a good paper.\nHere are my concerns:\n1. All the experiments are limited to BERT, which makes it unknown whether the proposed algorithms could be applied to other training tasks including cv and nlp. To justify the effectiveness of the proposed algorithm, I recommend to add at least 1 CV task, such as resnet on cifar100 or imagenet.\n2. My major concern is in the theoretical analysis. Some part of the theoretical analysis is confusing to me. I cannot find the definition of $V_{max}$, so I simply guess it is the maximum value of $V$. However, since the diagonal elements of $V$ are actually $1/v$, $V_{max}$ is actually assuming a lower bound of the gradient $\\|\\nabla f(x)\\|$. Such a lower bound seems unreasonable to me since the convergence of the algorithm actually means that $\\|\\nabla f(x)\\|$ should approach 0, which contradicts the lower bound of $\\|\\nabla f(x)\\|$. Could the authors resolve such an issue? ", "summary_of_the_paper": "In this paper, the authors propose 1-bit LAMB, which combines lamb optimizer with communication compression, with convergence analysis and experiment results.", "summary_of_the_review": "The paper is well-written. The idea of 1-bit lamb is well-supported by the convergence analysis and good empirical results. In overall, I think this is a good paper. Some other experiments is recommended. The main issue is in the assumption of the theoretical analysis.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "01m0v29qO6", "reviewer": "Reviewer_amJh", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Strengths\n\n1. An interesting insight that the norm of LAMB's variance is more unstable than Adam\n\n2. Comprehensive experimental results to validate the claims, including the motivation and proposed method.\n\nWeaknesses\n\n1. Given the prior work 1-bit Adam, the novelty becomes somewhat incremental. The novel part is that 1-bit LAMB will try to adjust the scaling coefficient of LAMB. The efficient communication of momentum follows existing communication-efficient momentum SGD method.\n\nThe reconstruction of the gradient means that other compressors (e.g., Top-K) may not be incorporated into the proposed method.\n\n2. Does 1-bit LAMB preserve the same performance as LAMB with different number of workers and batch size? It would be great to have a table summarizing the performance with different #worker and batch size settings. Figure 3 only gives the training throughput.\n\n3. The theoretical analysis does not seem to validate the coefficient scaling in 1-bit LAMB. It may need further heuristic and theoretical elaboration.\n\n4. The 3rd assumption in the theoretical analysis does not seem very common to me in communication-efficient distributed training. The author needs to validate it. A more commonly-used assumption about the compressor is $||C(x)-x||^2 \\leq (1-\\delta)||x||^2$, from which we may know more about the effect of the compression ratio. Therefore, the convergence analysis of 1-bit LAMB seems less solid to me.", "summary_of_the_paper": "The paper proposes a communication-efficient distributed LAMB optimizer with 1-bit compression. It follows previous work to first warm-up the variance, but proposes to inference the scaling factor based on reconstructed variance. Experiments show training speedup due to communication compression. The proposed 1-bit LAMB achieves similar model performance to full-precision LAMB.", "summary_of_the_review": "The paper propose a communication-efficient large-batch LAMB optimizer to accelerate distributed training. The experimental results looks promising, but there could be some improvements both empirically and theoretically.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose a communication-efficient distributed LAMB optimizer using a 1-bit compression. This work is similar in spirit to other prior work, eg 1-bit Adam. Although the algorithm works reasonably well it is a bit unclear how much compression is achieved. Overall, the algorithmic novelty is limited, given the prior work, and the benefits of the algorithm don't shine through as the experiments are quite limited in their data sets and models. The theoretical results are also of unclear usefulness due to the assumptions made.
Decision: Reject