大模型训练之流水线并行（Pipeline Parallelism）详解

大模型训练之流水线并行Pipeline Parallelism详解前言随着深度学习的飞速发展大语言模型LLM的参数规模从数亿迅速膨胀到数千亿甚至万亿级别。GPT-3 拥有 1750 亿参数PaLM 达到 5400 亿而 Switch Transformer 更是突破了万亿大关。如此庞大的模型单张 GPU 的显存早已无法容纳分布式并行训练成为了必然选择。在主流的并行策略中流水线并行Pipeline Parallelism是一种核心且优雅的方案。本文将从原理到实践全面介绍流水线并行的来龙去脉。一、并行策略全景图在深入流水线并行之前先简要了解大模型训练中的几种主要并行策略并行策略切分对象核心思想数据并行Data Parallelism数据每个设备持有完整模型副本切分训练数据张量并行Tensor Parallelism层内参数将单个算子/层的权重矩阵切分到多个设备流水线并行Pipeline Parallelism层间结构将模型的不同层分配到不同设备专家并行Expert ParallelismMoE 专家将不同专家分配到不同设备实际训练超大模型时通常会将这些策略组合使用形成3D 并行甚至4D/5D 并行。而流水线并行是其中不可或缺的一环。二、为什么需要流水线并行2.1 单卡显存瓶颈以一个拥有 1750 亿参数的模型为例仅模型参数FP16就需要约350 GB显存而当前主流的 NVIDIA A100 GPU 单卡仅有 80 GB 显存。即使不考虑优化器状态、激活值等额外开销也需要至少 5 张卡才能放下模型参数。2.2 数据并行的局限数据并行要求每个设备都持有完整的模型副本当单卡放不下整个模型时数据并行就无能为力了。2.3 张量并行的局限张量并行虽然能拆分单层的参数但它需要在每一层的前向和反向过程中进行高频通信AllReduce因此通常要求设备之间有极高的通信带宽如 NVLink跨节点使用效率较低。流水线并行的优势在于将模型按层切分到不同设备上设备之间只需在层的边界传递激活值前向和梯度反向通信量相对较小且通信模式为点对点P2P非常适合跨节点部署。三、朴素流水线并行Naive Pipeline Parallelism3.1 基本思想最简单的流水线并行思路非常直观将一个LLL层的模型分成KKK个阶段stage将每个阶段分配到一个 GPU 上前向传播时数据从 Stage 0 依次流向 Stage K-1反向传播时梯度从 Stage K-1 依次流回 Stage 0GPU 0: [Layer 0, Layer 1, ..., Layer l₁] GPU 1: [Layer l₁1, ..., Layer l₂] GPU 2: [Layer l₂1, ..., Layer l₃] GPU 3: [Layer l₃1, ..., Layer L]3.2 致命问题流水线气泡Pipeline Bubble朴素方案的最大问题是严重的设备空闲。来看一个 4 个 stage 的例子时间 → GPU 0: [ Forward ][ Idle ][ Backward ][Idle] GPU 1: [ Idle ][ Forward ][ Idle ][ Backward ][ Idle ] GPU 2: [ Idle ][ Forward ][ Idle ][ Backward ][ Idle ] GPU 3: [ Idle ][ Forward ][ Backward ][ Idle ]在任意时刻只有一个 GPU 在工作其余 GPU 全部空闲。如果有KKK个 stage理论设备利用率仅为1K\frac{1}{K}K1​这意味着 4 张卡只有 25% 的利用率——这完全不可接受。这些空闲时间被称为流水线气泡Pipeline Bubble是流水线并行需要解决的核心问题。四、GPipe微批次流水线并行4.1 核心思想GPipeGoogle, 2019是首个系统性解决流水线气泡问题的方案。其核心思想非常简洁将一个 mini-batch 切分成多个 micro-batch依次注入流水线让多个设备尽量同时工作。具体步骤将一个 mini-batchBBB切分为MMM个 micro-batchb1,b2,...,bMb_1, b_2, ..., b_Mb1​,b2​,...,bM​依次将各 micro-batch 注入流水线所有 micro-batch 的前向传播全部完成后再依次执行反向传播所有 micro-batch 的梯度累积后统一更新参数4.2 执行时序图以 4 个 stage、8 个 micro-batch 为例F Forward, B Backward时间步: 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 GPU 0: [F1] [F2] [F3] [F4] [F5] [F6] [F7] [F8] [B8] [B7] [B6] [B5] [B4] [B3] [B2] [B1] GPU 1: [F1] [F2] [F3] [F4] [F5] [F6] [F7] [F8] [B8] [B7] [B6] [B5] [B4] [B3] [B2] [B1] GPU 2: [F1] [F2] [F3] [F4] [F5] [F6] [F7] [F8] [B8] [B7] [B6] [B5] [B4] [B3] [B2] [B1] GPU 3: [F1] [F2] [F3] [F4] [F5] [F6] [F7] [F8] [B8] [B7] [B6] [B5] [B4] [B3] [B2] [B1]4.3 气泡分析GPipe 的气泡比例为Bubble ratioK−1K−1M\text{Bubble ratio} \frac{K - 1}{K - 1 M}Bubble ratioK−1MK−1​其中KKK是 stage 数量MMM是 micro-batch 数量。当M≫KM \gg KM≫K时气泡比例趋近于 0。例如K4,M32K4, M32K4,M32时气泡比例 ≈ 8.6%已经相当理想。4.4 GPipe 的代价GPipe 虽然有效减少了气泡但存在一个问题激活值的显存开销。由于前向传播和反向传播是分离的先做完所有 micro-batch 的前向再做反向所有 micro-batch 的中间激活值都需要保存在显存中显存占用与MMM成正比。解决方案激活重计算Activation Recomputation/CheckpointingGPipe 使用了激活重计算技术前向时只保存每个 stage 边界的激活值反向时需要某一层的激活值时重新前向计算得到。这用计算时间换显存空间通常增加约 33% 的计算量但可以大幅降低显存占用。五、PipeDream 与 1F1B 调度5.1 1F1B 调度策略PipeDreamMicrosoft, 2019提出了一种更高效的调度策略——1F1BOne Forward One Backward在稳定状态下每个设备交替执行一次前向和一次反向而不是像 GPipe 那样集中执行所有前向后再集中执行所有反向。5.2 非交错式 1F1B 时序图以 4 个 stage、8 个 micro-batch 为例时间 → GPU 0: [F1][F2][F3][F4] [B1][F5][B2][F6][B3][F7][B4][F8] [B5][B6][B7][B8] ← warmup → ← steady state (1F1B) → ← cooldown → GPU 1: [F1][F2][F3] [F4][B1][F5][B2][F6][B3][F7][B4][F8] [B5][B6][B7][B8] GPU 2: [F1][F2] [F3][F4][B1][F5][B2][F6][B3][F7][B4][F8] [B5][B6][B7][B8] GPU 3: [F1] [F2][F3][F4][B1][F5][B2][F6][B3][F7][B4][F8] [B5][B6][B7][B8]整个过程分为三个阶段Warmup 阶段流水线逐步填充各 stage 只做前向Steady State稳定阶段每个 stage 交替执行 1 次前向和 1 次反向Cooldown 阶段流水线逐步排空各 stage 只做反向5.3 1F1B 的优势显存节省1F1B 最大的优势在于显存效率。GPipe需要同时保存MMM个 micro-batch 的激活值1F1B每个 stage 在任意时刻最多只需保存KKK个 micro-batch 的激活值KKK为 stage 数这是因为在稳定状态下每完成一次反向就可以释放对应 micro-batch 的激活值形成先进先出的效果。当M≫KM \gg KM≫K时1F1B 的显存节省非常可观。5.4 气泡比例非交错式 1F1B 的气泡比例与 GPipe 相同Bubble ratioK−1K−1M\text{Bubble ratio} \frac{K - 1}{K - 1 M}Bubble ratioK−1MK−1​但其显存效率更高因此在实践中更受欢迎。六、交错式 1F1BInterleaved 1F1B6.1 核心思想Megatron-LM 团队进一步提出了交错式流水线调度。其核心思想是让每个 GPU 不再只持有连续的一段层而是持有多个不连续的模型块model chunk。例如一个 16 层模型分到 4 个 GPU 上非交错式GPU 0 持有 Layer 0-3GPU 1 持有 Layer 4-7…交错式2 chunksGPU 0 持有 Layer 0-1 和 Layer 8-9GPU 1 持有 Layer 2-3 和 Layer 10-11…非交错式: GPU 0: [Layer 0-3] GPU 1: [Layer 4-7] GPU 2: [Layer 8-11] GPU 3: [Layer 12-15] 交错式 (2 chunks per GPU): GPU 0: [Layer 0-1] [Layer 8-9] GPU 1: [Layer 2-3] [Layer 10-11] GPU 2: [Layer 4-5] [Layer 12-13] GPU 3: [Layer 6-7] [Layer 14-15]6.2 为什么能减少气泡交错式的关键洞察在于由于每个 GPU 持有多个 chunk虚拟 stage 数量增多从KKK变为K×vK \times vK×v其中vvv是 chunk 数每个虚拟 stage 的计算量变小单次前向/反向的执行时间缩短。气泡比例变为Bubble ratioK−1v(K−1M)\text{Bubble ratio} \frac{K - 1}{v(K - 1 M)}Bubble ratiov(K−1M)K−1​相比非交错式减少了vvv倍6.3 代价通信量增加。由于虚拟 stage 更多跨设备的点对点通信次数也相应增加。但在通信带宽充足的情况下整体吞吐量仍然有明显提升。七、进阶零气泡流水线并行Zero Bubble Pipeline Parallelism7.1 气泡的本质无论是 GPipe 还是 1F1B气泡都来源于流水线的启动和排空阶段。2024 年Qi 等人提出了Zero Bubble Pipeline Parallelism将气泡比例进一步压缩到接近零。7.2 核心思想拆分反向传播在标准训练中反向传播计算两部分梯度BBB输入梯度计算 loss 对输入激活值的梯度用于传递给上一层WWW权重梯度计算 loss 对本层权重的梯度用于参数更新传统方案将BBB和WWW绑定在一起计算而 Zero Bubble 的关键创新是将BBB输入梯度计算和WWW权重梯度计算解耦独立调度。BBB有严格的依赖关系下一层的BBB依赖上一层的BBB但WWW没有跨 stage 的依赖——它只需要本 stage 的激活值和输出梯度。因此WWW可以被灵活地安排到任何空闲时间槽中填补原本的气泡。7.3 调度示例时间 → GPU 0: [F1][F2][F3][F4][B1][F5][B2][F6][B3][F7][B4][F8][B5][B6][B7][B8][W1][W2]... GPU 1: [F1][F2][F3][B1][F4][B2][F5][B3][F6][W1][B4][F7][W2]... GPU 2: [F1][F2][B1][F3][B2][W1][F4][B3][W2]... GPU 3: [F1][B1][W1][F2][B2][W2][F3][B3][W3]...通过将WWW填入气泡理论上可以实现零气泡在满足内存约束的前提下。7.4 自动搜索最优调度由于手动设计调度方案非常复杂论文提出使用**整数线性规划ILP**来自动搜索最优调度策略在给定的内存约束下最小化总执行时间。八、实践中的关键考量8.1 如何切分模型Stage 划分理想情况下每个 stage 的计算量应尽量均衡否则最慢的 stage 将成为瓶颈增加气泡。常见策略均匀切分每个 stage 分配相同数量的 Transformer 层最常用基于 profiling 的切分考虑 embedding 层、最终分类层等非标准层的计算量差异自动搜索通过性能模型或实际测量来找最优划分8.2 Micro-batch 数量的选择MMM的选择涉及多个 trade-offMMM增大优势劣势✅气泡比例降低每个 micro-batch 更小可能降低计算效率✅更接近理想并行激活值显存占用增加GPipe需要更多梯度累积步骤经验法则M≥4KM \geq 4KM≥4KKKK为 stage 数通常能将气泡控制在较低水平。8.3 通信优化流水线并行中相邻 stage 之间传递的是激活值张量和梯度张量通信模式为点对点Send/Recv。优化手段通信与计算重叠在计算当前 micro-batch 时同时传输上一个 micro-batch 的结果激活值压缩使用 FP16 甚至更低精度传输激活值合理的设备放置将相邻 stage 放在网络拓扑中邻近的设备上8.4 与其他并行策略的组合在实际的大模型训练系统如 Megatron-LM、DeepSpeed中通常采用3D 并行3D 并行 数据并行 × 张量并行 × 流水线并行典型配置示例训练一个千亿参数模型使用 512 张 A100张量并行度 8在单节点的 8 张 GPU 之间NVLink 高带宽流水线并行度 8跨 8 个节点节点间带宽较低P2P 通信量小数据并行度 8512 / (8 × 8) 8 路数据并行这种层次化的设计充分利用了不同并行策略对通信带宽的不同需求。九、主流框架支持框架流水线并行支持调度策略Megatron-LM(NVIDIA)✅ 非常成熟1F1B、交错式 1F1BDeepSpeed(Microsoft)✅ 完善1F1B、与 ZeRO 结合FairScale(Meta)✅ 支持GPipe 风格PyTorch Native✅ torch.distributed.pipelining多种调度ColossalAI✅ 支持1F1B、交错式十、总结与展望方案气泡比例显存效率通信量实现复杂度朴素流水线(K−1)/K(K-1)/K(K−1)/K差低★GPipe(K−1)/(K−1M)(K-1)/(K-1M)(K−1)/(K−1M)较差需激活重计算低★★1F1B(K−1)/(K−1M)(K-1)/(K-1M)(K−1)/(K−1M)好低★★★交错式 1F1B(K−1)/[v(K−1M)](K-1)/[v(K-1M)](K−1)/[v(K−1M)]好中★★★★Zero Bubble≈ 0好低★★★★★流水线并行经过数年的发展已经从简单的模型切分演进为一套精细的调度艺术。从 GPipe 的 micro-batch 思想到 1F1B 的显存优化再到 Zero Bubble 的极致气泡消除每一步都在逼近理论最优。展望未来随着模型规模持续增长和硬件架构的演进如更高带宽的互联、更大的显存流水线并行的调度策略还将继续进化。特别是在以下方向值得关注异构流水线支持不同类型的加速器混合部署弹性流水线支持动态调整 stage 数量和分配与序列并行的深度结合处理超长序列场景自适应调度根据运行时状态动态调整调度策略参考文献Huang et al.,GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019Narayanan et al.,PipeDream: Generalized Pipeline Parallelism for DNN Training, SOSP 2019Narayanan et al.,Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021Qi et al.,Zero Bubble Pipeline Parallelism, ICLR 2024Rajbhandari et al.,ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC 2020本文作者注流水线并行是大模型训练中一个看似简单实则精妙的技术。理解它的关键在于认识到——计算资源的闲置就是最大的浪费而所有调度优化的本质都是在与气泡做斗争。希望本文能帮助读者建立起对流水线并行的完整认知。后记2026年4月17日16点10分于上海在opus 4.6辅助下完成。