Year: 2021

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 limitedbandwidth 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 communicationefficient 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 BERTLarge 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 timewise speedup, and the same sample-wise convergence speed (and same fine-tuning task accuracy) compared to uncompressed LAMB.

heading: 1 INTRODUCTION
text: Training large-scale deep learning models in a distributed fashion is computation-heavy and expensive (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 up 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 convergence speed, thereby reduce BERT pre-training communication volume by 5x (Tang et al., 2021).
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 distributed training with LAMB, even with the larger batch sizes. Previous study shows that Adam provides slower convergence speed compared to LAMB at batch sizes 16K or larger for BERT pre-training (You et al., 2020). Using the same methodology, our BERT experiments show that 1-bit Adam, similar
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 3).
The studies and challenges above motivate us to design a new algorithm called 1-bit LAMB (Section 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 a 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 backend can be applied to 1-bit LAMB, 1-bit Adam, and other communication compression algorithms. We evaluate 1-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 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 same GLUE/SQuAD fine-tuning task accuracy. The 1-bit LAMB optimizer as well as the NCCL-based communication backend has been open sourced in a deep learning optimization library (name hidden to maintain anonymity).

heading: 2 RELATED WORK AND BACKGROUND
text: 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 al., 2015), and gradient compression/quantization which we focus on in this paper. Before communication, we could compress the original gradient g into Cω[g], where Cω[·] is the compress operator1. 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; Agarwal et al., 2018; Yu et al., 2019; Spring et al., 2019; Ivkin et al., 2019; Shi et al., 2021). Several works focus on unbiased compression methods (original and compressed tensors have the same expectation), such as centralized compressed parallel SGD (Alistarh et al., 2017) and many others (Wangni et al., 2018; Shen et al., 2018; Zhang et al., 2017; Wen et 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 et al., 2019; Ivkin et al., 2019; Sun et al., 2019; Basu et al., 2019; Vogels et al., 2019; Tang et al., 2021).
The error-compensated compression is proposed in the 1-bit SGD work (Seide et al., 2014): instead of compressing the gradient at each iteration directly, they compress the sum of the gradient and the last step’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 a scale). Recent works provide theoretical guarantee of this method (Bernstein et al., 2019), and also demonstrate that it admits the same asymptotic convergence rate as the uncompressed one (Stich 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.
1Cω[·] could also include randomness.
Adam (Kingma & Ba, 2015) can be viewed as SGD with momentum and adaptive learning rate scaling on each coordinate of the gradient. It has demonstrated promising convergence speed and hyperparameter robustness on many deep learning tasks. Recently, Tang et al. (2021) proposed 1-bit Adam which combines the efficiency of error-compensated 1-bit compression with Adam’s convergence speed. They show that error-compensated 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 at an early stage of training. To this end, they design a 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 as a fixed precondition, and communicate based on the momentum applied with error-compensated 1-bit compression. Their experiments on up to 256 GPUs show that 1-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 at large scale, being able to support large minibatches while keeping the convergence speed is a 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 as Adam with adaptive layerwise learning rates. By using LAMB, they are able to scale the batch size of BERT 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 as2:
m (l) t =β1m (l) t−1 + (1− β1)g (l) t ,v (l) t = β2v (l) t−1 + (1− β2)(g (l) t ) 2
u (l) t =
m (l) t√
v (l) t + η
, c (l) t = clip ( ‖x(l)t−1‖ ‖u(l)t ‖ , cmin, cmax )
x (l) t =x (l) t−1 − γc (l) t u (l) t .
(1)
Here g(l)t = ∇F (xt; ξt),m (l) t , v (l) t , x (l) t denote the stochastic gradient, momentum, second moment (i.e., the variance), and the model parameters at the model’s l-th layer at step t; β1 and β2 are the decaying factor; γ is the learning rate; η is an additive constant to avoid division by 0; clip(x, a, b) := min{max{x, a}, b} is the clipping operation3; c(l)t is a layer-wise scaling factor that regulates the update of x(l)t into certain range. One thing to note is that within each layer, each tensor (e.g., weight and bias) will have its own scaling coefficient c(l)t . The underlying intuition of LAMB’s scaling coefficient is that when the update is relatively large compared to the parameter, we should apply a lower learning rate to that layer (and vice versa).

heading: 3 MOTIVATION AND INSIGHTS

heading: 3.1 1-BIT ADAM IS NOT SUFFICIENT FOR LARGE-BATCH DISTRIBUTED TRAINING
text: 1-bit Adam demonstrates the same convergence speed as Adam for BERT pre-training task with batch size 4K (Tang et al., 2021). On the other hand, the LAMB work shows that it is difficult to keep Adam’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 large-batch distributed training, we perform a similar experiment using BERT pre-training task at batch size 16K. Using You et al. (2020)’s training parameters and tuning procedure (details in Appendix A.1) for LAMB and Adam, we perform BERT pre-training with LAMB and 1-bit Adam, respectively. Then we use the two pre-trained BERT model to perform SQuAD 1.1 fine-tuning (details in Section 6). Results in Table 1 show that similar to Adam, 1-bit Adam has slower convergence speed compared to LAMB at larger batch size.
2Here (x)2, √ x and x
y all denote element-wise operations. For simplicity weight decay is omitted.
3In the LAMB paper the clip function is only applied to ‖x(l)t−1‖ without mentioning the exact clipping function configurations, and our experiments show that ‖x(l)t−1‖ varies a 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 we are able to achieve similar SQuAD accuracy compared to the original LAMB.

heading: 3.2 COMMUNICATION IS STILL AN OVERHEAD FOR LARGE-BATCH DISTRIBUTED TRAINING
text: Tang et al. (2021) 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 on clusters with Ethernet and InfiniBand networks. To investigate whether the communication overhead still exists, we conduct a 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 one with 8 V100 GPUs per node and 100 Gigabit InfiniBand EDR inter-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.

heading: 3.3 INVESTIGATING BERT PRE-TRAINING WITH BASELINE LAMB
text: Tang et al. (2021) find that Adam’s variance term becomes stable at 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(l)t in (1)) depends on the variance. Thus we investigate how LAMB’s scaling coefficient and variance change during training using BERT pre-training task.
Figure 1 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 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 1-bit Adam’s strategy affect convergence when directly applied to LAMB.

