第6章：流水线并行PP

📖 本章概述

流水线并行（Pipeline Parallelism, PP）将模型的不同层分配到不同设备，是跨机扩展的主要手段。本章聚焦于 PP 的核心挑战——Pipeline Bubble（流水线气泡），以及三代调度策略如何逐步减少 Bubble。

---

📑 章节结构

1. 流水线并行的基本思想

• 层间切分：将模型的 $L$ 层均匀分配到 $P$ 个 Stage，每个 Stage 负责 $\frac{L}{P}$ 层
• 类比：像工厂流水线——每个工位负责一道工序，产品依次经过各工位
• 核心问题（Bubble）：朴素实现中，一个 micro-batch 必须串行经过所有 Stage → 大量 GPU 在等待中空转
• Bubble 比例（朴素 PP）：$\frac{P-1}{P}$（$P$ 个 Stage 时，有 $P-1$ 个时间步是浪费的）

2. GPipe：微批次流水线

• 核心思想：将一个 mini-batch 切分为 $M$ 个 micro-batch，依次注入流水线
• 执行模式：所有 micro-batch 前向完成后，再统一做反向（F-F-F-…-B-B-B-…）
• Bubble 比例：$\frac{P-1}{M+P-1}$，$M$ 越大 Bubble 越小
• 缺点：
  • 前向完成后才能开始反向 → 需要保存所有 micro-batch 的激活值 → 显存占用大
  • $M$ 太大时 Bubble 虽小但显存爆炸

3. 1F1B（One Forward One Backward）

• 核心思想：交替执行前向和反向，尽早释放激活值
• 执行模式：Warmup 阶段（填充流水线）→ Steady 阶段（每做一次 F 紧接一次 B）→ Cooldown 阶段（排空流水线）
• Bubble 比例：与 GPipe 相同 $\frac{P-1}{M+P-1}$
• 关键优势：峰值激活值显存从 $M$ 份降至 $P$ 份（同一时刻最多只有 $P$ 个 micro-batch 的激活值存活）
• Megatron-LM 默认调度

4. Interleaved 1F1B（虚拟 Stage）

• 核心思想：每个设备不只负责连续的 $\frac{L}{P}$ 层，而是负责多段非连续的层（虚拟 Stage）
• 设计：每个设备承担 $v$ 个虚拟 Stage（每个虚拟 Stage 有 $\frac{L}{P \cdot v}$ 层）
• Bubble 比例：$\frac{P-1}{M \cdot v + P - 1}$，等效将 $M$ 扩大了 $v$ 倍
• 代价：通信频率增加 $v$ 倍（每个虚拟 Stage 之间都需要跨设备传输激活值）
• 适用场景：机间带宽充足时，用额外通信换更低的 Bubble

5. PP 的工程挑战

• 负载均衡：不同层的计算量可能不同（Embedding 层 vs Transformer 层），需要合理分配 Stage
• Embedding 和 Loss 的处理：Embedding 层在模型最前端，Loss 在最后端，通常放在第一个和最后一个 Stage
• 跨机激活值传输：Stage 间传输的数据量 = $b \times s \times h$（batch × seq_len × hidden），需 IB 带宽支撑
• 与 Activation Checkpointing 的配合：PP 天然与 recomputation 配合——反向时可重新计算而非存储激活

6. Bubble 比例对比与可视化

调度策略	Bubble 比例	峰值激活显存	通信频率
朴素 PP	$(P-1)/P$	1 份	低
GPipe	$(P-1)/(M+P-1)$	$M$ 份	低
1F1B	$(P-1)/(M+P-1)$	$P$ 份	低
Interleaved 1F1B	$(P-1)/(Mv+P-1)$	$P$ 份	高（$v$ 倍）

7. 动手实验

• 画出 4 Stage、8 Micro-batch 下 GPipe 和 1F1B 的时间线图
• 计算不同 $M$、$P$、$v$ 组合下的 Bubble 比例
• 分析 64 层模型、PP=4 时每个 Stage 的层分配方案

---

🎯 本章学习目标

• 能解释 Pipeline Bubble 的成因和量化公式
• 能画出 GPipe、1F1B、Interleaved 1F1B 三种调度的时间线图
• 能计算给定配置下的 Bubble 比例和峰值激活显存
• 能分析 Interleaved 策略”通信换 Bubble”的 trade-off