單步梯度延遲並非大規模非同步管線平行大型語言模型預訓練的障礙

摘要

現代大規模大型語言模型預訓練受益於使用管線並行技術；然而，同步實作會使 GPU 在管線氣泡期間閒置，浪費運算資源。非同步管線並行消除了這些氣泡，以梯度陳舊性為代價最大化吞吐量。在非同步排程中，PipeDream-2BW 尤其引人注目：與原始 PipeDream 排程不同，它確保無論管線深度如何，都保持恆定的單步梯度延遲。然而，由於普遍認為在陳舊性下進行最佳化本質上不穩定，其採用仍然有限。在本工作中，我們挑戰此一假設，證明了單步延遲下的性能下降高度依賴於最佳化器的選擇，而非固有的限制。我們提供了首次全面的實證分析，顯示雖然 PipeDream-2BW 提出時的主流最佳化器 AdamW 確實遭受嚴重性能下降，但近期方法如 Muon 在單步延遲下展現出強健的穩健性。我們引入了一種與最佳化器無關、靈感來自誤差反饋的修正方法，以進一步減輕延遲效應。我們提供了支持性的理論分析，證明了 Muon 在有無此修正下的收斂性。在參數量高達 100 億的模型上的廣泛評估證實了我們的策略彌合了與同步訓練的性能差距，突顯了大規模非同步管線並行的實際潛力。

English

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.