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

Abstract

To train large models (like BERT and GPT-3) with hundreds or even thousands of GPUs, the communication has become a major bottleneck, especially on commodity systems with limited-bandwidth TCP interconnects network. On one side large-batch optimization such as LAMB algorithm was proposed to reduce the number of communications. On the other side, communication compression algorithms such as 1-bit SGD and 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 on low-bandwidth Ethernet networks. 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 even when communication is compressed. In addition, we introduce a new system implementation for compressed communication using the NCCL backend of PyTorch distributed, which improves both usability and performance compared to existing MPI-based implementation. 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.6× communication volume reduction, up to 2.8× end-to-end speedup (in terms of number of training samples per second), and the same convergence speed (in terms of number of pre-training samples to reach the same accuracy on fine-tuning tasks) compared to uncompressed LAMB.

The effect of choosing optimizer algorithms to improve computer vision tasks: a comparative study

“…This momentum is accomplished by introducing two new variables, namely, velocity and friction, as given by Eqs. ( 3) and (4), respectively [46].…”

“…Then, the resulting output is propagated into the model to lessen the difference. The DL architecture adjusts the weights and repeats the process until a convergence is achieved [46,77]. An algorithm is searched to speed up the learning process while producing the best results.…”

Maximizing Communication Efficiency for Large-scale Training via 0/1 Adam

“…Bernstein et al (2018b); Sohn et al (2019); Le Phong & Phuong (2020); Lyu (2021) investigate the robustness of 1-bit SGD. Perhaps the closest works to this paper are (Tang et al, 2021;Li et al, 2021), which propose using two-stage training to enable 1-bit Adam and 1-bit Lamb, respectively. Among all the variants of 1-bit communication, the design with error feedback mechanism has shown to work best both empirically (Seide et al, 2014) and theoretically (Karimireddy et al, 2019).…”

“…t since the gradients are usually high-dimensional. Based on the profiling results from (Tang et al, 2021;Li et al, 2021), the communication of gradients could take up to 94% of the total training time on modern clusters. 1-bit compression (Liu et al, 2018) mitigates this problem by sending each gradient with only signs and a shared, usually the average over all the coordinates, magnitude.…”

MiCS: Near-linear Scaling for Training Gigantic Model on Public Cloud

“…To improve the training efficiency, more recent works advocate large-batch training [15,23,50,51,55]. However, due to the limited device memory, practitioners have to resort to gradient accumulation that divides a large batch into multiple micro-batches and accumulates the gradient w.r.t.…”