一步梯度延迟并非大规模异步流水线并行 LLM 预训练的障碍 One-Step Gradient Delay is Not a Barrier for Large-Scale Asynchronous Pipeline Parallel LLM Pretraining
证明用Muon优化器+Error Feedback修正可使异步PP达到同步训练质量
前置知识
Pipeline Parallelism (PP)
一种将深度学习模型垂直分割到多个 GPU 的分布式训练技术。每个 GPU 负责模型的一部分层,数据依次流经各个阶段。在同步 PP 中,需要等待所有阶段完成前向和反向传播才能更新参数,这会导致 pipeline bubbles(空闲时段),浪费计算资源。PipeDream-2BW 是一种异步 PP 调度策略,通过在每个 minibatch 的 M 个 micro-batches 完成反向传播后才更新参数,确保所有阶段都观察到恒定的一步梯度延迟,从而消除 bubbles。
理解 PP 是本文的基础,作者提出异步 PP 的核心动机就是消除同步 PP 中的 bubbles,提高 GPU 利用率。PipeDream-2BW 的恒定延迟特性是本文所有实验和理论分析的前提。
Gradient Staleness
在异步分布式训练中,当参数更新使用的是基于旧参数计算的梯度时,就出现了梯度过期现象。具体来说,如果参数在第 t 步使用的是第 t-1 步计算的梯度 g_{t-1},则存在一步延迟。延迟更新规则为 w_{t+1} = w_t - u_{t-1}(g_{t-1}),其中 u_{t-1} 是优化器更新函数。延迟梯度可能会偏离当前参数点的真实梯度方向,导致优化轨迹偏离最优路径,影响收敛性能和最终模型质量。
梯度过期是异步 PP 的核心挑战,也是本文研究的核心问题。作者证明了一步延迟的退化严重依赖于优化器选择,而非不可避免的限制。
Error Feedback
最初用于补偿量化误差的技术,本文将其应用于梯度延迟补偿。核心思想是跟踪实际应用更新与理想更新之间的差异,并将此误差反馈到后续更新中。在一步延迟设置下,标准延迟训练在第 t 步应用更新 -u_{t-1}(g_{t-1}),而 Error Feedback 修正为 x_{t+1} = x_t - 2u_{t-1}(g_{t-1}) + u_{t-2}(g_{t-2})。这可以理解为当前步应用了两倍的标准延迟更新,然后减去上一步已经应用过的更新,从而补偿缺失的信息。该方法只需要存储一个额外的模型级缓冲区,内存开销极小。
Error Feedback 是本文提出的核心缓解策略,实验证明它可以将同步-异步性能差距缩小 50-90%,是使异步 PP 在 10B MoE 模型上达到同步训练质量的关键技术。
Mixture-of-Experts (MoE)
一种稀疏模型架构,通过路由机制将输入分配给不同的专家子网络。相比密集模型,MoE 在大幅增加参数量的同时,保持推理时的计算量基本不变。然而,MoE 层显著增加了训练时的通信开销,因为需要在不同设备间传输激活值和梯度。这种较低的 compute-to-communication 比例导致数据并行效率下降,使得 PP 在大模型训练中重新受到重视。本文的 10B MoE 实验在 Qwen3-Next 架构上使用 Gated Delta Net 层,总参数 10B,训练时激活参数约 2.8B。
MoE 的通信特性是 PP 重新变得重要的驱动力,也是本文选择 MoE 模型进行大规模验证的原因。作者在 10B MoE 上证明了异步 PP 的实用性,展示了方法在真实场景中的价值。
研究动机
现代大规模 LLM 预训练受益于 Pipeline Parallelism,但同步实现由于需要保持同步参数更新,会在 pipeline 调度中引入 bubbles,在此期间部分 GPU 保持空闲,浪费计算资源。例如,在 16 个 pipeline 阶段的设置下,Ajanthan 等人报告相比同步训练验证损失增加了超过 0.2,这是一个实际训练中常用的规模。虽然异步 Pipeline Parallelism 可以完全消除 bubbles,最大化吞吐量,但代价是梯度延迟。PipeDream-2BW 是特别有吸引力的调度策略,与原始 PipeDream 调度不同,它确保无论 pipeline 深度如何,都保持恒定的一步梯度延迟。然而,其采用仍然有限,因为普遍认为在过期梯度下优化本质上是不稳定的。AdamW 是 PipeDream-2BW 引入时的主导优化器,确实遭受严重退化。
本文的目标是本文的具体目标是挑战梯度延迟本质不稳定的假设,证明一步延迟下的退化强烈依赖于优化器选择而非内在限制。作者提供首次全面的实证分析,展示虽然 AdamW 在一步延迟下确实严重退化,但像 Muon 这样的现代优化器表现出强大的鲁棒性。作者进一步引入一个优化器无关的、受 Error Feedback 启发的修正机制来进一步减轻延迟影响。最终,作者在 10B 参数的 Mixture-of-Experts 模型上验证发现,展示异步 PP 可以在不损失质量的情况下匹配同步训练,突显大规模异步流水线并行主义的实用潜力。
与已有工作不同的是,本文的独特切入角度在于之前的异步 PP 研究主要关注原始 PipeDream 调度,该调度存在可变的梯度延迟,随着 pipeline 深度增加导致严重收敛退化。这些工作认为延迟是主要问题,并尝试通过单一优化器来缓解。本文则指出真正的问题可能是延迟的异质性而非延迟本身。通过使用 PipeDream-2BW 的恒定延迟调度,作者将优化挑战简化为更清晰的设置,即固定一步延迟下的训练。更重要的是,本文首次对 LLM 训练中的延迟优化进行了全面的优化器基准测试,而不仅仅是测试单个优化器,这揭示了优化器选择的重要性,并识别出 Muon 等鲁棒优化器。
核心方法
方法整体思路是系统性地研究异步 PP 中一步梯度延迟的影响,并提出缓解策略。首先,作者使用 PipeDream-2BW 调度确保恒定一步延迟,将问题抽象为延迟更新规则 w_{t+1} = w_t - u_{t-1}(g_{t-1})。然后,作者在 135M 和 360M 模型上对多个现代优化器(AdamW、Muon、SOAP、Nadam、MARS、Adan、LION 等)进行基准测试,评估它们对延迟的鲁棒性。接着,作者研究超参数敏感性,特别是动量系数的影响,发现更高动量能减轻延迟退化。作者还评估了几种缓解策略,包括同步启动、同步冷却和 DC-ASGD,但发现它们效果有限或不一致。作者的核心贡献是提出一个受 Error Feedback 启发的更新级修正:x_{t+1} = x_t - 2u_{t-1}(g_{t-1}) + u_{t-2}(g_{t-2})。作者还提供了延迟 Muon 的理论收敛性分析。最后,作者在 2B 和 10B MoE 模型上扩展验证,证明 Muon + Error Feedback 可以达到与同步训练相同的最终质量。
核心创新点有两方面,一是揭示了梯度延迟的退化强烈依赖于优化器选择而非内在限制,发现 Muon 等现代优化器对一步延迟表现出强大鲁棒性,而 AdamW 则严重退化,二是将 Error Feedback 技术应用于 LLM 训练的梯度延迟补偿,提出优化器无关的更新级修正机制。与原始 PipeDream 的可变延迟不同,PipeDream-2BW 的恒定延迟简化了优化挑战,使作者能够专注于固定一步延迟的设置。与仅关注 Nadam 的先前工作不同,本文的全面基准测试识别出了更鲁棒的优化器。Error Feedback 是一个通用的、优化器无关的修正,适用于大多数优化器。此外,本文还提供了首个 LMO 算法在梯度延迟下的收敛性理论分析。
方法步骤详情
方法步骤的完整描述如下:第一步,设置延迟更新抽象,定义标准异步更新规则 x_{t+1} = x_t - u_{t-1}(g_{t-1}),其中 g_{t-1} 是延迟梯度,u_{t-1} 是优化器更新函数。第二步,基准测试优化器,在 135M 和 360M 模型上训练,使用 FineWeb-Edu 数据集,评估不同优化器在同步和异步设置下的验证损失。第三步,超参数敏感性分析,对每个优化器进行一维扫描,研究学习率、权重衰减、动量系数、Beta2、warmup 长度、梯度裁剪等参数的影响。第四步,评估缓解策略,测试同步启动、同步冷却、DC-ASGD 等基准策略。第五步,实现 Error Feedback 修正,修改更新规则为 x_{t+1} = x_t - 2u_{t-1}(g_{t-1}) + u_{t-2}(g_{t-2}),存储一个额外的模型级缓冲区。第六步,理论分析,对延迟 Muon 提供收敛性保证,证明在标准假设下,延迟设置的噪声界为开根号下 2 乘以 mu 乘以 rho 方乘以 sigma 方加上 8 乘以 L 乘以 eta 的平方。第七步,大规模验证,在 2B 和 10B MoE 模型上训练 200B tokens,验证 Muon 和 Muon + Error Feedback 的性能。
技术新颖性
技术新颖性体现在多个方面,首次对 LLM 训练中的延迟优化进行了全面的优化器基准测试,覆盖了 AdamW、Muon、SOAP、Nadam、MARS、Adan、LION 等多个现代优化器,揭示了动量系数对延迟鲁棒性的关键作用。首次将 Error Feedback 技术应用于 LLM 训练的梯度延迟补偿,提出优化器无关的更新级修正机制,只需一个额外的模型级缓冲区,内存开销极小。首次提供了 LMO 算法在梯度延迟下的收敛性理论分析,为延迟 Muon 提供了收敛保证。首次在 10B MoE 模型上成功演示异步 PP 达到与同步训练相同的最终质量,使用相同的超参数,最终验证损失都是 1.906。此外,作者还比较了 PipeDream-2BW 的恒定延迟与原始 PipeDream 的可变延迟,证明了恒定延迟对稳定性的重要性。
实验结果
核心发现包括多个实验的具体结果。首先,在 360M 模型上,AdamW 在一步延迟下严重退化,最终同步-异步损失差距为 0.278,而 Muon 的差距仅为 0.012。这揭示了优化器选择对延迟鲁棒性的关键影响。其次,在 135M 和 360M 模型的优化器基准测试中,大部分现代优化器保持鲁棒,损失差距在 0.03 以内。具体来说,Muon 的差距为 0.014(135M)和 0.012(360M),Adan 的差距为 0.006 和 0.010,NorMuon 的差距为 0.019 和 0.014。相反,AdamW 和 MARS 严重退化,AdamW 的差距为 0.350 和 0.278,MARS 的差距为 0.469 和 0.282。第三,超参数敏感性分析发现,动量系数对延迟鲁棒性有清晰且一致的负面影响,更高动量持续减少一步延迟引起的损失惩罚。第四,Error Feedback 修正为大多数优化器提供了一致的益处,对于 Muon、AdaMuon、SOAP 和 NorMuon 等鲁棒优化器,它恢复了延迟训练引入的约 50-70% 的退化。对于 MARS 和 AdamW 等最退化的优化器,它恢复了 85-90% 的差距。第五,在 2B MoE 模型上,同步和异步扩展曲线从 50B 到 200B tokens 几乎平行,表明一步延迟不会引入增长的同步-异步差距。第六,在 10B MoE 模型上训练 200B tokens,标准异步 PP 保持高度竞争,最终验证损失为 1.911,相比同步基线 1.906 的差距仅为 0.005。使用 Error Feedback,异步 PP 完全关闭了这一差距,达到相同的最终损失 1.906,使用完全相同的超参数。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 360M 模型语言建模 | 验证损失 | Muon Async: 2.590, Muon Async + EF: 2.583 | Muon Sync: 2.578 | 标准异步差距 0.012,EF 缩小差距 71% |
| 360M 模型语言建模 | 验证损失 | AdamW Async: 2.890, AdamW Async + EF: 2.640 | AdamW Sync: 2.612 | 标准异步差距 0.278,EF 缩小差距 90% |
| 10B MoE 模型预训练 | 验证损失(200B tokens) | Muon Async: 1.911, Muon Async + EF: 1.906 | Muon Sync: 1.906 | 标准异步差距 0.005,EF 完全关闭差距 |
| 135M 模型语言建模 | 验证损失 | Adan Async: 2.902 | Adan Sync: 2.896 | 差距仅 0.006,是最鲁棒的优化器 |
| 360M 模型语言建模 | 验证损失 | MARS Async: 2.897 | MARS Sync: 2.615 | 差距 0.282,严重退化 |
局限与改进
局限性分析包括作者承认的和我们自己的观察。作者承认他们缺乏关于为什么更高动量确实改善延迟训练的完整机制解释。此外,作者的 batch size 和 learning rate 网格仅限于 135M 模型,可能无法完全捕获大规模设置中的行为。作者只对 WPipe 风格的调度进行了有限的附录规模实验。我们自己的观察包括理论分析中的噪声界与同步设置相比多了一个项,这可能在某些情况下变得显著。实验主要集中在 FineWeb-Edu 数据集上,在其他数据集上的行为可能不同。Error Feedback 虽然对大多数优化器有益,但对 Adan 和 NAdam 等优化器略微退化,说明该方法不是普遍有益的。大规模实验仅在 MoE 模型上进行,密集模型的行为可能有所不同。同步-异步差距的评估主要基于验证损失,下游任务的性能虽然被验证为相当,但没有详细的下游任务比较。
独立分析的弱点
独立分析的弱点包括 Error Feedback 需要存储一个额外的模型级缓冲区,虽然作者声称内存开销极小,但在极端内存受限的环境中这可能成为问题。方法主要针对恒定一步延迟,对于可变延迟或更大步数的延迟,效果可能不同。超参数敏感性分析发现 batch size 对延迟训练有强影响,但 batch size 在大规模训练中通常受硬件利用率约束,这限制了在实际场景中调整它来改善异步训练的能力。同步启动策略在自适应优化器中在切换点引入了损失峰值,需要仔细处理。理论分析依赖于某些假设,这些假设在实际 LLM 训练中可能不完全满足。下游任务性能虽然被验证为相当,但详细的比较和更多的基准测试将更有说服力。改进方向可以包括研究 Error Feedback 与 batch size 的交互,探索可变延迟下的自适应策略,扩展到更多数据集和模型架构,以及开发对动量等超参数不敏感的延迟优化方法。
未来方向
未来研究方向包括作者建议在万亿 token 规模上验证这些发现的重要性,因为验证到 200B tokens 支持了异步 PP 在现实预训练运行中的可扩展性。基于本文成果,可以探索的方向包括研究为什么更高动量确实改善延迟训练的完整机制解释,这可能涉及分析动量如何缓解延迟引入的噪声。扩展 batch size 和 learning rate 网格到更大模型,以捕获大规模设置中的行为。研究 WPipe 风格调度的实用性,因为作者提到它可能是实践者的潜在更好替代方案。探索 Error Feedback 与其他缓解策略的组合,如与同步启动或基旋转的结合。研究可变延迟下的自适应策略,以应对异构 worker 运行时。扩展到更多优化器和模型架构。开发自动调优延迟训练超参数的方法,以减少手动调优的负担。研究 Error Feedback 在其他异步优化场景中的应用,如联邦学习或分布式强化学习。最后,探索延迟优化与模型压缩、量化等技术的交互,以在资源受限环境中实现高效训练。
复现评估
复现评估方面,论文声称使用 FineWeb-Edu 数据集,训练了 135M、360M、2B MoE 和 10B MoE 模型。10B MoE 实验使用 Qwen3-Next 架构和 Gated Delta Net 层,训练 200B tokens,全局 batch size 为 4M tokens,峰值学习率为 0.00225。作者没有明确声明是否开源代码和数据,这会影响复现性。算力需求方面,10B MoE 训练 200B tokens 需要大量计算资源,可能需要多个 GPU 集群。较小的 135M 和 360M 模型实验可能在单个或少个 GPU 上可以复现。难度评估为中等偏高,主要因为大规模实验的算力需求。然而,核心方法的实现相对简单,Error Feedback 只需存储一个额外的模型级缓冲区并修改更新规则。作者提供了详细的训练设置和超参数,包括权重衰减 0.1、梯度裁剪 1.0、余弦学习率调度衰减到峰值的 0.1、warmup 持续训练预算的 10%。这些信息有助于复现。总体而言,小规模实验的复现性较好,大规模实验的复现性取决于是否有足够的计算资源。
论文图表