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