heading: 3.4 EXISTING COMPRESSION STRATEGY AFFECTS LAMB’S CONVERGENCE
text: To combine 1-bit compression and large-batch training’s communication efficiency, first we attempt to directly apply 1-bit Adam’s strategy to 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 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 2 presents the BERT pre-training final validation loss when using vanilla LAMB, the experimental algorithm described in this section, and the proposed 1-bit LAMB algorithm (Section 4). We also use the pre-trained models to fine-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 as 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.

heading: 4 1-BIT LAMB ALGORITHM
text: 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 on a 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 x∈Rd
f(x) = 1
n n∑ i=1 Eξ(i)∼DiF (x; ξ (i))︸ ︷︷ ︸
:=fi(x)
, (2)
where d is the dimension of the input model x, Di is the data distribution of individual data sample ξ(i) on the i-th worker, F (x; ξ) is the loss function.
Notations and definitions Throughout this paper, we use the following notations:
• ∇f(·) denotes the gradient of a function f . • fi(x) := Eξ∼DiF (x; ξ). • ‖ · ‖p denotes the lp-norm for vectors and matrices. Notice p =∞ means the infinity norm. • Cω(·) denotes the randomized compressing operator, where ω denotes the random variable.
One example is the randomized quantization operator, for example, Cω(0.7) = 1 with probability 0.7 and Cω(0.7) = 0 with probability 0.3. • √ · denotes the square root of the argument. In this paper if the argument is a vector, then it
returns a vector taking the element-wise square root. • (x)2 denotes the element-wise square operation if x is a vector.
• ab or a/b denotes the element-wise division operation if both a and b are vectors and their dimension matches.
Algorithm 1 1-bit LAMB 1: Initialize: x(l)0 ,m (l) 0 = 0, v (l) 0 = 0, c (l) avg = 0 for each layer. Learning rate γ, initial error δ = 0, number
of total iterations T , warm-up steps Tw, three decaying factor β1, β2, β3 for LAMB’s momentum, variance, and scaling coefficient. r(l) = 1, rmin, rmax, rthreshold for 1-bit LAMB. We use 1-bit compression as the compressing operator, which means Cω[x] = ‖x‖‖Sign(x)‖Sign(x) where Sign(x) is the tensor that only keeps the sign of each element in the original tensor x.
2: Running the original LAMB in (1) for Tw steps, and at each step c (l) avg = β3c (l) avg + (1− β3)c(l)t . 3: At the end of step Tw, for each layer store the variance term (defined as v (l) t in (1)) v (l) Tw while still keep
updating v(l)t in the future steps. Also stop updating c (l) avg .
4: for t = Tw, . . . , T do 5: (On i-th node) 6: Randomly sample ξ(i)t and compute local stochastic gradient g (i) t := ∇Fi(x (i) t , ξ (i) t ), and update the
local momentumm(i)t according tom (i) t = β1m (i) t−1 + (1− β1)g (i) t .
7: Compress the fused momentum m(i)t into m̂ (i) t = Cω
[ m
(i) t + δ (i) t−1
] , and update the compression
error by δ(i)t =m (i) t + δ (i) t−1 − m̂ (i) t .
8: Send the m̂(i)t to the server. 9: (On server)
10: Take the average over all m̂(i)t it receives and compress it intomt = Cω [ 1 n ∑n j=1 m̂ (i) t + δt−1 ] , and
update the compression error accordingly by δt = 1n ∑n j=1 m̂ (i) t + δt−1 −mt.
11: Sendmt to all the workers. 12: (On j-th node) 13: Setmt =mt. 14: for the l-th layer do 15: Reconstruct global gradient g(l)t = (m (l) t − β1m (l) t−1)/(1− β1).
16: v(l)t = β2v (l) t−1 + (1− β2) ( g (l) t )2 .
17: r(l)t = ∥∥∥v(l)Tw /v(l)t ∥∥∥∞.
18: r(l)t = clip ( r (l) t , (1− rthreshold)× r (l) t−1, (1 + rthreshold)× r (l) t−1 ) .
19: r(l)t = clip ( r (l) t , rmin, rmax ) .
20: c(l)t = r (l) t c (l) avg .
21: Update model of the l-th layer x(l)t = x (l) t−1 − γc (l) t
m (l) t√
v (l) Tw
.
22: end for 23: end for 24: Output: x.
We summarize 1-bit LAMB in Algorithm 1. During the compression stage, we freeze the variance in order to apply the error compensation mechanism correctly (Tang et al., 2021). However, this brings two challenges for LAMB’s case: 1) We cannot update LAMB’s scaling coefficients (c(l)t in (1)) during compression stage based on 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 variance term is not stable for some layers but we freeze them during compression, we need to adjust scaling coefficients during compression to 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 of 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 to be used during compression stage. On the other hand, we still keep updating another “fresh” variance by reconstructing the global gradient based on this and last step’s compressed momentum.
To update LAMB scaling coefficient during compression stage, we compute a 1-bit LAMB scaling ratio which is the max element among the (frozen variance/fresh variance). We use the max element as the scaling ratio because it is effective and cheap to compute based on 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 A.9 provides a theoretical analysis for 1-bit LAMB.

