Pipeline Parallel,简称 PP,沿着模型深度切分。 如果一个 Transformer 有 80 层,4 个 pipeline stage 可以让每张 GPU 只放 20 层。 stage 之间传的不是参数,而是 micro-batch 的 activation; 反向时再把 activation gradient 反方向传回来。 这让 PP 特别适合跨节点:和 Tensor Parallel 相比,它通信频率低得多。

三个基础概念

同样是 PP,调度策略会决定显存峰值、bubble、通信复杂度。 GPipe 最好理解:先把所有 micro-batch forward 完,再 backward。 1F1B 更像真实生产默认:填满后每个 stage 交替做一个 forward 和一个 backward。 Megatron 的 Interleaved 1F1B 再把每张卡上的连续层切成多个 virtual stage,缩短 pipeline 单元粒度。

GPipe: all forward, then all backward

1F1B: one forward, one backward

假设 p 个 pipeline stage,每个 micro-batch 在每个 stage 上的 forward+backward 单位时间相同, 一次 global batch 被切成 m 个 micro-batch。 不考虑通信与 stage 不均衡时,常见的 pipeline bubble fraction 可以近似写成: (p - 1) / (m + p - 1)。 这个式子是 Day20 最重要的手算工具。

为什么 micro-batch 不能无限加

一个小 simulator

def bubble_fraction(p, m):
    return (p - 1) / (m + p - 1)

def table(stages=(2, 4, 8), micro_batches=(4, 8, 16, 32, 64)):
    for p in stages:
        print(f"pipeline stages = {p}")
        for m in micro_batches:
            frac = bubble_fraction(p, m)
            print(f"  m={m:2d}  bubble={frac*100:5.1f}%  useful={(1-frac)*100:5.1f}%")
        print()

table()

Interleaved 1F1B 的直觉是:每张物理 GPU 不只负责一段连续的大 stage, 而是负责多个更小的 model chunk,也叫 virtual pipeline stage。 这样 pipeline 深度在调度层面变细,每个 chunk 更早 ready,可以减少 bubble, 也能缓解某个物理 stage 太重的问题。

Megatron 参数形状

# 示例:TP=2, PP=4, 每个物理 stage 再拆 virtual chunks
torchrun --nproc_per_node=8 pretrain_gpt.py \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 4 \
  --num-layers 48 \
  --num-layers-per-virtual-pipeline-stage 3 \
  --micro-batch-size 2 \
  --global-batch-size 128

今天的交付物是一张 pipeline 实验表: 固定 global batch 和模型层数,改变 pipeline stage 数、micro-batch 数、是否 interleaved, 记录 bubble 估算、显存峰值和 step time。 没有多卡也可以先用纸面公式和 simulator 建立直觉。

from dataclasses import dataclass

@dataclass
class PPConfig:
    pp: int
    micro_batches: int
    micro_batch_size: int
    dp: int = 1
    tp: int = 1

    def global_batch(self):
        return self.micro_batches * self.micro_batch_size * self.dp

    def bubble(self):
        return (self.pp - 1) / (self.micro_batches + self.pp - 1)

configs = [
    PPConfig(pp=2, micro_batches=8, micro_batch_size=4),
    PPConfig(pp=4, micro_batches=8, micro_batch_size=4),
    PPConfig(pp=4, micro_batches=32, micro_batch_size=1),
    PPConfig(pp=8, micro_batches=64, micro_batch_size=1),
]

for cfg in configs:
    print(
        f"pp={cfg.pp} m={cfg.micro_batches:2d} "
        f"mbs={cfg.micro_batch_size} global={cfg.global_batch():3d} "
        f"bubble={cfg.bubble()*100:5.1f}%"
    )

实验记录模板

1. 画出 4 个 pipeline stage、4 个 micro-batch 的 GPipe 时序图,标出哪些格子是 bubble。
2. 解释 1F1B 的 warmup、steady state、cooldown 三个阶段,以及为什么它比 GPipe 更省 activation memory。
3. 推导并使用 (p - 1) / (m + p - 1) 估算 bubble fraction,分别计算 p=4,m=8 与 p=8,m=64。
4. 说明 Interleaved 1F1B 中 physical stage 与 virtual stage 的区别,以及 Megatron 的 --num-layers-per-virtual-pipeline-stage 大概控制什么。
5. 给定一个 48 层 Transformer、8 张 GPU,设计一个 TP×PP×DP 切分方案,解释为什么 TP 更适合节点内、PP 更适合跨节点。