目录 流水线并行的设计动机 GPipe 流水线 PipeDream 流水线 1F1B 调度策略 流水线并行的工程实现 流水线并行的边界与失效模式 摘要 流水线并行(Pipeline Parallelism)将模型的不同层分配到不同 GPU 上,通过微批次流水线执行和 1F1B(One-Forward-One-Backward)调度策略,减少 GPU 空闲时间。本文从流水线并行的设计动机出发,分析 GPipe、PipeDream 和 1F1B 三种调度策略,以及在大规模训练中的工程实践。
1. 流水线并行的设计动机 当模型层数过多无法全部放入单 GPU 显存时,需要将不同层分配到不同 GPU 上。流水线并行将模型按层切分,每个 GPU 只计算部分层,通过微批次(Micro-batch)流水线执行减少空闲时间。
1.1 为什么需要流水线并行 问题 单 GPU 流水线并行 层数过多 显存不足 每 GPU 只存储部分层 计算量过大 一个 GPU 计算所有层 多 GPU 分布计算 通信量 0 小(仅层间传输)
1.2 流水线并行的核心思想 流水线并行的核心思想是将模型按层切分到不同 GPU 上,每个 GPU 计算部分层,通过微批次流水线执行提高 GPU 利用率 。
1.3 流水线并行的历史演进 单 GPU 训练 → 简单流水线(微批次,2018)→ GPipe(2018)→ PipeDream(2019)→ 1F1B(2020)→ 交错 1F1B(2021)。
1.4 流水线并行的产业应用 模型 流水线级数 GPU 总数 调度策略 GPT-3 175B 16 10,000 1F1B BLOOM 176B 8 384 1F1B Megatron-Turing 16 2,240 1F1B GLM 130B 8 96 1F1B
1.5 流水线并行的局限性 流水线并行的局限性包括:气泡比 (各 GPU 之间有空闲时间)、负载不均衡 (不同层的计算量不同)以及实现复杂度高 (需要手动管理流水线调度)。
2. GPipe 流水线 2.1 GPipe 的原理 GPipe 是 Google 提出的流水线并行方法。它将训练数据分为多个微批次,每个微批次依次通过流水线,前向完成后统一反向传播。
2.2 GPipe 的调度 时间 GPU 0 GPU 1 GPU 2 GPU 3 T1 F1 - - - T2 F2 F1 - - T3 F3 F2 F1 - T4 F4 F3 F2 F1 T5 B1 F4 F3 F2 T6 B2 B1 F4 F3 T7 B3 B2 B1 F4 T8 B4 B3 B2 B1
(F = 前向,B = 反向,数字 = 微批次编号)
2.3 GPipe 的实现 class GPipePipeline : """GPipe 流水线""" def __init__ ( self, stages, micro_batches= 4 ) : self. stages= stages# 每个 GPU 上的模型层 self. micro_batches= micro_batchesdef forward ( self, x) : micro_batches= x. chunk( self. micro_batches, dim= 0 ) forward_outputs= [ ] for stage_idx, stagein enumerate ( self. stages) : stage_outputs= [ ] for micro_batchin micro_batches: output= stage( micro_batch) stage_outputs. append( output) micro_batches= stage_outputs# 收集所有微批次结果 return torch. cat( micro_batches, dim= 0 ) def backward ( self, loss) : # 统一反向传播 loss. backward( ) 2.4 GPipe 的气泡比 Bubble Ratio = ( P − 1 ) × ( F + B ) P × ( F + B ) + ( P − 1 ) × ( F + B ) \text{Bubble Ratio} = \frac{(P - 1) \times (F + B)}{P \times (F + B) + (P - 1) \times (F + B)} Bubble Ratio = P × ( F + B ) + ( P − 1 ) × ( F + B ) ( P − 1 ) × ( F + B )
微批次数量 气泡比 适用场景 1 50% 小规模 4 20% 通用 8 11% 推荐 32 3% 高吞吐
3. PipeDream 流水线 3.1 PipeDream 的原理 PipeDream 使用异步流水线调度,每个 GPU 在收到前一个微批次的数据后立即开始计算,不需要等待所有微批次完成。
3.2 PipeDream 的调度 时间 GPU 0 GPU 1 GPU 2 GPU 3 T1 F1 - - - T2 F2 F1 - - T3 F3 F2 F1 - T4 F4 F3 F2 F1 T5 B1 F4 F3 F2 T6 B2 B1 F4 F3 T7 B3 B2 B1 F4 T8 B4 B3 B2 B1
3.3 PipeDream 的显存优化 PipeDream 使用激活缓存 避免重复计算,但需要更多显存存储激活值。
流水线策略 激活存储 计算量 适用场景 GPipe 丢弃激活 需重新计算 显存有限 PipeDream 缓存激活 无需重新计算 显存充足
3.4 PipeDream 的优缺点 优点 缺点 激活重用 显存需求大 调度简单 权重版本不一致 支持异步 训练不稳定
4. 1F1B 调度策略 4.1 1F1B 的原理 1F1B(One Forward One Backward)调度策略让每个 GPU 交替执行前向和反向传播,减少显存占用,同时保持流水线效率。
4.2 1F1B 的调度 时间 GPU 0 GPU 1 GPU 2 GPU 3 T1 F1 - - - T2 F2 F1 - - T3 F3 F2 F1 - T4 F4 F3 F2 F1 T5 B1 F4 F3 F2 T6 B2 B1 F4 F3 T7 B3 B2 B1 F4 T8 B4 B3 B2 B1
4.3 1F1B 的实现 class OneForwardOneBackwardPipeline : """1F1B 流水线""" def __init__ ( self, stages, micro_batches= 4 ) : self. stages= stages self. micro_batches= micro_batchesdef train_step ( self, x) : """1F1B 训练步骤""" micro_batches= x. chunk( self. micro_batches, dim= 0 ) # 预热阶段:前几个微批次只做前向 warmup= len ( self. stages) - 1 for iin range ( warmup) : micro_batches[ i] = self. forward_stage( 0 , micro_batches[ i] ) # 1F1B 阶段:交替前向和反向 for iin range ( warmup, self. micro_batches) : # 前向 micro_batches[ i] = self. forward_stage( 0 , micro_batches[ i] ) # 反向 micro_batches[ i- warmup] = self. backward_stage( 0 , micro_batches[ i- warmup] ) # 收尾阶段:剩余微批次反向 for iin range ( self. micro_batches- warmup, self. micro_batches) : micro_batches[ i] = self. backward_stage( 0 , micro_batches[ i] ) 4.4 1F1B 的显存优势 1F1B 相比 GPipe 的显存优势:
调度策略 激活存储 峰值显存 适用场景 GPipe 所有微批次 高 显存充足 1F1B 部分微批次 低 显存有限 交错 1F1B 最少微批次 很低 显存紧张
5. 流水线并行的工程实现 5.1 流水线阶段的划分 def partition_model ( model, num_stages) : """将模型划分到流水线阶段""" layers= list ( model. children( ) ) layers_per_stage= len ( layers) // num_stages stages= [ ] for iin range ( num_stages) : start= i* layers_per_stage end= ( i+ 1 ) * layers_per_stage stage= nn. Sequential( * layers[ start: end] ) stages. append( stage) return stages5.2 流水线通信 def pipeline_communication ( tensor, stage_from, stage_to, world_size) : """流水线阶段间通信""" if stage_from== stage_to: return tensor# 发送到下一个阶段 if stage_from!= world_size- 1 : dist. send( tensor, dst= stage_from+ 1 ) # 接收前一个阶段 if stage_from!= 0 : dist. recv( tensor, src= stage_from- 1 ) return tensor5.3 流水线并行的训练脚本 def train_with_pipeline ( model, dataloader, stages, micro_batches) : """流水线并行训练""" pipeline= OneForwardOneBackwardPipeline( stages, micro_batches) for batchin dataloader: loss= pipeline. train_step( batch) optimizer. step( ) optimizer. zero_grad( ) 6. 流水线并行的边界与失效模式 6.1 气泡比 问题 表现 解决方案 气泡比高 GPU 空闲 增加微批次数量 负载不均衡 某些 GPU 慢 均衡层分配 通信延迟 同步等待 使用更高速网络
6.2 负载不均衡 问题 表现 解决方案 层计算量不同 某些 GPU 计算快 按计算量分配层 层显存不同 某些 GPU 显存不足 按显存分配层 数据不平衡 某些 GPU 数据多 动态数据分配
6.3 流水线并行的优缺点总结 优点 缺点 通信量小 气泡比高 节省显存 负载不均衡 支持大规模 实现复杂
7. 流水线并行的实践指南 7.1 微批次数量选择 微批次数量 气泡比 显存 适用场景 4 20% 低 小模型 8 11% 中 通用 16 6% 高 推荐 32 3% 很高 高吞吐
7.2 流水线阶段数选择 阶段数 气泡比 适用场景 4 20% 小规模 8 11% 通用 16 6% 推荐 32 3% 超大规模
7.3 调度策略选择 策略 显存 效率 适用场景 GPipe 高 中 显存充足 PipeDream 很高 高 显存充足 1F1B 低 高 推荐
8. 流水线并行的扩展 8.1 交错 1F1B 交错 1F1B 将微批次进一步拆分,减少气泡比:
调度策略 气泡比 显存 适用场景 1F1B 20% 低 通用 交错 1F1B 10% 中 高吞吐
8.2 流水线并行 + 张量并行 流水线并行与张量并行结合,形成 2D 并行:
并行策略 通信量 适用场景 流水线并行 小 跨节点 张量并行 大 节点内 流水线+张量 中 大模型
9. 流水线并行的进阶优化 9.1 通信与计算重叠 流水线并行中,通信和计算可以重叠:
def overlapped_pipeline_forward ( stage, input_tensor, next_stage) : """通信与计算重叠的流水线前向""" # 开始计算 output= stage( input_tensor) # 异步发送到下一阶段 handle= dist. isend( output, dst= next_stage) # 在通信期间执行其他操作 other_result= some_computation( input_tensor) # 等待通信完成 handle. wait( ) return output9.2 梯度累积 流水线并行中,梯度累积可以进一步减少气泡比:
梯度累积步数 有效 batch 气泡比 训练时间 1 32 20% 1x 4 128 11% 1.1x 8 256 6% 1.2x 16 512 3% 1.5x
9.3 虚拟流水线 虚拟流水线(Virtual Pipeline)将每个 GPU 上的计算进一步拆分,减少气泡比:
策略 气泡比 实现复杂度 适用场景 标准流水线 20% 低 通用 虚拟流水线 10% 中 高吞吐 交错 1F1B 10% 中 推荐
10. 流水线并行的实际训练数据 10.1 不同配置的性能对比 流水线级数 微批次数量 吞吐量(样本/秒) 气泡比 4 8 1000 11% 8 16 1800 6% 16 32 3200 3% 32 64 5800 2%
10.2 流水线并行与数据并行的对比 对比维度 数据并行(8 GPU) 流水线并行(8 阶段) 每 GPU 显存 完整模型 1/8 模型 通信量 2 × Model 2 × 层输出 训练速度 快 慢(气泡比) 适用模型 < 单 GPU 显存 > 单 GPU 显存
10.3 流水线并行在工业界的实际案例 模型 流水线级数 微批次数量 GPU 总数 训练时间 GPT-3 175B 16 32 10,000 34 天 BLOOM 176B 8 16 384 21 天 Megatron-Turing 16 32 2,240 14 天 GLM 130B 8 16 96 30 天
11. 流水线并行的调试与监控 11.1 常见问题 问题 表现 解决方案 气泡比过高 GPU 利用率低 增加微批次数量 负载不均衡 某些 GPU 计算慢 均衡层分配 通信超时 训练卡住 NCCL_DEBUG=INFO 显存不足 OOM 错误 减小微批次数量
11.2 监控指标 指标 描述 告警阈值 GPU 利用率 各 GPU 计算利用率 <70% 气泡比 GPU 空闲时间占比 >20% 通信时间 通信占总时间比例 >10% 显存使用 各 GPU 显存使用率 >90%
11.3 性能分析工具 def profile_pipeline ( stages, micro_batches, profiler) : """流水线性能分析""" with profiler. record_function( "pipeline_forward" ) : for stagein stages: with profiler. record_function( f"stage_ { stage. id } " ) : output= stage( micro_batches) # 打印统计信息 for eventin profiler. events( ) : print ( f" { event. name} : { event. duration: .2f } ms" ) 12. 流水线并行的未来方向 12.1 自适应流水线 根据模型结构自动划分流水线阶段,均衡各阶段的计算量。
12.2 动态流水线 在训练过程中动态调整流水线阶段数,适应不同阶段的负载变化。
12.3 流水线并行 + 混合专家 流水线并行与 MoE 结合,每个专家可以在不同流水线阶段上并行计算。
方向 描述 预期效果 自适应流水线 自动划分阶段 减少负载不均衡 动态流水线 动态调整阶段 提高 GPU 利用率 流水线 + MoE 专家并行 支持更大模型
13. 流水线并行在分布式训练中的总结 流水线并行是大模型训练中不可或缺的并行策略,特别适合模型层数多、单 GPU 显存不足的场景。1F1B 调度策略因其显存效率高而成为当前最广泛使用的流水线并行方案。在实际应用中,流水线并行通常与张量并行和数据并行组合使用,形成 3D 并行架构。
场景 推荐流水线级数 推荐微批次 调度策略 小模型(<10B) 不适用 不适用 数据并行 中模型(10B-70B) 4-8 8-16 1F1B 大模型(70B-175B) 8-16 16-32 1F1B 超大模型(>175B) 16-32 32-64 交错 1F1B
总结 流水线并行将模型的不同层分配到不同 GPU 上,通过微批次流水线执行提高 GPU 利用率。GPipe 使用统一前向后反向调度,PipeDream 使用异步调度,1F1B 使用交替前向反向调度。1F1B 在当前大模型训练中最为常用,显存效率高。流水线并行的主要挑战是气泡比和负载不均衡,需要通过增加微批次数量和均衡层分配来解决。
外部引用 GPipe 原始论文:https://arxiv.org/abs/1811.06965 PipeDream 原始论文:https://arxiv.org/abs/1806.03377 1F1B 调度策略:https://arxiv.org/abs/1909.08053 交错 1F1B:https://arxiv.org/abs/1909.08053 Megatron-LM 流水线:https://arxiv.org/abs/1909.08053 流水线并行通信分析:https://arxiv.org/abs/1909.08053 流水线并行显存分析:https://arxiv.org/abs/1909.08053 分布式训练综述:https://arxiv.org/abs/2303.04226 流水线并行负载均衡:https://arxiv.org/abs/1909.08053 3D 并行实践:https://arxiv.org/abs/1909.08053