[论文解读] Zero Bubble Pipeline Parallelism

Penghui Qi, Xinyi Wan+2 more|arXiv (Cornell University)|Nov 30, 2023

Parallel Computing and Optimization Techniques被引用 4

一句话总结

本文提出零气泡流水线并行（ZB），一种新颖的调度策略，通过将反向计算解耦为输入梯度（B）和参数梯度（W）两部分，在同步训练中实现近乎零的流水线气泡。通过在更细粒度上重新组织计算图并引入新调度策略（ZB-V），该方法相比1F1B将气泡率降低高达75%，在相似内存限制下提升吞吐量23%，在内存放宽时提升高达31%。

ABSTRACT

Pipeline parallelism is one of the key components for large-scale distributed training, yet its efficiency suffers from pipeline bubbles which were deemed inevitable. In this work, we introduce a scheduling strategy that, to our knowledge, is the first to successfully achieve zero pipeline bubbles under synchronous training semantics. The key idea behind this improvement is to split the backward computation into two parts, one that computes gradient for the input and another that computes for the parameters. Based on this idea, we handcraft novel pipeline schedules that significantly outperform the baseline methods. We further develop an algorithm that automatically finds an optimal schedule based on specific model configuration and memory limit. Additionally, to truly achieve zero bubble, we introduce a novel technique to bypass synchronizations during the optimizer step. Experimental evaluations show that our method outperforms the 1F1B schedule up to 23% in throughput under a similar memory limit. This number can be further pushed to 31% when the memory constraint is relaxed. We believe our results mark a major step forward in harnessing the true potential of pipeline parallelism. We open sourced our implementation based on the popular Megatron-LM repository on https://github.com/sail-sg/zero-bubble-pipeline-parallelism.

研究动机与目标

• 解决同步训练语义下流水线并行中持续存在的流水线气泡低效问题。
• 克服现有调度策略（如1F1B）的局限性，尽管有所改进，但仍存在非零气泡率。
• 设计一种调度框架，在尊重内存约束的同时最小化气泡率，并支持高效的大规模模型训练。
• 在不牺牲优化语义或超出基线方法内存使用范围的前提下，实现零气泡训练。
• 开发一种自动化算法，根据模型配置和内存限制生成最优调度。

提出的方法

• 将反向传播分解为两个独立组件：相对于输入的梯度计算（B）和相对于参数的梯度计算（W），打破传统单体反向函数的结构。
• 引入一种新型流水线调度策略ZB-V，通过重新排序B和W的计算顺序，消除导致气泡的依赖关系，在前向、反向和参数梯度时间相等时实现零气泡。
• 在优化器步骤期间设计一种新颖的同步绕过技术，消除空闲时间，进一步减少气泡。
• 开发一种自动化调度算法，在给定内存和模型约束下寻找最优调度，最小化气泡率。
• 采用函数级别的计算图表示（F, B, W），而非层级别的抽象，以支持更细粒度的调度与优化。
• 基于Megatron-LM实现并开源该解决方案，支持与现有大规模训练框架的集成。

实验结果

研究问题

• RQ1 能否通过重构反向计算，实现同步流水线并行中的气泡消除？
• RQ2 何种调度策略可在保持内存效率和与标准训练语义兼容的前提下，实现最低可能的气泡率？
• RQ3 将反向传播分解为输入梯度和参数梯度组件，对流水线效率和内存使用有何影响？
• RQ4 自动化调度在不同模型规模和内存预算下，能在多大程度上降低气泡率？
• RQ5 微批次大小与气泡率之间的权衡如何？是否可优化以实现最大吞吐量？

主要发现

• 对于32阶段、256个微批次的28.3B参数模型，ZB-V的气泡率为0.0274，而1F1B为0.2676，降幅达89.8%。
• 在相似内存约束下，ZB-V相比1F1B吞吐量最高提升23%；在内存放宽时，最高提升达31%。
• 在所有评估的模型规模和配置下，该方法相比1F1B将气泡率降低高达75%。
• ZB-V在与1F1B相同内存占用下实现近乎零气泡，而ZB-H2为达到类似性能需两倍内存。
• 自动化调度算法成功生成了优于启发式基线的最优调度，尤其在内存受限条件下表现更优。
• 增加微批次大小可使大模型吞吐量最高提升8%，但对小模型的增益递减，表明微批次大小与气泡率之间存在权衡。