流水线并行与微批次调度
流水线并行将模型按层切分到不同 GPU 上,每个 GPU 负责一部分层的计算。微批次调度策略决定了流水线的效率和显存占用。
基本原理
将 N 层模型分配到 P 个 GPU:
GPU 0: Layer 0 ~ Layer N/P-1
GPU 1: Layer N/P ~ Layer 2N/P-1
...
GPU P-1: Layer (P-1)N/P ~ Layer N-1每个微批次通过流水线依次在各 GPU 上计算。
调度策略
GPipe 调度
先执行所有微批次的前向传播,再执行所有反向传播:
时间 → GPU0 GPU1 GPU2 GPU3
t1 F0,0
t2 F0,1 F0,0
t3 F0,2 F0,1 F0,0
t4 F0,3 F0,2 F0,1 F0,0
t5 F0,3 F0,2 F0,1
...
t9 B0,3
t10 B0,2 B0,3
...- 优点:实现简单
- 缺点:显存峰值高(需存储所有微批次的激活值)
1F1B 调度
交替执行前向和反向传播,减少显存峰值:
python
class OneForwardOneBackward:
"""1F1B 流水线调度"""
def schedule(self, num_stages, num_micro_batches):
schedule = []
# 预热阶段
for i in range(num_stages):
schedule.append(('F', i))
# 稳态:1前向1反向
for i in range(num_micro_batches - num_stages):
schedule.append(('F', num_stages + i))
schedule.append(('B', i))
# 冷却阶段
for i in range(num_stages):
schedule.append(('B', num_micro_batches - num_stages + i))
return scheduleInterleaved 1F1B
每个 GPU 负责不连续的层块(多个虚拟阶段),减少气泡:
GPU0: Layer 0-9, Layer 40-49 (虚拟阶段 0, 4)
GPU1: Layer 10-19, Layer 50-59 (虚拟阶段 1, 5)
...| 策略 | 气泡率 | 显存峰值 | 通信量 |
|---|---|---|---|
| GPipe | (P-1)/(M+P-1) | O(M×激活) | P2P |
| 1F1B | (P-1)/(M+P-1) | O(1×激活) | P2P |
| Interleaved | (P-1)/(M×V+P-1) | O(V×激活) | 2×P2P |
其中 P = 阶段数,M = 微批次数,V = 虚拟阶段数。
微批次数量
微批次数 M 越大,气泡率越低。推荐 M ≥ 4P,使气泡率低于 25%。但 M 过大增加延迟,需在吞吐和延迟间平衡。
显存与气泡权衡
Interleaved 1F1B 减少气泡但增加通信量和显存峰值。在 NVLink 互联环境中推荐使用(通信开销小);跨节点场景建议标准 1F1B。