[论文解读] 1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed

Hanlin Tang, Shaoduo Gan+7 more|arXiv (Cornell University)|Feb 4, 2021

Advanced Neural Network Applications参考文献 39被引用 15

一句话总结

本文提出1-bit Adam，一种通信高效的优化器，通过在预热阶段使用Adam对动量SGD进行预条件处理，将Adam的收敛速度与1-bit梯度压缩相结合。在BERT-Large和ResNet训练中，其通信量最多减少5倍，吞吐量最高提升3.3倍，同时保持与原始Adam相同的收敛速度和精度。

ABSTRACT

Scalable training of large models (like BERT and GPT-3) requires careful optimization rooted in model design, architecture, and system capabilities. From a system standpoint, communication has become a major bottleneck, especially on commodity systems with standard TCP interconnects that offer limited network bandwidth. Communication compression is an important technique to reduce training time on such systems. One of the most effective methods is error-compensated compression, which offers robust convergence speed even under 1-bit compression. However, state-of-the-art error compensation techniques only work with basic optimizers like SGD and momentum SGD, which are linearly dependent on the gradients. They do not work with non-linear gradient-based optimizers like Adam, which offer state-of-the-art convergence efficiency and accuracy for models like BERT. In this paper, we propose 1-bit Adam that reduces the communication volume by up to $5 imes$, offers much better scalability, and provides the same convergence speed as uncompressed Adam. Our key finding is that Adam's variance (non-linear term) becomes stable (after a warmup phase) and can be used as a fixed precondition for the rest of the training (compression phase). Experiments on up to 256 GPUs show that 1-bit Adam enables up to $3.3 imes$ higher throughput for BERT-Large pre-training and up to $2.9 imes$ higher throughput for SQuAD fine-tuning. In addition, we provide theoretical analysis for our proposed work.

研究动机与目标

• 解决在有限网络带宽的通用硬件上大规模分布式训练中的通信瓶颈问题。
• 克服误差补偿压缩与非线性优化器（如Adam）之间的不兼容性，而这类优化器对训练最先进的模型（如BERT）至关重要。
• 在实现激进的1-bit梯度压缩以减少通信量的同时，保持Adam的快速收敛速度。
• 设计一种混合优化策略，在预热阶段利用Adam的自适应学习率，主训练阶段采用压缩的动量SGD。
• 在使用标准TCP互连的大型模型（如BERT和ResNet）上实现高可扩展性和端到端吞吐量提升。

提出的方法

• 提出两阶段优化框架：首先使用全精度Adam进行预热阶段（通常占训练步数的20%），然后切换到1-bit压缩的动量SGD以完成剩余训练。
• 利用观察结果：Adam的方差（非线性项）在训练早期即趋于稳定，使其能够作为后续压缩优化中动量SGD的预条件器。
• 对SGD中的动量项而非梯度应用误差补偿压缩，确保在1-bit压缩下仍能保持收敛稳定性。
• 设计一种基于MPI的自定义集体通信原语，以高效处理减少后的通信量并实现高可扩展性。
• 将1-bit Adam优化器和通信后端集成到DeepSpeed中，支持生产部署并开源实现代码。

实验结果

研究问题

• RQ1 能否有效将误差补偿压缩应用于非线性优化器（如Adam），而这类优化器对训练现代模型（如BERT）至关重要？
• RQ2 Adam中的非线性方差项是否在训练早期保持稳定，从而可作为压缩优化的预条件器？
• RQ3 结合Adam与压缩动量SGD的混合优化器能否在减少通信量的同时，实现与全精度Adam相同的收敛速度？
• RQ4 在256张GPU上，此类混合优化器在大规模模型（如BERT和ResNet）上的可扩展性和端到端吞吐量提升如何？
• RQ5 1-bit Adam与直接对Adam梯度进行朴素1-bit压缩相比性能如何？是否保持了最终模型的精度？

主要发现

• 1-bit Adam相比未压缩的Adam，通信量最多减少5倍，且不损失收敛速度或最终模型精度。
• 在64张GPU上使用4K批量大小进行BERT-Large预训练时，1-bit Adam在51.5小时内完成训练，相比基线Adam（174.3小时）实现3.4倍加速。
• 在最多256张GPU上，1-bit Adam相比未压缩Adam实现高达3.3倍的端到端吞吐量提升，且在以太网上的可扩展性优于InfiniBand。
• 在BERT预训练、SQuAD微调、CIFAR-10上的ResNet-18以及DCGAN训练中，该方法在样本级收敛速度上与全精度Adam保持一致。
• 在10Gbps以太网上，1-bit Adam的吞吐量可与100Gbps InfiniBand上的Adam相媲美，展现出强大的硬件效率。
• 实验表明，对Adam梯度进行朴素1-bit压缩会失败收敛，而采用预条件动量SGD的1-bit Adam则保持了鲁棒的性能。