147 lines
4.5 KiB
Markdown
147 lines
4.5 KiB
Markdown
# 流水线并行（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
|
||
```
|
||
|
||
### 划分策略
|
||
|
||
1. **均匀切分**：每Stage层数相等（如24层模型切4 Stage，每Stage 6层）
|
||
2. **非均匀切分**：根据计算量和内存需求动态分配（如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-batch
|
||
- `accumulate_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] ← 交错气泡
|
||
...
|
||
```
|
||
|
||
### 解决方案
|
||
|
||
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平衡 | 企业级支持 | |