论文
一步梯度延迟不是大规模异步流水线并行 LLM 预训练的障碍
One-Step Gradient Delay is Not a Barrier for Large-Scale Asynchronous Pipeline Parallel LLM Pretraining
摘要
现代大规模 LLM 预训练受益于利用管道并行性；然而，同步实现会使 GPU 在管道泡沫期间处于空闲状态，从而浪费计算资源。异步管道并行消除了这些气泡，以梯度陈旧为代价最大化吞吐量。在异步调度中，PipeDream-2BW 特别有吸引力：与原始 PipeDream 调度不同，它确保恒定的一步梯度延迟，无论管道深度如何。然而，由于人们普遍认为在陈旧状态下进行优化从根本上来说是不稳定的，因此它的采用仍然受到限制。在这项工作中，我们挑战了这一假设，证明一步延迟下的退化很大程度上取决于优化器的选择，而不是一种内在的限制。我们提供了第一个全面的实证分析，表明虽然 PipeDream-2BW 引入时的主要优化器 AdamW 确实遭受了严重的退化，但像 Muon 这样的最新方法在一步延迟下表现出很强的鲁棒性。我们引入了与优化器无关的错误反馈修正，以进一步减轻延迟影响。我们提供支持性理论分析，证明经过和不经过此校正的 Muon 的收敛性。对高达 10B 参数的模型的广泛评估证实，我们的策略弥补了同步训练的性能差距，突出了大规模异步管道并行的实际潜力。