提出一种高效分布式训练Transformer的符号动量方法，减少通信开销。
Distributed Sign Momentum with Local Steps for Training Transformers
- 局部用任意优化器训练，全局用符号动量聚合梯度差。
- 在非凸光滑函数下达到最优 $O(1/T^{1/4})$ 的 $oldsymbol{ ext{l}_1}$ 梯度收敛率。
- 适合大规模Transformer从头预训练，显著优于现有分布式多步方法。
预训练Transformer模型资源消耗大，近期研究表明符号动量是训练大规模深度学习模型（尤其是Transformer）的有效技术。然而其在分布式训练中的应用仍不充分。本文提出一种新型通信高效的分布式符号动量方法，支持多个本地步骤，适用于无法每步通信的场景。该方法允许各类基础优化器用于本地步骤，并在全局步骤中使用符号动量，动量由本地步骤累积差异生成。针对通用基础优化器，通过将符号算子近似为随机版本（期望下为连续类比），给出一般收敛性分析，特定实例达到 $O(1/ ext{√}T)$ 收敛率；当本地步骤为随机梯度下降时，证明在非凸光滑代价函数下 $oldsymbol{ ext{l}_1}$ 梯度范数可实现最优 $O(1/T^{1/4})$ 收敛率。在不同规模GPT-2模型的从头预训练上进行了广泛评估，实验结果表明，相比其他多本地步分布式方法有显著提升。
原文摘要 · Abstract (English)
Pre-training Transformer models is resource-intensive, and recent studies have shown that sign momentum is an efficient technique for training large-scale deep learning models, particularly Transformers. However, its application in distributed training remains underexplored. This paper investigates a novel communication-efficient distributed sign momentum method with multiple local steps, to cope with the scenarios where communicating at every step is prohibitive. Our proposed method allows for a broad class of base optimizers for local steps, and uses sign momentum in the global step, where momentum is generated from differences accumulated during local steps. For generic base optimizers, by approximating the sign operator with a randomized version that acts as a continuous analog in expectation, we present a general convergence analysis, which specializes to an $O(1/\sqrt{T})$ rate for a particular instance. When local step is stochastic gradient descent, we show an optimal $O(1/T^{1/4})$ rate in terms of $\ell_1$ gradient norm for nonconvex smooth cost functions. We extensively evaluate our method on the pre-training of various sized GPT-2 models from scratch, and the empirical results show significant improvement compared to other distributed methods with multiple local steps.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。