面向分布式训练 Transformer 的局部步骤分布式符号动量

预训练 Transformer 模型计算资源密集型，最近的研究表明符号动量是训练大规模深度学习模型（尤其是 Transformer）的高效技术。然而，其在分布式训练中的应用仍鲜有探索。本文研究了一种 novel communication-efficient distributed sign momentum 方法，通过 multiple local steps 来应对通信在每个步骤上都不可行的场景。我们提出的方法允许多种 base optimizer 用于 local steps，并在 global step 中使用符号动量，其中动量来自在 local steps 期间累积的差异。对于通用的 base optimizer，通过将符号算子近似为一个随机化版本，用作连续类比的期望，我们提出了一般的收敛分析，这一分析在特定实例下可特化为 $O(1/\sqrt{T})$ 的收敛速率。当 local step 使用随机梯度下降时，我们证明了对于非凸光滑代价函数，以 $\ell_1$ 梯度范数为标准，可达到最优的 $O(1/T^{1/4})$ 速率。我们在从头预训练 various 大小的 GPT-2 模型时广泛评估了该方法，实证结果表明相较于其他多 local steps 的分布式方法，本方法具有显著的改进。