heading: 5 PROPOSED SYSTEM DESIGN
text: To realize the compressed communication at system level, Tang et al. (2021) designed a custom collective primitive, called “compressed allreduce”, using Message Passing Interface (MPI). In addition to implementing 1-bit LAMB using this MPI-based backend, we introduce a 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 1. 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 NCCLbased implementation significantly improves the usability since NCCL is integrated with PyTorch distributed. In addition, evaluations show that the performance of the NCCL-based implementation is better than the MPI-based for Ethernet-based systems and on-par for InfiniBand-based systems. Thus we will mainly present the results with NCCL-based backend, but we also include a comparison between MPI and NCCL-based implementations in Appendix A.6.

heading: 6 EVALUATION
text: Dataset and models We evaluate the convergence and performance of 1-bit LAMB and uncompressed LAMB for BERT-Large (L = 24, H = 1024, A = 16, 340M params) pre-training task. We use the same dataset as Devlin et al. (2019), which is a concatenation of Wikipedia and BooksCorpus with 2.5B and 800M words respectively. Compared to the original BERT model, one notable change is that we applied PreLN instead of PostLN for better training stability (Zhang & He, 2020; Xiong et al., 2020). We use the GLUE fine-tuning benchmark (Wang et al., 2018) and SQuAD 1.1 fine-tuning task4 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 to 256 GPUs for BERT pre-training tasks to measure 1-bit LAMB’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 8 and 16 GPUs on the Ethernet cluster, respectively.
Training parameters For BERT pre-training, we set the parameters in (1) as β1 = 0.9, β2 = 0.999, cmin = 0.01 and cmax = 0.3 for LAMB and 1-bit LAMB. For 1-bit LAMB, we set β3 = 0.9, rmin = 0.5, rmax = 4.0, and rthreshold = 0.1 in Algorithm 1. 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× 10−3, exponentially increases to 12× 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 × 10−3 as a warmup in the first 150 steps, then decays into 0.9 of the original after
4https://rajpurkar.github.io/SQuAD-explorer/
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 on when do the LAMB’s scaling coefficients become stable as discussed in Section 3.3 thus could potentially be automated. For the other parameters (β3, rmin, rmax, and rthreshold in Algorithm 1) we find tuning them does not lead to very different convergence since they are just used to avoid extreme value/dramatic change of LAMB’s scaling coefficients. Similarly the clipping parameters cmin and cmax in (1) for LAMB and 1-bit LAMB do not require non-trivial tuning.
For GLUE benchmarks we use 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 finetune for 3 epochs for all GLUE tasks. For each task, we select the best learning rate among {2× 10−5, 3× 10−5, 4× 10−5, 5× 10−5}. For SQuAD fine-tuning we use Adam optimizer and the same parameters as published by HuggingFace (batch size = 24, learning rate = 3× 10−5, dropout = 0.1, 2 epochs).
Convergence analysis Figure 2 presents the BERT pre-training sample-wise convergence results. For both seqlen 128 and 512, 1-bit LAMB provides the same convergence speed as LAMB, while it takes much less time to process each batch due to communication compression. We already presented the BERT pre-training validation loss and SQuAD fine-tuning results in Section 3.4 Table 2. And Table 3 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 as 1/(warmup_ratio + (1 - warmup_ratio)/16) for FP16 training, 1-bit LAMB offers 4.6× and 4.1× 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 a throughput analysis where we run the warmup (i.e., baseline 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.5× speedup during the compression stage, which is equivalent to 2.8× end-to-end speedup (computed as 1/(warmup_ratio + (1 - warmup_ratio)/compression_stage_speedup)). For seqlen 512, 1-bit LAMB provides up to 3.0× speedup during the compression stage, which is equivalent to 2.2× end-to-end speedup. This demonstrates 1-bit LAMB’s better scalability compared to LAMB. It is also worth mentioning that 1-bit LAMB on Ethernet (4.1 Gbps effective bandwidth, 4 GPUs per node) is able
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, we also measure the total runtime of pre-training at batch size 64/32K for both 1-bit LAMB and LAMB. As shown in Table 4, overall 1-bit LAMB is able to provide 2.2× and 1.5× speedup for seqlen 128 and 512. These numbers are consistent with the endto-end speedup calculated in the throughput analysis (2.0× and 1.5× based on results in Figure 3(d) and 3(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 larger communication volume makes it more sensitive to the occasional fluctuation of the actual network bandwidth.

heading: 7 CONCLUSION
text: To reduce both the frequency and volume of communications for large-scale training, we propose an error-compensated 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.6× communication volume reduction and up to 2.8× 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.

heading: A APPENDIX

heading: A.1 TRAINING PARAMETERS FOR 1-BIT ADAM EXPERIMENTS IN SECTION 3.1
text: For experiments in Section 3.1, for both LAMB and 1-bit Adam we 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−3, 10% LR warmup, clipping configs (cmin and cmax in (1)) as 0.1 and 1. For 1-bit Adam, we use learning rates ∈ {1 × 10−4, 2 × 10−4, 3 × 10−4}, LR warmup ∈ {5%, 10%, 20%}. All of these training parameters (except LAMB clipping configs) are from the LAMB paper. For 1-bit Adam, following the original work’s strategy we set the number of warmup steps as 4000 (out of total 28125 steps) for seqlen 128 and 475 (out of 3125) for seqlen 512.

heading: A.2 DETAILED RESULTS FOR PROFILING EXPERIMENTS IN SECTION 3.2
text: 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 is smaller, when the network bandwidth is lower.

heading: A.3 DETAILED RESULTS FOR LAMB EXPERIMENTS IN SECTION 3.3
text: Figure 4 presents LAMB’s scaling coefficients for different layers during BERT pre-training sequence length 128 (sequence length 512 has similar patterns). Only the cls.seq_relationship.bias has a very unstable scaling coefficient. This is because this bias only has two elements, representing the two states of whether the two sequences are next to each other.
Figure 5 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 varaince terms are less stable: 1) 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 in Adam. And each separate variance could have less stable norm compared to a single fused variance.

heading: A.4 DESIGN OF THE EXPERIMENTAL ALGORITHM IN SECTION 3.4
text: Because LAMB’s scaling coefficient is essentially part of the learning rate (and it’s just a scalar), updating this coefficient would not affect communication compression’s error compensation mechanism
as long as the change is small enough between two steps.5 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 at 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 scaling coefficients during compression stage based
5In fact as described in Section 4, the proposed 1-bit LAMB algorithm does require updating LAMB’s scaling coefficient during compression stage, but in a way different from original LAMB algorithm.
on original LAMB algorithm would produce suboptimal results and slower convergence speed. Thus for Section 3.4’s experimental algorithm, we stop updating LAMB’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 we also tried communicating based on 1-bit compressed gradient, but it leads to much slower convergence speed. We believe it’s because gradients are less stable, which could lead to higher compression error and slower convergence.
A.5 1-BIT LAMB ALGORITHM DESIGN AND IMPLEMENTATION DETAILS
Figure 6 presents 1-bit LAMB scaling ratios (r(l)t 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 a result 1-bit LAMB is able to adaptively update LAMB scaling coefficients during compression according to the difference between the frozen and fresh variance.

heading: A.5.1 REDUCING NUMBER OF COMMUNICATIONS BY MOMENTUM FUSION
text: 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 we fuse all layers’ momentum into a contiguous 1D buffer and just do one communication over this fused momentum. This fused momentum is not a duplication, but a different “view” of all layers’ momentum. We implement this momentum fusion by torch._utils._flatten_dense_tensors and torch._utils._unflatten_dense_tensors.

heading: A.5.2 REDUCING COMPRESSION ERROR BY MOMENTUM SCALING
text: We find that after momentum fusion, the compression error increase a 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 to the same scale due to the nature of
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 layerwise adaptive learning rate, each layer is learning at 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 a 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, it won’t affect 1-bit compression’s error compensation mechanism.

heading: A.6 COMPARING MPI AND NCCL-BASED COMMUNICATION BACKEND
text: 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 NCCLbased 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).

heading: A.7 ADDITIONAL CONVERGENCE ANALYSIS ABOUT NUMBER OF WORKERS AND BATCH SIZE
text: 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 of 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 Section 6. Results show that both LAMB and 1-bit LAMB achieve similar SQuAD F1 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 at batch size 16K and 4K. At batch size 16K, for LAMB we use the same training parameters described in Appendix A.1, 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 112500/12500 steps for seqlen 128/512, weight decay = 0.01, linear LR warmup and decay, learning rate = 1.77× 10−3, 2.5% LR warmup, clipping configs (cmin and cmax in (1)) 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 other hyperparameters in Algorithm 1, we keep them the same under all batch sizes: β3 = 0.9, rmin = 0.5, rmax = 4.0, and rthreshold = 0.1. As shown in Table 7, results show that 1-bit LAMB provides better SQuAD F1 scores than LAMB under all batch sizes.

heading: A.8 ADDITIONAL CONVERGENCE ANALYSIS: RESNET-50 ON CIFAR100
text: This section provides additional analysis of training ResNet-50 on CIFAR100. Because ResNet-50 is much smaller than BERT-Large and because we use 8 GPUs on a single node for the training, 1-bit LAMB and LAMB have very similar throughput. Thus this experiment focus on 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 ∈ {0.005, 0.01, 0.02}, linear LR warmup in the first 1K steps, no LR decay, weight decay = 0.01. For LAMB we set the clipping configs (cmin and cmax in (1)) as 0.01 and 0.3. For 1-bit LAMB we set the number of warmup steps as 1K (8%), and set other hyperparameters in Algorithm 1 as: β3 = 0.9, rmin = 0.5, rmax = 2.0, and rthreshold = 0.1. Figure 8 presents the sample-wise training loss and testing accuracy, and Table 8 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.

heading: A.9 THEORETICAL ANALYSIS
text: Notice that for 1-bit LAMB, we only use original LAMB at warm-up, and then we essentially run error-compensated momentum SGD with coordinate-dependent learning rate γ√
v Tw
. Therefore here
we consider the LAMB-based warm-up phase as a way to find a good precondition variance term v Tw to be used in the compression phase. Below we are going to introduce the convergence rate for the compression phase after warm-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(·) is assumed to be with L-Lipschitzian gradients, which means ‖∇f(x)−∇f(y)‖ ≤ L‖x− y‖, ∀x,∀y,
2. Bounded variance: The variance of the stochastic gradient is bounded
Eζ(i)∼Di‖∇F (x; ζ (i))−∇f(x)‖2 ≤ σ2, ∀x,∀i.
3. Bounded magnitude of error for Cω[·]: The magnitude of worker’s local errors δ(i)t and the server’s global error δt, are assumed to be bounded by a constant
n∑ k=1 Eω ∥∥∥δ(i)t ∥∥∥ ≤ 2 , n∑ i=1 Eω ∥∥δt∥∥ ≤ 2 , ∀t, ∀i.
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 one non-zero element in the tensor. In (Tang
et al., 2019), authors proved that our assumption is equivalent to the signal-to-noise assumption when the gradient magnitude is upper bounded.
Next we present the main theorem for 1-bit LAMB. Theorem 1. Under Assumption 1, for 1-bit LAMB, we have the following convergence rate(
1− γmaxLVmax − 8γ3maxL 2
(1− β)2γmin ) T∑ t=0 E‖∇f(xt)‖2
≤2Ef(x1)− 2Ef(y ∗)
γmin +
24γ3maxL 2V 3max 2T
(1− β)2γmin + LVmaxγ
2 maxσ 2T
nγmin +
8γ3maxL 2V 2maxσ 2T
n(1− β)2γmin , (3)
where V = diag ( 1/v
(1) Tw , 1/v (2) Tw , · · · , 1/v(d)Tw
) is a diagonal matrix spanned by v
Tw .
Here 1/v is a moving sum over all history (∇f(x))2, not ‖∇f(x)‖. Therefore it will not approach to 0 for the following two reasons: i) We stop updating v after the warmup stage, which means ∇f(x) at this time will not approach 0 since it has not converged yet; ii) Since β2 = 0.999, this means v is heavily dependent to the history gradients, therefore even though ∇f(x) approaches to 0, v itself will still taking a time to decay to 0. Therefore it is reasonable for us to assume an upper bound for Vmax
Given the generic result in Theorem 1, we obtain the convergence rate for 1-bit LAMB with appropriately chosen learning rate γ. Corollary 1. Under Assumption 1, for 1-bit LAMB, choosing γ = γmax = γmin =
1−β
16LVmax+σ √ T n +T 1 3 2 3 , we have the following convergence rate
Vmax T T−1∑ t=0 E‖∇f(xt)‖2V . σ√ nT + 2 3 T 2 3 + 1 T ,
where we treat f(x1)− f∗, β and L as constants.
This result suggests that: 1-bit LAMB essentially admits the same convergence rate as distributed SGD in the sense that both of them admit the asymptotical convergence rate O(1/ √ nT ), which means we can still achieve linear speedup w.r.t. the number of workers n.
For theoretical analysis, we follow the framework of previous work (Tang et al., 2021). Below we first summarize the updating rule of 1-bit LAMB, followed by the proof of Theorem 1, then the proof of Corollary 1.

