改进异步流水线训练中的梯度延迟问题,提升大模型训练效率。
Nesterov Method for Asynchronous Pipeline Parallel Optimization
- 基于Nesterov加速梯度思想,调整前瞻步骤以缓解梯度延迟。
- 理论证明在固定延迟下可实现次线性收敛速度。
- 在10亿参数语言模型上超越现有异步方法,甚至优于同步基线。
流水线并行(Pipeline Parallelism, PP)通过将模型分割为多个阶段,使小型互联设备能够训练大型神经网络。为最大化流水线利用率,异步优化因其天然实现100%利用率而具有吸引力。然而,其固有挑战在于权重与梯度不再同步,导致梯度出现滞后(或延迟)。为缓解此问题,本文提出一种适用于异步优化的Nesterov加速梯度(NAG)变体。具体地,我们修改了NAG中的前瞻步骤,以有效应对梯度滞后的现象。理论上证明,在梯度存在固定延迟的情况下,该方法仍可实现次线性收敛速率。在使用最多10亿参数的仅解码器架构的大规模语言建模任务上的实验表明,该方法显著优于现有异步方法,甚至超越了同步基线性能。
原文摘要 · Abstract (English)
Pipeline Parallelism (PP) enables large neural network training on small, interconnected devices by splitting the model into multiple stages. To maximize pipeline utilization, asynchronous optimization is appealing as it offers 100% pipeline utilization by construction. However, it is inherently challenging as the weights and gradients are no longer synchronized, leading to stale (or delayed) gradients. To alleviate this, we introduce a variant of Nesterov Accelerated Gradient (NAG) for asynchronous optimization in PP. Specifically, we modify the look-ahead step in NAG to effectively address the staleness in gradients. We theoretically prove that our approach converges at a sublinear rate in the presence of fixed delay in gradients. Our experiments on large-scale language modelling tasks using decoder-only architectures with up to 1B parameters, demonstrate that our approach significantly outperforms existing asynchronous methods, even surpassing the synchronous baseline.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。