[论文解读] Maximizing Communication Efficiency for Large-scale Training via 0/1 Adam

Yucheng Lu, Conglong Li+3 more|arXiv (Cornell University)|Feb 12, 2022

Stochastic Gradient Optimization Techniques被引用 4

一句话总结

本文提出 0/1 Adam，一种面向大规模基于 Adam 训练的通信高效优化方法，通过使用过时的方差和动量近似线性化 Adam 更新，实现 1-bit 梯度压缩与本地步长的并行使用。该方法在 128 个 GPU 上实现了高达 87% 的通信量减少、54% 的通信轮次减少，以及 2 倍的训练吞吐量提升，且与 1-bit Adam 相比模型精度无损失。

ABSTRACT

1-bit gradient compression and local steps are two representative techniques that enable drastic communication reduction in distributed SGD. Their benefits, however, remain an open question on Adam-based large model pre-training (e.g. BERT and GPT). In this paper, we demonstrate the non-linearity in Adam causes slow convergence even when 1-bit compression or local steps are individually applied. To alleviate this limitation, we propose 0/1 Adam that linearizes each Adam step via approximating its optimizer states using their stale estimates and linear correlation. 0/1 Adam performs an Adam-like step to preserve the adaptivity, while its linearity allows utilizing 1-bit compression and local steps simultaneously for wall-clock time speed up. We provide convergence guarantee for 0/1 Adam on smooth non-convex objectives. On various large-scale benchmarks such as BERT-Base, BERT-Large, GPT-2 pre-training and ImageNet, we demonstrate on up to 128 GPUs that 0/1 Adam is able to reduce up to 87% of data volume, 54% of communication rounds, and achieve up to 2$ imes$ higher training throughput and end-to-end training time reduction compared to the state-of-the-art baseline 1-bit Adam; while enjoying the same statistical convergence speed and end task model accuracy on GLUE dataset and ImageNet validation set.

研究动机与目标

• 解决在基于 Adam 的大规模模型训练中应用激进的 1-bit 梯度压缩与本地步长的挑战，其中优化器状态中的非线性会阻碍通信效率。
• 克服现有方法（如 1-bit Adam）的局限性，后者需要完整的高精度预热阶段，且难以与本地步长轻松集成。
• 通过对方差和动量的过时估计线性化更新过程，实现在 Adam 中同时使用 1-bit 压缩与本地步长。
• 为平滑非凸目标函数下的 0/1 Adam 提供理论收敛保证。
• 在 BERT、GPT-2 和 ImageNet 基准测试中实现最先进的训练加速，且不牺牲模型精度。

提出的方法

• 提出 0/1 Adam 作为统一的单阶段优化方法，通过自适应冻结优化器状态的方差，避免 1-bit Adam 所需的两阶段预热。
• 使用对方差和动量的过时估计来近似动量和参数更新，使跨工作节点的 Adam 步骤可实现线性近似。
• 在方差冻结的假设下，利用动量对梯度的线性依赖关系，仅对动量值进行 1-bit 压缩，从而大幅降低通信量。
• 通过允许工作节点使用相同的过时方差和动量估计执行多次本地更新，实现本地步长，降低同步频率。
• 利用相邻优化器状态更新之间的线性相关性，界定近似误差并维持收敛稳定性。
• 提出一种收敛性分析框架，界定线性近似与量化带来的误差，证明在平滑非凸目标函数下 0/1 Adam 的收敛性。

实验结果

研究问题

• RQ1 能否在不损害收敛性或精度的前提下，有效结合基于 Adam 的训练中的 1-bit 梯度压缩与本地步长？
• RQ2 Adam 更新规则中的非线性是否从根本上阻止了在大规模预训练中使用激进压缩与本地步长？
• RQ3 能否通过统一的单阶段训练过程替代 1-bit Adam 的两阶段预热，同时保持性能？
• RQ4 在 1-bit 压缩与本地更新下，对线性化版 Adam 能提供何种理论保证？
• RQ5 在使用 0/1 Adam 的大规模 BERT、GPT-2 和 ImageNet 训练中，通信量与通信轮次最多可减少多少？

主要发现

• 在最多 128 个 GPU 上，0/1 Adam 相较于 1-bit Adam 最多可将通信量减少 87%，显著降低数据传输成本。
• 通信轮次最多减少 54%，通过减少同步开销加速训练。
• 0/1 Adam 在 BERT-Base、BERT-Large、GPT-2 和 ImageNet 基准测试中，实现最高 2 倍的训练吞吐量提升与端到端训练时间缩短。
• 该方法在 GLUE 和 ImageNet 验证集上的统计收敛速度与最终任务精度与 1-bit Adam 保持一致。
• 收敛性分析证明，当近似误差与量化误差有界时，0/1 Adam 可在平滑非凸目标函数上实现收敛。
• 0/1 Adam 优化器与训练脚本已在 DeepSpeed 中开源，支持更广泛的应用与可复现性。