One-Step Gradient Delay不是大规模异步流水线并行LLM预训练的障碍

One-Step Gradient Delay is Not a Barrier for Large-Scale Asynchronous Pipeline Parallel LLM Pretraining

精选理由

这篇论文用Muon优化器颠覆了'异步流水线并行梯度延迟必然拉胯训练'的老观念,10B参数实验证明性能和同步训练一样好,做大型LLM预训练的同学应该看看。

AI 摘要

该研究挑战了异步流水线并行中梯度陈旧导致不稳定的普遍假设。论文发现,在PipeDream-2BW调度下,一步延迟的退化强烈依赖于优化器选择:AdamW遭受严重退化,而Muon表现出强鲁棒性。作者提出了一种优化器无关的Error Feedback修正进一步缓解延迟影响,并在高达10B参数的模型上验证了与同步训练的性能差距已被消除。理论分析证明了Muon在有无修正下的收敛性。

原文 · arXiv cs.LG

One-Step Gradient Delay is Not a Barrier for Large-Scale Asynchronous Pipeline Parallel LLM Pretraining

Modern large-scale LLM pretraining benefits from utilizing Pipeline Parallelism; however, synchronous implementations leave GPUs idle during pipeline bubbles, wasting computational resources. Asynchronous Pipeline Parallelism eliminates these bubbles, maximizing throughput at the cost of gradient staleness. Among asynchronous schedules, PipeDream-2BW is particularly appealing: unlike the original PipeDream schedule, it ensures a constant one-step gradient delay regardless of pipeline depth. However, its adoption remains limited due to the common belief that optimizing under staleness is fundamentally unstable. In this work, we challenge this assumption, demonstrating that degradation under one-step delay depends strongly on optimizer choice rather than being an intrinsic limitation. We provide the first comprehensive empirical analysis showing that while AdamW, the predominant optimizer at the time when PipeDream-2BW was introduced, indeed suffers from severe degradation, recent methods like Muon exhibit strong robustness under a one-step delay. We introduce an optimizer-agnostic Error Feedback-inspired correction to further mitigate delay effects. We provide supporting theoretical analysis demonstrating convergence for Muon with and without this correction. Extensive evaluation on models up to 10B parameters confirms that our strategies bridge the performance gap with synchronous training, highlighting the practical potential of asynchronous pipeline parallelism at scale.