Skip to content

流水线并行与微批次调度

流水线并行将模型按层切分到不同 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 schedule

Interleaved 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。

相关资源

最近更新