流水线并行原理：模型切分策略与调度机制
大家好，我是你们的老朋友。今天我们来聊聊流水线并行中最核心的几个概念。说实话，我刚开始接触分布式训练时，觉得流水线并行不就是把模型切成几段嘛，有什么难的？结果第一次上手就踩了不少坑。嗯，今天我把这些经验都整理出来，希望能帮你少走弯路。
模型切分策略：按层切还是按算子切？
模型切分，说白了就是决定「从哪里下刀」。我个人习惯把切分策略分成两种：按层切分和按算子切分。
按层切分（Layer-wise Partitioning）
这是最直观的方式。比如一个12层的Transformer，切成3段，每段4层，分别放到3张GPU上。我在项目中遇到过一个大模型，就是按层切分的。优点是实现简单，代码改动少。但缺点也很明显——负载不均衡。
为什么会这样？你想想看，有些层计算量大（比如注意力层），有些层计算量小（比如LayerNorm）。如果切得不均匀，就会出现「有的GPU忙死，有的GPU闲死」的情况。
按算子切分（Operator-wise Partitioning）
这种策略更精细。不是按层切，而是按算子切。比如把一个Transformer层里的Self-Attention和FFN拆开，放到不同的GPU上。我在做GPT-3级别的模型时试过这种方式，负载均衡效果更好，但通信开销也更大。
举个例子：
# 按层切分示意
stage_0: layers[0:4] → GPU 0
stage_1: layers[4:8] → GPU 1
stage_2: layers[8:12] → GPU 2
# 按算子切分示意
stage_0: embedding + self_attention → GPU 0
stage_1: ffn + layer_norm → GPU 1
stage_2: output_projection → GPU 2微批次（Micro-batch）调度：让流水线转起来
模型切好了，接下来就是怎么调度数据。这里引入一个关键概念——微批次（Micro-batch）。
假设你的全局batch size是64，你可以把它拆成4个micro-batch，每个16条数据。为什么要这么做？说白了就是为了减少气泡。
我画了一张图，帮你理解微批次调度的核心逻辑：
看到图里的灰色区域了吗？那就是气泡。在流水线启动和收尾阶段，总有一些GPU在等待。微批次越多，气泡占比就越小，但也不是越多越好——微批次太多会增加通信开销。
前向与反向传播的流水线气泡
气泡（Bubble）是流水线并行中绕不开的话题。说白了，就是GPU在空转等待。我刚开始做流水线并行时，看到GPU利用率只有60%，急得不行。后来才明白，这是流水线固有的问题。
气泡产生的原因很简单：
- 启动阶段：前面的GPU在计算，后面的GPU在等待
- 收尾阶段：后面的GPU在计算，前面的GPU在等待
- 反向传播依赖：反向传播必须等前向传播完成后才能开始
怎么减少气泡？我总结了三个方向：
- 增加微批次数量——让流水线更「满」
- 均衡各stage的计算量——避免某个stage成为瓶颈
- 使用1F1B调度策略——这个下面会讲
GPipe vs PipeDream：两种主流方案对比
说到流水线并行，就绕不开两个经典方案：GPipe和PipeDream。我在实际项目中都用过，各有千秋。
| 对比维度 | GPipe | PipeDream |
|---|---|---|
| 调度策略 | 同步流水线：所有micro-batch前向完成后，统一开始反向 | 异步流水线：前向和反向可以交错执行（1F1B策略） |
| 气泡大小 | 较大（O(P-1) / O(M) 比例，P=stage数，M=micro-batch数） | 较小（通过1F1B策略减少气泡） |
| 显存占用 | 较高（需要保存所有micro-batch的中间激活值） | 较低（只保存当前正在计算的激活值） |
| 实现复杂度 | 简单，容易上手 | 较复杂，需要处理权重版本问题 |
| 适用场景 | 小规模流水线（2-4个stage） | 大规模流水线（8+个stage） |
GPipe的同步调度
GPipe的思路很直接：先把所有micro-batch的前向传播跑完，再统一跑反向传播。优点是实现简单，梯度更新是同步的，模型收敛稳定。缺点就是气泡大，显存占用高。
# GPipe调度示意（伪代码）
for micro_batch in micro_batches:
forward(micro_batch) # 所有stage依次前向
for micro_batch in micro_batches:
backward(micro_batch) # 所有stage依次反向
optimizer.step() # 统一更新参数PipeDream的1F1B调度
PipeDream用了更聪明的策略——1F1B（一个前向，一个反向）。什么意思呢？就是每个GPU在完成一个micro-batch的前向后，立刻开始另一个micro-batch的反向。这样流水线更紧凑，气泡更小。
# PipeDream 1F1B调度示意（伪代码）
# 假设有4个micro-batch，3个stage
for i in range(num_micro_batches):
forward(i) # 启动阶段
# 1F1B核心循环
for i in range(num_micro_batches):
backward(i) # 反向一个
forward(num_micro_batches + i) # 前向一个（如果有剩余）
# 收尾阶段
for i in range(num_micro_batches):
backward(num_micro_batches + i)我的选择建议
说了这么多，到底该用哪个？我个人习惯这样选：
- stage数 ≤ 4：用GPipe，简单稳定，调试方便
- stage数 ≥ 8：用PipeDream，显存效率高，气泡小
- 显存紧张：优先PipeDream，能省不少显存
- 追求收敛稳定性：GPipe更稳妥，毕竟是同步更新
嗯，今天的内容就到这里。流水线并行看起来简单，但真正调优起来，细节还是挺多的。希望这些经验能帮你少踩几个坑。
想要更多资料，联系可可