4.5 KiB
4.5 KiB
流水线并行(Pipeline Parallelism)
概述
流水线并行(Pipeline Parallelism,PP)是将模型按层垂直切分到多个GPU上,每个GPU负责模型的一部分连续层(前向传播和反向传播)。与张量并行不同,流水线并行关注的是层与层之间的划分,而非单层内部的计算并行。
Stage划分
基本概念
将深度学习模型按层边界划分为多个Stage:
模型层次: [Embedding → Layer0 → Layer1 → ... → LayerN → Output]
↓ ↓
Stage 0 Stage 1 ... Stage P-1
GPU 0 GPU 1 GPU P-1
划分策略
- 均匀切分:每Stage层数相等(如24层模型切4 Stage,每Stage 6层)
- 非均匀切分:根据计算量和内存需求动态分配(如Embedding层单独Stage,Transformer层均匀分配)
设备映射
假设有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. GPipe(Google 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-batchaccumulate_steps:梯度累积步数
优点:
- 显存利用率高(通过micro-batch轮转)
- 易于实现
缺点:
- 流水线启动和结束时存在大量空闲(气泡)
气泡率分析:
气泡比例 = (P - 1) / (num_micro_batches + P - 1)
当 micro-batch 数量足够多时,气泡占比可忽略。
2. PipeDream(微软研究院)
特点:采用1F1B(One 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] ← 交错气泡
...
解决方案
- 增加micro-batch数量:气泡占比 = (P-1)/(N+M-1),其中N为micro-batch数,P为stage数,M为反向阶段深度
- 前向backward重叠:在等待反向传播时继续前向
- 异步流水线:允许非严格顺序执行
- 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平衡 | 企业级支持 |