[Paper Review] APMSqueeze: A Communication Efficient Adam-Preconditioned Momentum SGD Algorithm

Hanlin Tang, Shaoduo Gan+5 more|arXiv (Cornell University)|Aug 26, 2020

Vehicle License Plate Recognition35 references4 citations

TL;DR

APMSqueeze is a communication-efficient optimization algorithm that preconditions momentum SGD using Adam for the first phase of training, then applies error-compensated gradient compression to reduce communication by up to 97% (32× less data), achieving up to 10× end-to-end speedup on BERT and ResNet-18 without sacrificing convergence or accuracy.

ABSTRACT

Adam is the important optimization algorithm to guarantee efficiency and accuracy for training many important tasks such as BERT and ImageNet. However, Adam is generally not compatible with information (gradient) compression technology. Therefore, the communication usually becomes the bottleneck for parallelizing Adam. In this paper, we propose a communication efficient {\bf A}DAM {\bf p}reconditioned {\bf M}omentum SGD algorithm-- named APMSqueeze-- through an error compensated method compressing gradients. The proposed algorithm achieves a similar convergence efficiency to Adam in term of epochs, but significantly reduces the running time per epoch. In terms of end-to-end performance (including the full-precision pre-condition step), APMSqueeze is able to provide {sometimes by up to $2-10 imes$ speed-up depending on network bandwidth.} We also conduct theoretical analysis on the convergence and efficiency.

Motivation & Objective

• Address the communication bottleneck in distributed training of large models like BERT and ResNet-18, where Adam is essential but incompatible with standard compression.
• Overcome the incompatibility of gradient compression with Adam’s non-linear update rule, which breaks convergence in existing methods.
• Develop a communication-efficient algorithm that maintains Adam’s convergence rate and final model accuracy while drastically reducing data transmission.
• Enable scalable, high-performance distributed training of state-of-the-art models under limited network bandwidth using a novel preconditioning and compression pipeline.

Proposed method

• Precondition momentum SGD using a few epochs of full-precision distributed Adam optimization before switching to compressed training.
• Apply error-compensensated gradient compression to the momentum update direction after the preconditioning phase, ensuring convergence is preserved.
• Use a randomized compressing operator $\bm{C}_{\omega}(\cdot)$ that supports both quantization and sparsification, with error compensation to correct for compression loss.
• Maintain the momentum vector $\bm{m}_t$ and adaptive learning rate $\bm{v}_t$ from Adam during the preconditioning phase, then freeze $\bm{v}_t$ and compress $\bm{m}_t$ during the main training phase.
• Introduce a hybrid training strategy: first 15% of training uses full-precision Adam, followed by 85% of compressed momentum SGD with error compensation.
• Theoretical analysis shows that the compressed algorithm achieves the same asymptotic convergence rate as the uncompressed version under standard assumptions.

Experimental results

Research questions

• RQ1 Can gradient compression be effectively applied to Adam-based training without degrading convergence or model accuracy?
• RQ2 Does preconditioning momentum SGD with a few epochs of Adam enable stable and efficient communication compression in subsequent training stages?
• RQ3 What is the theoretical convergence guarantee for a compressed momentum SGD algorithm when preconditioned by Adam?
• RQ4 How much communication overhead reduction and end-to-end speedup can be achieved in large-scale models like BERT and ResNet-18 using this method?
• RQ5 Is the proposed method robust across different network bandwidths and compression techniques?

Key findings

• APMSqueeze reduces communication by up to 97% (equivalent to 32× compression) while maintaining the same convergence behavior and final accuracy as full-precision Adam on BERT-Base and BERT-Large.
• On 128 GPUs, APMSqueeze achieves up to 8× speedup in per-epoch runtime and up to 10× end-to-end speedup under 1Gbps network conditions.
• With 10Gbps bandwidth, APMSqueeze achieves a 2× end-to-end speedup, demonstrating strong scalability across network conditions.
• The algorithm maintains the same training loss and test accuracy as Adam across all evaluated tasks, including BERT-Base, BERT-Large, and ResNet-18 on CIFAR-10.
• Theoretical analysis confirms that APMSqueeze achieves the same asymptotic convergence rate as the uncompressed baseline, with linear speedup in the number of workers.
• APMSqueeze is the first communication-efficient distributed algorithm that successfully supports training models as complex as BERT using gradient compression with Adam-like optimizers.