流水线并行 PP 与 interleaved schedule

8.1 为什么 TP 之外还要 PP

TP 的痛点：

• 受限于单节点 NVLink 域，TP 度 ≤ 8；
• 跨节点 TP 性能会塌方；
• 对 175B 这种级别的模型，TP=8 还不够（每卡仍要 22B 的状态）。

PP 把模型的不同层放到不同 GPU 上，每张卡只算一段。175B 经典配置 PP=16，每张卡装 6 层，配上 TP=8 就 8×16=128 卡支撑一份完整模型。

8.2 PP 的核心问题：bubble

如果只把一个 batch 送进 pipeline，会出现严重空闲：

每条 stage 在 forward / backward 大部分时间都在等别人。这就是 pipeline bubble。

8.3 解决之一：micro batch 流水

把一个 global batch 拆成 $m$ 个 micro batch，让它们像水管里的水一样接连进入 pipeline：

bubble 比例：

$$ \text{bubble ratio} = \frac{S - 1}{m + S - 1} $$

$S$ = stage 数，$m$ = micro batch 数。$m$ 越大 bubble 越小，但单步显存压力越大（要存 $m$ 份激活）。

8.4 1F1B：让显存可控的调度

"all-forward-all-backward"（GPipe 风格）虽然简单，但每个 stage 要存 $m$ 份激活，显存爆炸。Megatron 用 "1F1B（one-forward-one-backward）"调度：前向到第 $S$ 个 micro batch 后开始穿插反向。

每张卡同时只持有 ~$S$ 份激活（不是 $m$ 份）。

8.5 Interleaved（virtual pipeline）：进一步降 bubble

1F1B 已经把激活控制住了，但 bubble 仍存在。Megatron 提出 virtual pipeline（也叫 interleaved 1F1B）：让每张物理卡持有多段不连续的层。

例如 PP=4，每张卡原本装 $L/4$ 层；启用 virtual pipeline=2 后，每张卡装 两段 $L/8$ 层，分别对应 virtual stage 0 和 4。这样一个 micro batch 在 pipeline 里转两圈，bubble 变小：

$$ \text{bubble ratio}_\text{interleaved} = \frac{1}{v} \cdot \frac{S - 1}{m + S - 1} $$

$v$ 是 virtual pipeline 大小。代价是通信次数变 $v$ 倍，所以对网络带宽要求更高。

8.6 在 Megatron 里怎么打开

# 基本 PP
--pipeline-model-parallel-size 16

# Interleaved（virtual pipeline）
--pipeline-model-parallel-size 16 \
--virtual-pipeline-model-parallel-size 2 \
--num-layers-per-virtual-pipeline-stage 3        # 96 层 / (16 × 2) = 3 层每 vstage

# 让 embedding 单独占一个 stage（loss balancing）
--standalone-embedding-stage

8.7 layer 怎么分配到 stage

Megatron 提供几种 layer layout 策略：

8.8 关键源码

8.9 forward_backward_func：用户可见入口

外部代码（包括 pretrain_gpt.py）调的入口是：

from megatron.core.pipeline_parallel.schedules import get_forward_backward_func

forward_backward_func = get_forward_backward_func()

losses = forward_backward_func(
    forward_step_func=forward_step_func,
    data_iterator=data_iterator,
    model=[model],
    num_microbatches=num_microbatches,
    seq_length=seq_length,
    micro_batch_size=micro_batch_size,
    forward_only=False,
)

这个函数返回的不是单一 forward / backward，而是整套 1F1B 调度的执行结果。算法工程师只要传 forward_step_func 这一个 callback，剩下都交给 schedule。

8.10 PP 通信量

每个 micro batch 跨 stage 的 p2p 通信 = $B_\mu \cdot S \cdot H \cdot 2\,\text{bytes}$（bf16）。
全 step 通信量约 $2 m \cdot B_\mu \cdot S \cdot H \cdot 2$（前向 + 反向）。

相比 TP 的 all-reduce，PP 通信量小很多，所以 PP 能跨节点（infiniband 即可），TP 不行（要 NVLink）。

8.11 PP 跨节点 vs TP 跨节点

经典放置法：单节点内 TP=8，跨节点 PP=16，再跨更大范围 DP。

8.12 PP 的典型坑

8.13 schedules.py：三个 forward_backward 函数

所有 PP 调度集中在 megatron/core/pipeline_parallel/schedules.py（104KB）。入口和三种调度：

另外新文件 pipeline_parallel/hybrid_cp_schedule.py（2025-Q2）专门为 PP + CP 联合 设计——超长序列下原 1F1B 的 buffer 数与 chunk 模型不再够用。

8.14 ParallelState 关键 API

调试 PP 时一定会用到，megatron/core/parallel_state.py（1790 行）的核心几个：

整个文件有约 53 个 get_*_world_size / get_*_rank，覆盖所有并行维度（TP/PP/DP/CP/EP/MoE 内 DP/MoE 内 EP/...）。

8.15 这章你需要带走什么

• PP 把不同层放到不同 GPU；
• 1F1B 调度把显存占用从 $m$ 份激活降到 $S$ 份；
• bubble 比例 = $(S{-}1)/(m{+}S{-}1)$，调大 micro batch 数 $m$ 来降；
• Interleaved（virtual pipeline）进一步降 bubble，代价是通信变多；
• PP 可跨节点，TP 不行；
• forward_backward_func 是用户可见入口。