heading: A.9.1 PROOF TO THE UPDATING FORM
text: Since our algorithm is equivalent to running a parameter-server prototype communication on each chunk of the gradient, so below we will assume a parameter-server model (which means the tensor is not required to be divided into n chunks) for simplicity. Although we use a layer-wise updating strategy in 1-bit LAMB, for the simplicity of theoretical analysis, 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 1, at iteration t + 1, the updating step of the momentum termmt+1 can be divided into two steps:
1. Local Update and Compress: each worker locally updatemt and use the error-compensate strategy for compressing.
m (i) t =βmt + (1− β)g (i) t
m (i)
t+ 12 =Cω
[ m
(i) t + ct−2 ct−1 δ (i) t ] δ (i) t+1 =m (i) t + δ (i) t −m (i)
t+ 12 .
2. All workers send its m(i) t+ 12 to the server. The server takes the average over them and compress it again using error-compensation.
mt+ 12 = 1
n n∑ i=1 m (i) t+ 12
mt+1 =Cω [ mt+ 12 + ct−2 ct−1 δt ] δt+1 =mt+ 12 + δt −mt+1.
3. The server broadcastmt+1 to all workers, and all workers update the local model according to
xt+1 = xt − γct−1 mt+1√ v2Tw .
Notice that here we do not include the updating rule for {ct} because for the following part we will only view it as a scaling factor.
So actually the updating rule above can be summarized as mt+1 =Cω [ mt+ 12 + ct−2 ct−1 δt ] =mt+ 12 + ct−2 ct−1 δt − δt+1 (from the definition of δt+1)
= 1
n n∑ i=1 Cω [ m (i) t + ct−2 ct−1 δ (i) t ] + ct−2 ct−1 δt − δt+1
= 1
n n∑ i=1 ( m (i) t + ct−2 ct−1 δ (i) t − δ (i) t+1 ) + ct−2 ct−1 δt − δt+1 (from the definition of δ(i)t+1)
=βmt + 1− β n n∑ i=1 g (i) t + ct−2 ct−1
( 1
n n∑ i=1 δ (i) t + δt
) − ( 1
n n∑ i=1 δ (i) t+1 + δt+1
) .
Denote
gt = 1
n n∑ i=1 g (i) t
δt = 1
n n∑ i=1 δ (i) t + δt,
the update rule ofmt can be summarized as
mt = βmt−1 + (1− β)gt + ct−2 ct−1 δt−1 − δt,
and
xt+1 = xt − γV ct−1mt,
where V = diag(1/ √ v1, 1/ √ v2, · · · , 1/ √ vd) is the a diagonal matrix that spanned with vTw . In order to simlify the notation, we define γt := γct−1 as a time-varying learning rate, which reduces the updating rule above into
mt =βmt−1 + (1− β)gt + γt−1 γt δt−1 − δt,
xt+1 =xt − γtVmt.

