4. DeepSpeed流水线并行架构
DeepSpeed 的流水线并行并不是单一模块，而是一套由引擎、通信组与调度器共同协作的运行时架构。理解这三者的职责边界，是后续调优气泡率、显存占用和吞吐量的前提。本节从组件分层入手，说明它们如何协同完成一次完整的流水线迭代。
一、核心原理：三层职责分离
DeepSpeed 流水线并行的运行时可以抽象为三个层次：引擎层负责模型切分与前后向执行，通信层负责跨 stage 的激活与梯度传递，调度层负责决定每个 micro-batch 在何时进入哪个 stage。三者通过明确的接口解耦，使得调度策略可以独立替换，而不影响模型切分逻辑。
引擎层由 PipelineEngine 承载，它继承自 DeepSpeed 的 DeepSpeedEngine，在初始化阶段完成三件事：按 partition_method 将模型层划分到各 stage、为每个 stage 构建本地 nn.Module 序列、注册与流水线相关的通信句柄。通信层依赖 PyTorch 的 ProcessGroup，DeepSpeed 会为流水线单独创建一个 pipeline 通信组，与数据并行组、张量并行组相互正交。调度层则由 PipelineScheduler 及其子类实现，常见的有 Schedule1F1B、ScheduleGPipe 和 ScheduleInterleaved1F1B。
二、通信组：正交划分与初始化
DeepSpeed 在初始化时会根据并行配置创建多个正交通信组。流水线组只包含同一数据并行副本内、不同 stage 的 rank；数据并行组只包含同一 stage 内、不同数据副本的 rank。这种正交划分保证了流水线通信与数据并行通信互不干扰，也决定了 send/recv 的 peer 关系。
# 典型 8 卡配置：2 路流水线 × 4 路数据并行
# pipeline_group: [0,4] [1,5] [2,6] [3,7]
# data_group: [0,1,2,3] [4,5,6,7]
# 每个 rank 同时属于一个 pipeline 组和一个 data 组
deepspeed --num_gpus=8 train.py \
--deepspeed_config ds_config.json
在 ds_config.json 中，pipeline 字段的 stages 决定切分粒度，partition_method 决定按层还是按参数切分。通信组一旦建立，调度器便可在其上编排 micro-batch 的收发顺序。
三、调度器角色：从 1F1B 到交错调度
调度器的核心任务是决定每个 micro-batch 在 stage 上的执行顺序，以最小化流水线气泡。以 Schedule1F1B 为例，其状态机在预热阶段连续执行多个前向，填满流水线后进入稳定阶段，交替执行一次前向和一次后向，最后进入冷却阶段排空剩余后向。
| 调度器 | 执行模式 | 气泡特征 | 适用场景 |
|---|---|---|---|
| ScheduleGPipe | 全前向 → 全后向 | 气泡最大，显存峰值高 | 小规模验证、调试 |
| Schedule1F1B | 预热 → 1F1B 交替 → 冷却 | 气泡适中，显存平稳 | 常规训练首选 |
| ScheduleInterleaved1F1B | 虚拟 stage 交错 1F1B | 气泡更小，通信更频繁 | 大模型高吞吐场景 |
四、案例场景：一次 1F1B 迭代的完整链路
假设 4 个 stage、4 个 micro-batch。rank0 在预热阶段依次将 micro-batch 0~2 的前向激活发送给 rank1；rank1 收到后继续向 rank2 传递。当 rank0 完成第 3 个前向后，进入稳定阶段：每收到一个来自 rank1 的反向梯度，就执行一次本地后向，同时发起下一个 micro-batch 的前向。冷却阶段则只处理后向，直到所有 micro-batch 的梯度回传完毕。
关键观察：1F1B 的显存优势来自“前向激活在对应后向完成后立即释放”。若调度器未能及时配对前后向，激活会持续驻留，导致显存峰值接近 GPipe。
五、误区与复盘
- 误区一：认为流水线组越大越好。stage 数增加会放大通信次数与气泡比例，需结合
micro_batch数量权衡。 - 误区二：忽略通信组正交性。若 pipeline 组与 data 组 rank 重叠错误，会出现梯度规约与激活传递互相阻塞。
- 误区三：调度器与切分策略不匹配。按参数切分时，1F1B 的配对逻辑需要额外处理参数边界，否则后向无法正确触发。
排查流水线效率时，优先确认三件事：pipeline_group 的 rank 顺序是否与 stage 编号一致；micro_batch 数量是否大于 stage 数；调度器是否与 partition_method 匹配。这三项决定了气泡率的下限。
综上，DeepSpeed 流水线并行的架构可以概括为：引擎层负责“切什么、怎么算”，通信层负责“传给谁、怎么传”，调度层负责“何时算、何时传”。三者通过正交通信组与可替换调度器解耦，使得同一套模型切分可以适配不同的效率优化策略。