Zero Bubble Pipeline Parallelism

管道并行是大规模分布式训练的关键组成部分之一，但其效率受到管道气泡的影响，这被认为是不可避免的。在本文中，我们介绍了一种调度策略，据我们所知，这是第一次在同步训练语义下成功实现零管道气泡。这种改进背后的关键思想是将反向计算分为两部分，一部分计算输入的梯度，另一部分计算参数的梯度。基于这个想法，我们手工设计了新颖的管道调度方案，显著优于基线方法。我们进一步开发了一种算法，根据特定的模型配置和内存限制自动找到最优的调度方案。此外，为了真正实现零气泡，我们引入了一种新颖的技术，在优化器步骤中绕过同步。实验评估表明，我们的方法在类似的内存限制下，比1F1B调度方案的吞吐量高出多达23％。当内存约束放宽时，这个数字可以进一步提高到31％。我们相信我们的结果标志着在利用管道并行的真正潜力方面迈出了重要的一步。我们在https://github.com/sail-sg/zero-bubble-pipeline-parallelism上基于流行的Megatron-LM存储库开源了我们的实现。

解决问题：论文旨在解决大规模分布式训练中的管道泡沫问题，提出一种调度策略以实现零管道泡沫。

关键思路：论文将反向计算分为计算输入梯度和计算参数梯度两部分，手工设计了新的管道调度方案，并开发了一种算法自动找到最优调度方案。

其他亮点：论文通过绕过优化器步骤中的同步来实现真正的零泡沫，提出的方法在相似的内存限制下比1F1B调度的吞吐量高出23％，在内存限制更宽松时可以进一步提高至31％。开源了实现代码。

相关研究：该领域的相关研究包括管道并行训练的其他调度策略和泡沫消除技术，如Megatron-LM、GPipe、PipeDream等。