通过 0/1 Adam 最大化大规模训练的通信效率
机器学习
2022-05-24 v3
摘要
1-bit 梯度压缩与本地步长是分布式 SGD 中实现显著通信削减的两类代表性技术。然而，在基于 Adam 的大模型预训练（如 BERT 和 GPT）上，其收益仍是一个开放问题。本文中，我们证明了 Adam 中的非线性即使在单独应用 1-bit 压缩或本地步长时也会导致缓慢收敛。为缓解此局限，我们提出 0/1 Adam，通过使用陈旧估计与线性相关性近似其优化器状态，从而线性化每个 Adam 步。0/1 Adam 执行类 Adam 步以保留自适应性，而其线性使其能同时利用 1-bit 压缩与本地步长以实现挂钟时间加速。我们给出了 0/1 Adam 在光滑非凸目标上的收敛保证。在 BERT-Base、BERT-Large、GPT-2 预训练与 ImageNet 等多种大规模基准上，我们在最多 128 块 GPU 上证明，与最先进的基线 1-bit Adam 相比，0/1 Adam 能够减少多达 87% 的数据量、54% 的通信轮数，并实现多达 2 更高的训练吞吐量与端到端训练时间缩减；同时在 GLUE 数据集与 ImageNet 验证集上享有相同的统计收敛速度与末端任务模型精度。
引用
@article{arxiv.2202.06009,
title = {Maximizing Communication Efficiency for Large-scale Training via 0/1 Adam},
author = {Yucheng Lu and Conglong Li and Minjia Zhang and Christopher De Sa and Yuxiong He},
journal= {arXiv preprint arXiv:2202.06009},
year = {2022}
}