1-bit LAMB: Communication Efficient Large-Scale Large-Batch Training with LAMB's Convergence Speed

1 INTRODUCTION . 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 can not 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 ) . 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 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 ) . 3 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 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 .

Summary: 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.