流水线并行的“气泡”怎么消？带你拆解 1F1B 调度算法的精妙之处

导语：为什么流水线并行会产生“气泡”？

在训练超大规模深度学习模型时（如GPT系列），单个GPU的显存往往无法容纳整个模型。我们不得不采用模型并行策略，其中，流水线并行（Pipeline Parallelism, PP）是一种常用的方法，它将模型的不同层切分到不同的GPU上（称为Stage）。

然而，经典的流水线并行调度存在一个核心问题：“气泡”（Bubble），即由于数据依赖导致的GPU空闲时间。当一个Stage完成所有前向传播（Forward, F）后，它必须等待最后一个Stage完成其反向传播（Backward, B）任务才能再次启动，这导致了大量的资源浪费。

解决这个问题的关键在于：1F1B调度算法，它通过巧妙地交错前向和反向计算，几乎消除了“气泡”。

1. 经典流水线调度的低效性

假设我们将大批次（Mini-Batch）切分为 $N$ 个微批次（Micro-Batch），模型被切分为 $M$ 个 Stage (S1到SM)。

在标准同步调度中，一个微批次必须完成所有 $M$ 个 Stage 的前向计算后，才能开始反向计算。

经典调度的执行顺序（以 M=3, N=4 为例）：

在所有微批次的F操作完成后，才开始B操作。问题在于 Stage 1 和 Stage 2 在等待 Stage 3 完成 F 之后，以及在 B 操作开始前，会产生巨大的空闲时间——这就是“气泡”。

气泡的总长度约为 $2 \times (M-1) \times T_{micro}$，其中 $T_{micro}$ 是处理一个微批次的时间。

2. 1F1B 调度算法的精妙之处

1F1B (One Forward, One Backward) 调度是一种异步交错策略，它核心思想是：一旦某个 Stage 完成了一个微批次的前向计算，并且所需的梯度信息可用，它就立即执行下一个微批次的前向计算，或者如果条件满足，就执行反向计算。

算法核心原则：

1. 最大化Overlap： 确保Stage在任何时刻，只要数据依赖满足，就立即执行 F 或 B 任务。
2. 交错执行： 在 Steady State（稳定状态）下，每个 Stage 都在同时执行一个前向操作（F）和一个反向操作（B）。

1F1B 的执行序列（以 M=3, N=4 为例）：

1. 填充阶段 (Fill Phase): 仅执行 F 操作，用于填充流水线。
   $$F_1(1) \to F_2(1) \to F_3(1) \to F_1(2) \to F_2(2) \to F_3(2) \to …$$
2. 稳定阶段 (Steady State): 当第一个微批次的 $F_3$ 完成后，它允许 $B_3(1)$ 开始。此时，Stage 3 刚刚释放，可以立即接收 $F_3(2)$。同时，Stage 2 完成了 $F_2(2)$，它接收 $B_2(1)$。关键在于 F 和 B 任务的交替执行。
3. 排空阶段 (Drain Phase): 当所有 F 操作完成后，仅执行 B 操作，排空流水线。

稳定状态下，Stage 1 保持执行 $F_1$ 直到所有 $N$ 个微批次完成；Stage M 保持执行 $B_M$ 直到所有 $N$ 个微批次完成。中间的 Stage 则交替执行 $F_i$ 和 $B_{i-1}$。

3. Python 模拟 1F1B 执行序列

我们使用一个简化的 Python 脚本来展示 1F1B 调度是如何在 3 个 Stage 和 4 个 Micro-Batch 之间安排任务，从而保持GPU的高利用率。

假设 $T_{F} \approx T_{B}$，我们只需要关注任务的调度顺序。

import collections

# 定义阶段数和微批次数
M = 3  # Stages: S1, S2, S3
N = 4  # Micro-Batches: MB1, MB2, MB3, MB4

# 调度队列：存储待执行的任务 (stage_id, batch_id, type)
schedule = collections.deque()

# 初始化任务列表：首先是所有的 F 任务，然后是所有的 B 任务
all_tasks = []
for i in range(1, N + 1):
    for j in range(1, M + 1):
        all_tasks.append((j, i, 'F'))

for i in range(1, N + 1):
    for j in range(M, 0, -1): # B 任务反向执行
        all_tasks.append((j, i, 'B'))

# 跟踪每个 Stage 上次执行的任务类型 (用于简化依赖检查)
stage_status = {i: None for i in range(1, M + 1)}

# 模拟 1F1B 调度执行过程
print(f"--- 1F1B 调度模拟 (M={M}, N={N}) ---")

