起因#
上一篇搞清楚了 Pipeline Parallelism 的调度机制：GPipe 先 Forward 后 Backward，1F1B 通过交替调度降低 Activation 暂存压力，但理想条件下 Bubble Time 相同。
理解了调度之后，我以为接下来只要调好参数就行了。但继续看 MIT 6.S191 Lecture 9 的材料，发现有个问题没那么直觉：
假设模型已经均匀切成 4 个 Stage，用 1F1B 调度，Micro-batch 也配好了。如果再加 4 张 GPU 把 Pipeline 拉长到 8 个 Stage——每个 Stage 的计算量减半，总算力翻倍——训练不一定更快。
这让我意识到，调度只是 PP 的一个维度。从调度理论到实际训练吞吐之间，还有一层工程权衡需要想清楚。
Forward 传 Activation，Backward 传 Input Gradient#
先把 PP 中 Stage 之间到底传递什么搞清楚。这个问题我在看 TP 和 PP 的通信区别时想了比较久。
用最简单的两层线性网络：
$$ h = W_0 x, \quad y = W_1 h, \quad L = \ell(y, t) $$GPU 0 负责第一层（\(W_0\)），GPU 1 负责第二层（\(W_1\)）。
Forward 很直接：GPU 0 算出 \(h\)，发送给 GPU 1。GPU 1 用 \(h\) 计算 \(y\) 和 Loss \(L\)。跨 GPU 传递的是 Activation \(h\)。
Backward 需要更仔细地看。GPU 1 需要计算两个东西：
- \(\frac{\partial L}{\partial W_1}\)——用于更新自己的参数。这个梯度留在 GPU 1 本地，不需要传给任何人。
- \(\frac{\partial L}{\partial h}\)——Loss 对边界 Activation \(h\) 的梯度。这个需要发回 GPU 0，因为 GPU 0 的参数更新依赖它。
GPU 0 收到 \(\frac{\partial L}{\partial h}\) 后，计算：
$$ \frac{\partial L}{\partial W_0} = \frac{\partial L}{\partial h} x^T $$这里用的是列向量约定（\(h = W_0 x\)，\(\frac{\partial L}{\partial W_0}\) 是外积形式）。
有一点容易混淆：Backward 传递的是 Loss 对边界 Activation 的梯度 \(\frac{\partial L}{\partial h}\)，不是 \(\frac{\partial y}{\partial h}\)。前者是链式法则一路乘到 Loss 的结果，后者只是相邻两层的局部 Jacobian。也不是参数梯度——参数梯度（Weight Gradient）始终留在计算它的 Stage 本地。
每个 Stage 不需要拿到其他 Stage 的 Weight Gradient。 参数更新是各 Stage 独立完成的：每个 Stage 用自己收到的 Input Gradient 和本地保存的 Activation 计算出自己的 Weight Gradient，然后直接更新本地参数。
PP 主要依赖 P2P，TP 经常需要 Collective#
理解了 PP 传递什么之后，一个自然的问题是：PP 和 TP 的通信模式为什么不同？
PP 的通信模式来自模型的层间依赖：
$$ h_2 = f_1(f_0(x)) $$Layer 0 的输出是 Layer 1 的输入。这是一条串行链。跨 Stage 的通信只发生在相邻 Stage 之间——GPU 0 把 Activation 发给 GPU 1，GPU 1 把 Input Gradient 发回 GPU 0。这是点对点的 Send/Recv（P2P）。
TP 那篇讨论的 Row Parallel：
$$ Y = X_1 W_1 + X_2 W_2 $$每张 GPU 算出局部结果后，需要跨设备求和——这是一个 AllReduce。TP 的通信来自单个算子内部的计算分解，需要 Collective 操作把局部结果聚合成完整输出。
总结区别：
- PP：跨 Layer 的数据流，主要是相邻 Stage 之间的 P2P。
- TP：单个 Layer 内部的计算分片，经常需要 AllReduce 等 Collective。
但不能写成"PP 永远没有 Collective"——实际系统组合了 PP、TP 和 DP，训练过程中仍然存在 Collective 通信（比如 DP 的梯度同步）。逻辑计算依赖决定了通信模式，而不是开发者任意选择一种通信 API。
最慢的 Stage 决定稳态吞吐#
这一节让我对 PP 的工程直觉有了一个重要修正。
假设 4 个 Stage 对单个 Micro-batch 的 Forward + Backward 总耗时分别是：
1[2, 2, 6, 2] ms在忽略通信、计算时间稳定的假设下，Pipeline 进入稳态后，每个 Micro-batch 必须经过所有 Stage。而 Stage 2 需要 6 ms 才能完成——其他 Stage 再快，也得等 Stage 2 腾出来。
理想稳态吞吐上限：
$$ \text{Throughput} \lesssim \frac{1}{\max_i T_i} $$Stage 2 是瓶颈。每 6 ms 最多处理一个 Micro-batch，不管其他 Stage 是 2 ms 还是 0.1 ms。
现在假设新增一张 GPU，把 4 Stage 扩展到 5 Stage。有两种切法：
方案 A：[1, 1, 2, 6, 2] ms——把前面的 Stage 切得更细。但最慢的 Stage 2 仍然是 6 ms，瓶颈没有改善。
方案 B：[2, 2, 3, 3, 2] ms——把最慢的 Stage 拆成两段。瓶颈从 6 ms 降到 3 ms。在不考虑其他开销的情况下，理想稳态吞吐上限从 \(\frac{1}{6}\) 提升到 \(\frac{1}{3}\) MB/ms。
需要说明：这只是稳态理想上限。实际的 Iteration Time 还受 Warmup/Cooldown Bubble、通信延迟等因素影响。而且新增一个 Stage 本身会增加 Pipeline 深度 \(p\)，可能引入更多 Bubble。
优化 Stage Balance 不是追求所有 Stage 一样快，而是消除极端瓶颈。 Transformer 的不同模块（Embedding、Attention、FFN、LayerNorm）计算量不同，均匀按层数切分不一定是最优的。
Pipeline Bubble 的公式和两种分母#
上一篇建立了 Bubble 的直觉。这里把公式写清楚。
定义：
- \(p\)：Pipeline Stage 数量
- \(m\)：Micro-batch 数量
- \(T_F\)：单 Stage 单 Micro-batch 的 Forward 时间
- \(T_B\)：单 Stage 单 Micro-batch 的 Backward 时间
在各 Stage 均衡、忽略通信、普通 Pipeline Flush 调度的假设下，一次 Iteration 的总时间：
$$ T_{\text{iteration}} \approx (m + p - 1)(T_F + T_B) $$其中有效计算时间是 \(m(T_F + T_B)\)，Bubble Time 是 \((p - 1)(T_F + T_B)\)。
Bubble 占实际 Iteration Time 的比例：
$$ \text{Bubble Ratio} = \frac{p - 1}{m + p - 1} $$但 Megatron-LM 论文里经常用另一种定义——Bubble 相对于理想计算时间：
$$ \text{Bubble Overhead} = \frac{p - 1}{m} $$两个公式的分母不同：前者是 \(m + p - 1\)（实际时间），后者是 \(m\)（理想时间）。不能混为同一个指标。
用具体数字感受一下：
| \(p\) | \(m\) | \(\frac{p-1}{m+p-1}\) | \(\frac{p-1}{m}\) |
|---|---|---|---|
| 4 | 8 | 27% | 38% |
| 8 | 8 | 47% | 88% |
| 8 | 32 | 18% | 22% |
几个观察：
- \(p = 8, m = 8\)：Bubble 占了接近一半的实际时间。Stage 越多，固定 \(m\) 下 Bubble 越严重。
- \(p = 8, m = 32\)：增加 Micro-batch 数量可以有效摊薄 Bubble。
- 但这只是理想公式，不是实际性能的准确预测。
曲线展示了不同 \(p\) 下，Bubble Ratio 随 \(m\) 增加的下降趋势。\(p\) 越大，需要更多的 \(m\) 才能把 Bubble 压到可接受的范围。
为什么不能无限增加 Micro-batch 数量#
看完上面的曲线，一个自然的想法是：既然 Micro-batch 越多 Bubble 越小，为什么不用尽可能多的 Micro-batch？
这里有一个关键约束。假设固定了 Global Batch Size \(B_{\text{global}}\) 和 Data Parallel Degree \(DP\)：
$$ B_{\text{global}} = DP \times m \times b $$其中 \(b\) 是单个 Micro-batch 的样本数。固定 \(B_{\text{global}}\) 和 \(DP\)，增加 \(m\) 意味着减小 \(b\)。
\(b\) 变小带来的工程代价：
- GEMM 规模缩小：GPU 上的矩阵乘法效率与矩阵规模密切相关。batch 维度缩小，GEMM 无法充分利用 GPU 的并行计算单元，Kernel 执行效率下降。
- Kernel Launch 与调度开销占比上升：每个 Micro-batch 都需要启动一系列 CUDA Kernel。\(m\) 变大意味着更多次 Kernel Launch，固定的调度开销被放大。
- P2P 通信次数增加：每个 Micro-batch 都需要在相邻 Stage 间传递 Activation 和 Gradient。更多 Micro-batch 意味着更多次通信，较小 Tensor 的通信启动延迟占比更高。
更少 Bubble 不等于更高 GPU Kernel 效率，更不等于更高训练吞吐。 实际的最优 \(m\) 是 Bubble 减少和 Kernel 效率下降之间的平衡点。
另一种情况：如果不固定 \(B_{\text{global}}\)，而是固定 \(b\) 只增加 \(m\)——Global Batch Size 会随之增大。这不仅是性能参数变化，还会影响梯度累积语义和优化器更新频率。这已经不是同一个训练配置下的性能比较了。
Activation Memory：AFAB 与 1F1B 的工程含义#
简洁回顾上一篇的结论。
固定单个 Micro-batch 的 Activation 大小为 \(A\)：
- AFAB：在途 Activation 峰值约 \(m \cdot A\)。所有 Forward 完成后才开始 Backward，\(m\) 个 Micro-batch 的 Activation 同时驻留。
- 1F1B：在途 Activation 峰值约 \(p \cdot A\)。Warmup 后交替执行 F/B，及时释放已完成 Backward 的 Activation。
当 \(m = 32, p = 4\) 时，1F1B 的 Activation 暂存量约为 AFAB 的 \(\frac{1}{8}\)。
这个差距在工程上意味着：1F1B 让我们有更多的显存空间来增大 \(m\) 或增大 \(b\)，间接提升训练吞吐。AFAB 在 \(m\) 较大时可能因 Activation Memory 撑爆显存，1F1B 则不容易遇到这个瓶颈。
但实际峰值显存不只是在途 Activation。参数、梯度、Optimizer States、通信 Buffer、CUDA Runtime 开销都占显存。之前的文章详细算过这些账。不要把 \(O(p)\) 的在途 Activation 数量级直接当作整个 GPU 显存占用。
为什么增加 GPU 不一定更快#
从 Critical Path 的角度整理增加 PP Degree 可能带来的工程代价：
Bubble 增加。\(p\) 变大，固定 \(m\) 下 Bubble Ratio \(\frac{p-1}{m+p-1}\) 上升。即使增加 \(m\) 来补偿，也受到 Kernel 效率的约束。
Stage Balance 更难保证。Stage 越多，模型被切得越细。Transformer 不同模块的计算量不同——Embedding 层、Attention 层、FFN 层的 FLOPs 差异可能导致某些 Stage 明显快于或慢于其他 Stage。而最慢的 Stage 决定稳态吞吐上限。
P2P 通信边界增加。每增加一个 Stage 边界，就多一组 Forward Activation 和 Backward Gradient 的传输。如果 GPU 之间的互连带宽有限（比如跨节点走以太网而非 NVLink），通信开销可能进入关键路径。
Micro-batch 过小时 Kernel 效率下降。Stage 增加可能要求更多 Micro-batch 来压低 Bubble，但固定 Global Batch 下 Micro-batch Size 变小，GEMM 效率下降。
GPU 间拓扑不均匀。同一节点内的 GPU 通过 NVLink 连接，跨节点的 GPU 走网络。PP 的 Stage 边界如果落在跨节点链路上，通信延迟可能显著高于节点内。
端到端瓶颈可能不在计算。DataLoader、数据预处理、Checkpoint 写入等环节可能是实际的训练瓶颈，增加 GPU 不会改善这些环节。
优化 PP 不是追求 Stage 越多越好，也不是追求 Bubble 越小越好，而是在完整训练工作负载下最大化有效吞吐。
PP 性能优化的实验框架#
最后给出一个可执行的思路。假设固定了模型、序列长度和硬件总规模，需要选择最优的 PP 配置。
核心实验设计：
- 固定训练语义（有效 Global Batch Size 不变）
- 对比不同 PP Degree（\(p = 1, 2, 4, 8\)）
- 对比不同 Micro-batch 配置（\(m\) 和 \(b\) 的组合）
- 保持其他关键设置（TP Degree、DP Degree、Activation Checkpointing）尽可能一致
需要关注的指标：
| 指标 | 来源 | 说明 |
|---|---|---|
| End-to-end Iteration Time | 训练循环计时 | 最终决策指标 |
| Samples/s 或 Tokens/s | 上述推算 | 有效训练吞吐 |
| 各 Stage Forward/Backward 耗时 | Profiler | 判断 Stage Balance |
| Pipeline Bubble 占比 | Profiler / 计算 | 理论 vs 实际 |
| P2P Communication Time | Profiler | 是否在关键路径 |
| Peak GPU Memory | torch.cuda.max_memory_allocated | 是否有显存余量 |
| GPU SM Active Time | NSight / Profiler | 实际计算效率 |
几个实践要点：
- 不能只看
nvidia-smi上的 GPU Utilization。它统计的是"GPU 上有 Kernel 在跑"的时间比例，不区分 Kernel 是否高效。一个很小的 GEMM 和一个充分利用计算单元的 GEMM 在nvidia-smi上都是 100% utilization。 - 应该比较实际训练吞吐，而不是只比较理论 Bubble。两种配置的理论 Bubble 差 5%，但 Kernel 效率差 20%，最终吞吐可能反转。
- Profile 时注意区分计算瓶颈和通信瓶颈。如果某个 Stage 的 Idle 主要在等通信完成，优化方向是通信（更快的互连、更大的 Tensor 来分摊启动延迟、通信与计算 Overlap）。如果 Idle 是因为 Stage Balance 差，优化方向是模型划分。
写在最后#
| 方法 | 主要收益 | 主要代价 |
|---|---|---|
| 增加 PP Degree | 单 Stage 计算和 Model States 减少 | Bubble 增加、通信边界增加 |
| 增加 Micro-batch Count | 减少理想 Bubble 比例 | 固定 Batch 下单个 Micro-batch 更小 |
| 优化 Stage Balance | 提升稳态吞吐上限 | 切分和调度更复杂 |
| 使用 1F1B | 降低 Activation 暂存压力 | 不会自动消除 Bubble |
| 通信计算 Overlap | 隐藏部分通信开销 | 受依赖关系和实现限制 |
回看这两篇，我觉得 PP 最本质的问题可以用四个维度概括：
- Why：模型太大单卡放不下，或者希望减少单卡的计算量。PP 把模型按层切分到多张 GPU。
- How：模型切分定义 Stage，Micro-batch 让 Pipeline 流动，调度策略（GPipe / 1F1B）决定 Forward 和 Backward 的执行顺序。
- Trade-off：Bubble 是固有代价，通信是额外成本，Stage Balance 是稳态吞吐的天花板。
- Boundary：当 Micro-batch 过小导致 Kernel 效率太低、当 Stage 数量过多导致 Bubble 和通信吃掉所有收益、当端到端瓶颈不在 GPU 计算——PP 可能反而拖慢训练。
参考资料#
- Pipeline Parallelism：从 GPipe 到 1F1B，为什么 GPU 还是在等待？
- Tensor Parallelism：为什么矩阵乘法可以拆分到多张 GPU 上？
- 从 DDP 到 ZeRO：为什么分布式训练要对参数、梯度和优化器状态分片？
- 7B 模型只有 14 GB 权重，一张 80 GB GPU 为什么仍可能 OOM？
- 从 AllReduce 到 Ring：分布式训练中的通信原语
- Huang et al. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. NeurIPS 2019.
- Narayanan et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. SC 2021.
- NVIDIA Megatron-LM Pipeline Parallel Source
- NVIDIA: Scaling Language Model Training
- MIT 6.S191: Introduction to Deep Learning — Lecture 9: Secrets to Massively Parallel Training. https://introtodeeplearning.com/