Files
2026-05-16 17:16:51 +08:00

4.5 KiB
Raw Permalink Blame History

流水线并行Pipeline Parallelism

概述

流水线并行Pipeline ParallelismPP是将模型按层垂直切分到多个GPU上每个GPU负责模型的一部分连续层前向传播和反向传播。与张量并行不同流水线并行关注的是层与层之间的划分而非单层内部的计算并行。

Stage划分

基本概念

将深度学习模型按层边界划分为多个Stage

模型层次: [Embedding → Layer0 → Layer1 → ... → LayerN → Output]
                ↓          ↓
           Stage 0    Stage 1    ...    Stage P-1
           GPU 0      GPU 1             GPU P-1

划分策略

  1. 均匀切分每Stage层数相等如24层模型切4 Stage每Stage 6层
  2. 非均匀切分根据计算量和内存需求动态分配如Embedding层单独StageTransformer层均匀分配

设备映射

假设有P个GPU模型有L层

  • Stage i 负责层 i * L/P(i+1) * L/P - 1
  • Stage 0 在GPU 0上执行Embedding和初始层
  • Stage P-1 在GPU P-1上执行输出层和Loss计算

调度方法

1. GPipeGoogle Pipe理

特点将mini-batch切分为多个micro-batch串行流水执行。

调度流程

Micro-batch 1: [F0, F1, F2, F3, F4] → [B4, B3, B2, B1, B0]
Micro-batch 2:     [F0, F1, F2, F3, F4] → [B4, B3, B2, B1, B0]
...

关键参数

  • num_micro_batches将一个batch分为多少个micro-batch
  • accumulate_steps:梯度累积步数

优点

  • 显存利用率高通过micro-batch轮转
  • 易于实现

缺点

  • 流水线启动和结束时存在大量空闲(气泡)

气泡率分析

气泡比例 = (P - 1) / (num_micro_batches + P - 1)

当 micro-batch 数量足够多时,气泡占比可忽略。

2. PipeDream微软研究院

特点采用1F1BOne Forward One Backward调度尽量保持各设备忙碌。

调度流程

时间步 1: F0
时间步 2: F1, B0
时间步 3: F2, B1, F0
时间步 4: B2, F1, B0
...

优势

  • 比GPipe更早开始反向传播
  • 更好利用机器资源

变种

  • PipeDream-1F1B标准1F1B
  • PipeDream-2F1B每个backward对应两个forward提高吞吐量
  • PipeDream-Flush结合GPipe的气泡消除策略

3. 其他调度策略

策略 特点 适用场景
Interleaved 1F1B 每个设备处理多个stage 多stage场景
Chimney 针对不均衡计算量优化 异构硬件
Megatron-Balance 结合TP的调度 超大规模模型

流水线气泡问题

问题描述

在流水线启动阶段第一个micro-batch前向传播和结束阶段最后一个micro-batch反向传播完成存在设备空闲。

理想状态: [GPU0 busy] [GPU1 busy] [GPU2 busy] [GPU3 busy]
实际状态: [GPU0 busy] [idle    ] [idle    ] [idle    ]  ← 启动气泡
         [GPU0 busy] [GPU1 busy] [GPU2 busy] [GPU3 busy]
         [idle    ] [GPU0 busy] [GPU1 busy] [GPU2 busy]  ← 交错气泡
         ...

解决方案

  1. 增加micro-batch数量:气泡占比 = (P-1)/(N+M-1)其中N为micro-batch数P为stage数M为反向阶段深度
  2. 前向backward重叠:在等待反向传播时继续前向
  3. 异步流水线:允许非严格顺序执行
  4. PipeDream-Flush定期flush流水线减少气泡

气泡率计算

对于GPipe假设有P个stage划分为M个micro-batch

气泡率 = (P - 1) / (M + P - 1)

当 M=32, P=4 时,气泡率约 8.6%,可接受。

实现要点

显存管理

  • ** Activation Recomputation**:通过重新计算激活值而非存储来节省显存
  • 梯度累积不跨pipeline并行同步梯度而是本地累积后同步
  • 模型分片每个stage只保存对应层的参数

通信模式

  • 跨stage通信前向时将激活值传给下一stage反向时将梯度回传
  • 集合通信AllReduce用于梯度同步在1F1B中体现

负载均衡

  • 计算量评估Transformer层计算量约 4 * batch_size * seq_len * hidden_size * intermediate_size
  • 根据实测调整stage边界避免木桶效应

框架支持

框架 流水线策略 备注
PyTorch NaivePipeline, 1F1B, Interleaved 1F1B 需手动处理调度
DeepSpeed pipe schedule支持GPipe, 1F1B 自动优化
Megatron-LM schedule支持Interleaved 结合Tensor并行
NVIDIA Nemo 自动pipeline平衡 企业级支持