[论文解读] APMSqueeze: A Communication Efficient Adam-Preconditioned Momentum SGD Algorithm

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

Vehicle License Plate Recognition参考文献 35被引用 4

一句话总结

APMSqueeze 是一种通信高效的优化算法，通过在训练初期使用 Adam 对动量 SGD 进行预条件处理，随后应用误差补偿的梯度压缩，将通信量减少高达 97%（即 32 倍数据减少），在 BERT 和 ResNet-18 上实现了高达 10 倍的端到端加速，且不损失收敛性或精度。

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.

研究动机与目标

• 解决 BERT 和 ResNet-18 等大规模模型分布式训练中的通信瓶颈问题，其中 Adam 至关重要但与标准压缩方法不兼容。
• 克服梯度压缩与 Adam 的非线性更新规则之间的不兼容性，现有方法因此破坏了收敛性。
• 开发一种通信高效的算法，在显著减少数据传输的同时，保持 Adam 的收敛速度和最终模型精度。
• 通过新颖的预条件处理与压缩流水线，实现在有限网络带宽下大规模、高性能的 SOTA 模型分布式训练。

提出的方法

• 在切换到压缩训练前，使用少量周期的全精度分布式 Adam 优化对动量 SGD 进行预条件处理。
• 在预条件阶段后，对动量更新方向应用误差补偿的梯度压缩，确保收敛性得以保持。
• 使用一个支持量化和稀疏化的随机压缩算子 $\bm{C}_{\omega}(\cdot)$，并通过误差补偿校正压缩损失。
• 在预条件阶段保持 Adam 的动量向量 $\bm{m}_t$ 和自适应学习率 $\bm{v}_t$，随后在主训练阶段冻结 $\bm{v}_t$ 并压缩 $\bm{m}_t$。
• 引入一种混合训练策略：前 15% 的训练使用全精度 Adam，随后 85% 的训练采用带误差补偿的压缩动量 SGD。
• 理论分析表明，在标准假设下，压缩后的算法能达到与未压缩版本相同的渐近收敛速率。

实验结果

研究问题

• RQ1 能否在不损害收敛性或模型精度的前提下，将梯度压缩有效应用于基于 Adam 的训练？
• RQ2 使用少量 Adam 训练周期对动量 SGD 进行预条件处理，是否能实现在后续训练阶段的稳定且高效的通信压缩？
• RQ3 当使用 Adam 进行预条件处理时，压缩动量 SGD 算法的理论收敛保证是什么？
• RQ4 使用该方法在 BERT 和 ResNet-18 等大规模模型上，能实现多大程度的通信开销减少和端到端加速？
• RQ5 所提出的方法在不同网络带宽和压缩技术下是否具有鲁棒性？

主要发现

• APMSqueeze 在 BERT-Base 和 BERT-Large 上将通信量减少高达 97%（相当于 32 倍压缩），同时保持与全精度 Adam 相同的收敛行为和最终精度。
• 在 128 个 GPU 上，APMSqueeze 在每轮训练时间上实现高达 8 倍加速，在 1Gbps 网络条件下实现高达 10 倍的端到端加速。
• 在 10Gbps 带宽下，APMSqueeze 实现了 2 倍的端到端加速，表明其在各种网络条件下均具有出色的可扩展性。
• 该算法在所有评估任务中均保持与 Adam 相同的训练损失和测试精度，涵盖 BERT-Base、BERT-Large 和 CIFAR-10 上的 ResNet-18。
• 理论分析证实，APMSqueeze 达到了与未压缩基线相同的渐近收敛速率，并在工作节点数量上实现了线性加速。
• APMSqueeze 是首个成功支持使用类似 Adam 的优化器进行梯度压缩训练 BERT 等复杂模型的通信高效分布式算法。