heading: A.9.2 PROOF TO THEOREM 1
text: Notice that in for 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 1 2 z),
also
h(z) = f(V 1 2 z),
where V is a diagonal matrix.
In this case we have
V 1 2∇f(V 12 z) = ∇h(z).
Therefore the updating rule of 1-bit LAMB in the view of h(·) is
V 1 2 zt+1 = V 1 2 zt − γV 1 2 ( V 1 2mt ) .
It can be easily verified that
mt =(1− β) t∑
s=0
βt−sgs + t∑ s=0 βt−s(δs−1 − δs)
=(1− β) t∑
s=0
βt−s 1
n n∑ i=1 ∇F (V 12 zt; ξ(i)t ) + t∑ s=0 βt−s(δs−1 − δs)
which means
V 1 2mt =(1− β) t∑ s=0 βt−s 1 n n∑ i=1 V 1 2∇F (V 12 zt; ξ(i)t ) + t∑ s=0 βt−sV 1 2 (δs−1 − δs)
=(1− β) t∑
s=0
βt−s 1
n n∑ i=1 ∇H(V 12 zt; ξ(i)t ) + t∑ s=0 βt−sV 1 2 (δs−1 − δs)
=(1− β) t∑
s=0
βt−sgs(z) + t∑ s=0 βt−sV 1 2 (δs−1 − δs),
where gs(z) is the corresponding averaged stochastic gradient computed in the view of loss function h(·).
Then, if we definemt(z) = V 1 2mt, the updating rule ofmt(z) admits
mt(z) = βmt−1(z) + (1− β)gt(z) + V 1 2 δt−1 − V 1 2 δt, (4)
and
V 1 2 zt+1 =V 1 2 zt − γtV 1 2mt(z)
zt+1 =zt − γtmt(z). (5)
From (4) and (5) 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 error-compensated momentum SGD using a coordinate-independent time-varying learning rate.
Below are some critical lemmas for the proof of Theorem 1.
Lemma 1. Given two non-negative sequences {at}∞t=1 and {bt}∞t=1 that satisfying
at = t∑ s=1 ρt−sbs, (6)
with ρ ∈ [0, 1), we have
Dk := k∑ t=1 a2t ≤ 1 (1− ρ)2 k∑ s=1 b2s.
Proof. From the definition, we have
Sk = k∑ t=1 t∑ s=1 ρt−sbs = k∑ s=1 k∑ t=s ρt−sbs = k∑ s=1 k−s∑ t=0 ρtbs ≤ k∑ s=1 bs 1− ρ , (7)
Dk = k∑ t=1 t∑ s=1 ρt−sbs t∑ r=1 ρt−rbr
= k∑ t=1 t∑ s=1 t∑ r=1 ρ2t−s−rbsbr
≤ k∑ t=1 t∑ s=1 t∑ r=1 ρ2t−s−r b2s + b 2 r 2
= k∑ t=1 t∑ s=1 t∑ r=1 ρ2t−s−rb2s
≤ 1 1− ρ k∑ t=1 t∑ s=1 ρt−sb2s
≤ 1 (1− ρ)2 k∑ s=1 b2s, (due to (7))
which completes the proof.
Lemma 2. Under Assumption 1, for any sequence that follows the updating rule of
xt+1 =xt − γtmt mt =βmt−1 + (1− β)gt +
γt−1 γt δt−1 − δt,
if
Egt = ∇f(xt), E‖gt −∇f(xt)‖2 ≤ σ2
n , E‖δt‖2 ≤ 2, ∀t,
‖∇f(x)−∇f(y)‖ ≤ L‖x− y‖, ∀x,∀y,
then we can guarantee that( 1− γL− 2γ 2L2
(1− β)2 ) T∑ t=0 E‖∇f(xt)‖2
≤2Ef(x1)− 2Ef(x ∗)
γ +
6γ2L2 2T (1− β)2 + Lγσ2T n + 2γ2L2σ2T n(1− β)2
Proof. Instead of investigating xt directly, we introduce the following sequence
yt =xt − γtmt + γt−1δt
1− β .
The updating rule of yt admits
yt+1 − yt =xt+1 − xt − γt
1− β (mt+1 −mt) + γtδt+1 − γt−1δt 1− β
=− γtmt − γt
1− β (βmt + (1− β)gt+1 − δt+1 + γt−1 γt δt) + γtδt+1 − γt−1δt 1− β
=− γtgt+1. Since f(·) is with L-Lipschitzian, we have Ef(yt+1)− Ef(yt)
≤E 〈∇f(yt),yt+1 − yt〉+ L
2 E ‖yt+1 − yt‖2
=− γtE 〈∇f(yt), gt+1〉+ Lγ2t 2 E‖gt+1‖2 =− γtE 〈∇f(yt),∇f(xt+1)〉+ Lγ2t 2 E‖gt+1‖2
=− γt 2 E‖∇f(xt+1)‖2 − γt 2 E‖∇f(yt)‖2 + γt 2 E‖∇f(xt+1)−∇f(yt)‖2 + Lγ2t 2 E‖gt+1‖2
≤− γt 2 E‖∇f(xt+1)‖2 + Lγt 2 E‖xt+1 − yt‖2 + Lγ2t 2 E‖gt+1‖2
=− γt 2 E‖∇f(xt+1)‖2 +
γ3tL 2
2 E ∥∥∥∥∥ (2− β)mt1− β + γt−1 γt δt−1 1− β ∥∥∥∥∥ 2 + Lγ2t 2 E‖gt+1‖2
≤− γt 2 E‖∇f(xt+1)‖2 +
4γ3tL 2
(1− β)2 E‖mt‖2 +
4γ3tL 2
(1− β)2 E‖δt−1‖2 + Lγ2t 2 E‖gt+1‖2 (rthreshold < 1)
≤− γt 2 E‖∇f(xt+1)‖2 +
4γ3tL 2
(1− β)2 E‖mt‖2 +
4γ3t 2 (1− β)2 + Lγ2t 2 E‖gt+1‖2
≤− γt 2 E‖∇f(xt+1)‖2 +
4γ3tL 2
(1− β)2 E‖mt‖2 +
4γ3t 2 (1− β)2 + Lγ2t 2 E‖∇f(xt+1)‖2 + Lγ2t σ 2 2n .
Summing up the equation above from t = 0 to t = T we get Ef(yT+1)− Ef(y0) ≤ T∑ t=0 − (1− γtL)γt 2 E‖∇f(xt)‖2 + T∑ t=0 4γ3tL 2 (1− β)2 E‖mt‖2 + 4L2 2 ∑T t=0 γ 3 t (1− β)2 + Lσ2 ∑T t=0 γ 2 t 2n .
Therefore if we γt < 12L and it’s within certain range γt ∈ [γmin, γmax], the equation above leads to
(1− γmaxL) T∑ t=0 E‖∇f(xt)‖2
≤2Ef(y0)− 2Ef(yT+1) γmin + 8γ3maxL 2 (1− β)2γmin T∑ t=0 E‖mt‖2 + 8γ3maxL 2 2T (1− β)2γmin + Lγ2maxσ 2T nγmin . (8)
Notice that we have
mt =(1− β) t∑
s=0
βt−sgs + t∑ s=0 βt−s(δs−1 − δs)
which by using Lemma 1, we have T∑ t=0 ‖mt‖2 ≤ T∑ t=0 ‖gt‖2 + 2 (1− β)2 T∑ t=0 ‖δt‖2 ≤ T∑ t=0 ‖∇f(xt)‖2 + σ2T n + 2 2T (1− β)2 . (9)
Combing (8) and (9) together we get( 1− γmaxL− 8γ3maxL 2
(1− β)2γmin ) T∑ t=0 E‖∇f(xt)‖2
≤2Ef(y0)− 2Ef(yT+1) γmin + 24γ3maxL 2 2T (1− β)2γmin + Lγ2maxσ 2T nγmin + 8γ3maxL 2σ2T n(1− β)2γmin ≤2Ef(x1)− 2Ef(y ∗)
γmin +
24γ3maxL 2 2T (1− β)2γmin + Lγ2maxσ 2T nγmin + 8γ3maxL 2σ2T n(1− β)2γmin .
Proof to Theorem 1 Since using a per-coordinate learning rate for loss function f(·) is equivalent to use a constant learning for all coordinates but for loss function h(·), the only two thing that change are
• Different L-Lipschitzian coefficient: the L-Lipschitzian coefficient for h(·) is ‖∇h(x)−∇h(y)‖2 = ∥∥∥V 12∇f(V 12x)− V 12∇f(V 12y)∥∥∥2
= ∥∥∥∇f(V 12x)−∇f(V 12y)∥∥∥2
V ≤L2 ∥∥∥V 12x− V 12y∥∥∥2
V
=L2‖x− y‖2V 2 ≤L2V 2max‖x− y‖2.
Therefore the effective L-Lipschitzian coefficient of h(x) is LVmax • Different definition of δt: from (4) we shall see that actually the compression error in the
view of h(·) is V 12 δt, so in this case we have
E‖V 12 δt‖2 ≤ Vmax 2
Proof. From Lemma 2, we have( 1− γmaxLVmax − 8γ3maxL 2
(1− β)2γmin ) T∑ t=0 E‖∇f(xt)‖2
≤2Ef(x1)− 2Ef(y ∗)
γmin +
24γ3maxL 2V 3max 2T
(1− β)2γmin + LVmaxγ
2 maxσ 2T
nγmin +
8γ3maxL 2V 2maxσ 2T
n(1− β)2γmin .

heading: A.9.3 PROOF TO COROLLARY 1
text: Proof. By choosing γmax = γmin = 1−β 16LVmax+σ √ T n +T 1 3 2 3 , we can guarantee that
1− γL− 8γ 2L2V 2max (1− β)2 ≥1 2 .
So (3) leads to
T∑ t=0 E‖∇f(xt)‖2V ≤ 2 (Ef(y0)− f(y∗)) (1− β)
( 16LVmax + σ √ T
n + T
1 3 2 3 ) + ( (1− β)L √ T + 2L2V 2max ) σ√ n + 32L2 2 3T 1 3V 3max
1
T T∑ t=0 E‖∇f(xt)‖2V ≤ 2 (Ef(y0)− f(y∗)) (1− β) ( 16LVmax T + σ√ nT + T − 2 3 2 3 ) + ( (1− β)L+ 2L
2V 2max√ T ) σ√ nT + 32L2 2 3T− 2 3V 3max.
Treating f(y1)− f∗, β and L as constants, from the inequality above we get
1
T T∑ t=0 E‖∇f(xt)‖2 . σ√ nT + 2 3 T 2 3 + 1 T .
It completes the proof.

Review summaries:
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.

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.