📚 课程目录第4章方法流程

4. 方法流程

在 DeepSpeed 中，流水线并行并不是简单地“把模型按层切开”，而是由配置、划分策略与调度算法三者共同决定训练效率。配置决定流水线的形状与通信方式，划分策略决定每个 stage 的负载是否均衡，调度算法决定 micro-batch 在设备间的推进节奏。三者任一处失衡，都会让 GPU 出现空泡，训练吞吐显著下降。

一、DeepSpeed 流水线配置详解

DeepSpeed 的流水线并行通过 deepspeed.init_inference 或训练配置中的 pipeline 字段启用。核心参数包括流水线切分粒度、stage 数量、micro-batch 数量以及通信后端。下面是一段典型的训练配置片段：

{
  "train_batch_size": 64,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 4,
  "pipeline": {
    "enabled": true,
    "pipe_parallel_size": 4,
    "activation_checkpointing": {
      "partition_activations": true,
      "cpu_checkpointing": false
    }
  },
  "zero_optimization": {
    "stage": 1
  }
}

其中 pipe_parallel_size 表示流水线 stage 数量，train_micro_batch_size_per_gpu 与 gradient_accumulation_steps 共同决定 micro-batch 总数。micro-batch 越多，流水线填充与排空阶段占比越小，但显存占用与调度开销也会上升。配置阶段还需关注 activation_checkpointing，它通过重计算换取显存，使更多 micro-batch 得以驻留，从而提升流水线利用率。

配置检查顺序建议：先确认 stage 数与 GPU 数匹配，再确认 micro-batch 数 ≥ stage 数，最后开启 activation checkpointing 观察显存与吞吐的平衡点。

二、负载均衡策略：均匀划分与动态规划划分

流水线并行的第一道效率关卡是“层划分”。如果每个 stage 分到的计算量差异过大，最慢的 stage 会成为瓶颈，其他 stage 只能等待。

均匀划分：按 Transformer 层数平均分配，例如 24 层模型切 4 个 stage，每个 stage 6 层。实现简单，但忽略了不同层（如 embedding、attention、FFN）计算量与显存差异，容易造成实际负载不均。
动态规划划分：以每层实测计算时间或参数量为代价，用动态规划求解最小化最大 stage 负载的切分点。DeepSpeed 支持通过 profiling 得到每层耗时，再自动搜索最优切分。

下面用一张图对比两种划分策略在 8 层模型、3 个 stage 下的负载分布差异：

动态规划划分将最大 stage 耗时从 160ms 降到约 115ms，流水线空泡明显减少。实际使用中，可先跑一轮 profiling 采集每层前向与反向耗时，再调用 DeepSpeed 的自动划分接口。

三、动态调度算法：1F1B、PipeDream 与 PipeDream-2BW

调度算法决定 micro-batch 如何在 stage 间流动。DeepSpeed 默认采用 1F1B（One Forward One Backward），并在部分场景支持 PipeDream 系列调度。

调度算法	核心思想	显存特点	适用场景
1F1B	前向与反向交替执行，稳态阶段每个 stage 同时只保留一个 micro-batch 的激活	显存占用低且稳定	通用训练，DeepSpeed 默认
PipeDream	前向优先填满流水线，反向延迟执行，允许权重版本不一致	需保存多版本权重，显存高	对吞吐极致敏感、可接受异步更新
PipeDream-2BW	限制权重版本为 2 份，兼顾异步吞吐与收敛稳定性	显存介于 1F1B 与 PipeDream 之间	大模型异步流水线训练

1F1B 的稳态时序可以概括为：填充阶段连续执行前向，直到最后一个 stage 收到第一个 micro-batch；随后进入稳态，每个 stage 执行“一次前向 + 一次反向”；最后排空阶段处理剩余反向。其优势是激活显存与 stage 数无关，只与 micro-batch 数相关。

选择建议：若追求稳定收敛与低显存，优先 1F1B；若流水线 stage 数很多、micro-batch 很大，可评估 PipeDream-2BW 以提升填充效率。

四、案例与常见误区复盘

假设在 8 卡上训练一个 24 层模型，采用 pipe_parallel_size=4、train_micro_batch_size_per_gpu=2、gradient_accumulation_steps=8，则 micro-batch 总数为 16。若均匀划分为每 stage 6 层，但第 1 个 stage 包含 embedding 且计算量偏大，实测该 stage 耗时比其他高 30%，整体吞吐下降约 18%。改用动态规划划分后，将 embedding 与首个 attention 层单独成组，最大 stage 耗时下降，吞吐恢复。

常见误区：只调大 micro-batch 数而不检查 stage 负载，会导致显存溢出或最慢 stage 更慢；忽略 activation checkpointing 与 1F1B 的配合，也会让显存成为瓶颈而非计算。

复盘时建议按以下清单逐项核对：

stage 数是否整除 GPU 数，通信组是否对齐；
micro-batch 总数是否远大于 stage 数，填充排空占比是否可接受；
各 stage 实测耗时差异是否超过 10%，是否需要动态规划重划分；
调度算法与显存预算是否匹配，是否开启 activation checkpointing；
日志中 pipe 相关耗时与空泡比例是否在预期范围。

只有把配置、划分与调度三者联动调优，DeepSpeed 流水线才能真正接近线性加速，而不是被最慢 stage 和空泡拖住整体吞吐。