# 1. 填充阶段：仅执行 F 直到第一个 B 任务具备条件
fill_tasks = []
for i in range(1, M):
    # F_i(1) 必须完成才能执行 F_{i+1}(1)
    fill_tasks.append(f"F{i}(1)")

# 稳定阶段和排空阶段：开始交错
time_step = 0
executed_tasks = []

# 理想的 1F1B 序列生成器
# 序列长度为 2 * M + N - 2

for k in range(1, 2 * M + N - 1):
    time_step += 1
    tasks_at_step = []

    for i in range(1, M + 1):
        # 计算在当前时间步 k，Stage i 应该执行哪个任务

        # F 任务：(k - i + 1) 是 micro-batch 索引
        mb_f = k - i + 1
        # B 任务：(k - 2*M + i + N) 是 micro-batch 索引 (需要确保 F 已完成)
        mb_b = k - (2*M) + i + N

        current_task = None

        # 优先执行 F 任务（只要 micro-batch 索引在范围内且尚未执行）
        if 1 <= mb_f <= N:
            # 检查 F 任务是否已完成 (简化检查：只要它在序列中出现即可)
            current_task = f"F{i}({mb_f})"

        # 检查是否可以执行 B 任务
        # B 任务启动的条件是：该微批次的 F 已经全部完成，且上游 Stage 的 B 任务已完成 (或该 Stage 是 SM)
        # 简化检查：当 Stage M 开始执行 B 任务后，整个流水线进入稳定状态

        if M <= k <= M + N - 1:
            # 在稳定阶段，Stage M 应该执行 B
            if i == M:
                mb_b_index = k - M + 1
                if 1 <= mb_b_index <= N:
                    current_task = f"B{M}({mb_b_index})"

        if M + 1 <= k <= M + N:
            # Stage M-1 开始执行 B
            if i == M - 1:
                 mb_b_index = k - M
                 if 1 <= mb_b_index <= N:
                    current_task = f"B{M-1}({mb_b_index})"

        # 由于复杂的依赖关系难以用简单的公式表达，我们直接打印理想序列的核心模式

        if i == 1:
            # Stage 1 始终是 F 任务（直到 F 任务耗尽）
            if 1 <= mb_f <= N: tasks_at_step.append(f"S{i}: F{i}({mb_f})")
            elif 1 <= mb_b <= N: tasks_at_step.append(f"S{i}: B{i}({mb_b})")

        elif i == M:
            # Stage M 始终是 B 任务（直到 B 任务耗尽）
            if M <= k <= M + N - 1:
                mb_b_index = k - M + 1
                if 1 <= mb_b_index <= N: tasks_at_step.append(f"S{i}: B{i}({mb_b_index})")
            elif 1 <= mb_f <= N:
                tasks_at_step.append(f"S{i}: F{i}({mb_f})")

        elif i == 2:
            # Stage 2 处于交错状态
            if 1 <= mb_f <= N and k < M + N:
                tasks_at_step.append(f"S{i}: F{i}({mb_f})")

            # B 任务开始晚于 F 任务
            if M+1 <= k <= M + N - 1:
                 mb_b_index = k - M
                 if 1 <= mb_b_index <= N: tasks_at_step.append(f"S{i}: B{i}({mb_b_index})")

    if tasks_at_step:
        print(f"T={time_step}: {' | '.join(tasks_at_step)}")

print("\n...此调度大大减少了 Stages 之间的空闲周期（气泡），尤其是在稳定状态下，Stage M-1 和 Stage M 之间的计算被最大化重叠。\n")

4. 1F1B 带来的性能提升

通过 1F1B 调度（或类似 GPipe、PipeDream 等算法），流水线在进入稳定阶段后，每个 Stage 几乎可以保持 100% 的利用率，从而将总训练时间 $T_{total}$ 从原来的 $T_{sync}$ 显著缩短到 $T_{1F1B}$：

$$T_{sync} \approx N \times M \times T_{micro} \times 2 + \text{气泡时间}
$$

$$T_{1F1B} \approx N \times T_{micro} \times 2 + (2M – 2) \times T_{micro}
$$

其中，$(2M – 2) \times T_{micro}$ 对应于填充和排空阶段的开销。对于大型任务（N 很大），这个开销相对于总训练时间几乎可以忽略不计，从而实现接近理想的线性加速比。

总结

1F1B 调度算法通过将大批量数据切分为微批次，并在流水线 Stage 之间交错执行前向和反向计算，有效地隐藏了数据依赖延迟，消除了流水线中的“气泡”。这是实现高效大规模分布式深度学习训练的关键技术之一。