【分布式训练⑤】流水线并行：按层切分模型与气泡率优化

【分布式训练⑤】流水线并行：按层切分模型与气泡率优化

本文是「分布式训练」系列的第五篇，带你从零理解流水线并行（Pipeline Parallelism, PP）的核心原理。当模型深到单卡放不下所有层时，我们可以把模型按深度"切开"，像工厂流水线一样让数据逐层流过不同的 GPU。但流水线也有自己的烦恼——气泡问题。本文从 GPipe 到 1F1B，从气泡率公式到优化策略，一次讲透。

---

一、开篇引入：流水线并行——把模型按深度切开

上一篇我们讲了张量并行（TP），它解决的是"单层太宽装不下"的问题——把矩阵按维度切开，分到多张卡上。

但还有另一种情况：模型太深了，一层能放下，但几十上百层堆在一起就放不下了。

比如一个 80 层的 Transformer 大模型，单张卡只能放下 20 层——这时候怎么办？

答案就是流水线并行（Pipeline Parallelism, PP）：把模型按层切分成多个 stage（阶段），每个 stage 放到一张或一组 GPU 上，数据像流水一样依次经过各个 stage。

1.1 一个生活化的类比：工厂流水线

想象一条汽车装配流水线：

• 工位 1：装底盘和发动机
• 工位 2：装车身和车门
• 工位 3：装内饰和座椅
• 工位 4：装轮胎和检测

一辆汽车从工位 1 流到工位 4，每个工位只做自己的事。第一辆车从工位 1 走了之后，工位 1 就可以开始装第二辆——这样四个工位可以同时工作，效率大大提高。

流水线并行就是同一个道理：

• GPU 0 负责第 1~20 层
• GPU 1 负责第 21~40 层
• GPU 2 负责第 41~60 层
• GPU 3 负责第 61~80 层

数据从 GPU 0 流到 GPU 3，每张卡只算自己那几层。

1.2 PP vs TP：到底有什么不一样？

特性	张量并行（TP）	流水线并行（PP）
切分维度	单层参数（宽度）	层（深度）
通信频率	每层都通信	只在 stage 边界通信
通信量	大（每层都有 AllGather/AllReduce）	小（只传激活值，P2P 通信）
并行度限制	受 hidden/head 数限制	受层数限制
气泡问题	没有	有（GPU 空闲时间）
典型部署	节点内（NVLink）	节点间（IB）或节点内

一句话总结：TP 切"宽度"，通信多但没有气泡；PP 切"深度"，通信少但有气泡。两者通常配合使用。

---

二、流水线并行基础

2.1 按层切分怎么切？

假设模型有 8 层 Transformer，用 4 张卡做流水线并行：

GPU 0 (Stage 0): Layer 0, Layer 1
GPU 1 (Stage 1): Layer 2, Layer 3
GPU 2 (Stage 2): Layer 4, Layer 5
GPU 3 (Stage 3): Layer 6, Layer 7

数据流动方向：

输入 → GPU 0 → GPU 1 → GPU 2 → GPU 3 → 输出

关键细节：

• 每个 stage 可以包含多层，不一定是一层对应一张卡
• 划分时尽量让每个 stage 的计算量均衡，避免有的卡闲有的卡忙
• stage 之间只需要传递激活值（P2P 通信，send/recv）

2.2 为什么不直接一张卡放所有层？

问得好。如果模型小到单卡能放下，那当然不用 PP。但对于大模型：

• 70B 模型：约 80 层，单卡 A100 80GB 放不下 FP16 的权重（140GB）+ 激活值 + 梯度 + 优化器状态
• 175B 模型：约 96 层，权重就 350GB，更是不可能单卡放下

PP 让我们可以把模型"摊开"到多张卡上，每张卡只存一部分层的参数——这就是 PP 的核心价值。

2.3 Micro-batch 是什么？

在流水线并行中，一个 batch 的数据会被进一步切成更小的 micro-batch。为什么？

因为如果整个 batch 作为一个整体流过流水线，那么同一时间只有一个 stage 在工作，其他 stage 都在等——这和单卡没区别，还多了通信开销。

切成 micro-batch 之后，流水线就可以"填满"了：

时刻 t0: m1 在 GPU 0
时刻 t1: m1 在 GPU 1, m2 在 GPU 0
时刻 t2: m1 在 GPU 2, m2 在 GPU 1, m3 在 GPU 0
时刻 t3: m1 在 GPU 3, m2 在 GPU 2, m3 在 GPU 1, m4 在 GPU 0

看到了吗？从 t3 开始，4 张卡都在同时工作！这就是流水线的魔力。

---

三、GPipe：fill-flush 方式与气泡问题

GPipe 是 Google 在 2018 年提出的流水线并行方案，也是最经典、最容易理解的一种。

3.1 GPipe 的工作流程

GPipe 的流程可以概括为 fill → full → flush 三个阶段：

Forward 阶段（前向传播）：

GPU 0: [m1][m2][m3][m4]
GPU 1:    [m1][m2][m3][m4]
GPU 2:       [m1][m2][m3][m4]
GPU 3:          [m1][m2][m3][m4]
 时间: ─────────────────────────▶
        ├─ fill ─┤  full  ├─ flush ─┤

• fill（填充）：前 P-1 个时刻，流水线逐渐填满，后面的 GPU 还没开始工作
• full（满载）：所有 GPU 都在工作，流水线效率最高
• flush（排空）：最后 P-1 个时刻，流水线逐渐清空，前面的 GPU 已经没事干了

Backward 阶段（反向传播）：

Forward 全部完成之后，才开始 Backward。方向相反，从最后一个 stage 往回传：

GPU 3: [m4][m3][m2][m1]
GPU 2:    [m4][m3][m2][m1]
GPU 1:       [m4][m3][m2][m1]
GPU 0:          [m4][m3][m2][m1]

3.2 什么是"气泡"？

在 fill 和 flush 阶段，有些 GPU 是空闲的——这些空闲时间在时序图上看起来就是"空泡"，所以叫做气泡（Bubble）。

气泡是流水线并行的"天生缺陷"——你不可能让所有 GPU 所有时刻都在工作。但气泡的大小是可以优化的。

3.3 气泡率公式

怎么衡量气泡有多严重？我们用气泡率（Bubble Rate）——GPU 空闲时间占总时间的比例。

对于 GPipe（forward 全部完成后再 backward）：

$$\text{Bubble Rate}=\frac{(P-1)}{M}\times \frac{t_{fw}}{t_{fw}+t_{bw}}$$

其中：

• $P$：stage 数（流水线阶段数）
• $M$：micro-batch 数
• $t_{fw}$：一个 micro-batch 的 forward 时间
• $t_{bw}$：一个 micro-batch 的 backward 时间

通常 backward 时间大约是 forward 的 2 倍（因为要计算梯度），所以 $t_{fw}:t_{bw}\approx 1:2$。

举个例子：P=4, M=4, fw:bw = 1:2

$$\text{Bubble Rate}=\frac{4-1}{4}\times \frac{1}{1+2}=\frac{3}{4}\times \frac{1}{3}=25\%$$

也就是说，GPU 有 25% 的时间在空转！这还是 M=4 的情况，如果 M 更小，气泡更大。

---

四、PipeDream / 1F1B 调度

GPipe 气泡太大了，有什么办法减少吗？当然有——**1F1B（One Forward One Backward）**调度就是最经典的优化方案。

4.1 1F1B 是什么意思？

GPipe 是"先把所有 forward 做完，再做所有 backward"。

1F1B 更聪明：做完一个 micro-batch 的 forward 之后，能做 backward 就尽快做 backward，不用等全部 forward 完成。

用时间线来看更直观（P=4, M=4）：

GPU 0: f1 f2 f3 f4 b1 b2 b3 b4
GPU 1:    f1 f2 f3 f4 b1 b2 b3 b4
GPU 2:       f1 f2 f3 f4 b1 b2 b3 b4
GPU 3:          f1 f2 f3 f4 b1 b2 b3 b4

注：f = forward, b = backward

对比 GPipe 的时间线：

GPipe:  f1 f2 f3 f4 ........ b4 b3 b2 b1
           (等待所有f完成)    (反向倒序)

1F1B:   f1 f2 f3 f4 b1 b2 b3 b4
           (紧接着就开始b)

1F1B 把 backward 往前"挪"了，让 forward 和 backward 交错执行，减少了 GPU 的等待时间。

4.2 1F1B 的气泡率

1F1B 的气泡率大约是：

$$\text{Bubble Rate}=\frac{P-1}{M+P-1}$$

还用刚才的例子：P=4, M=4

$$\text{Bubble Rate}=\frac{4-1}{4+4-1}=\frac{3}{7}\approx 42.9\%$$

哎？怎么看起来比 GPipe 的 25% 还高？别急，这里的公式对比的基准不一样。

GPipe 的公式里有个 $\frac{t_{fw}}{t_{fw}+t_{bw}}$ 的因子（约 1/3），因为 forward + backward 加起来才是完整的 step。

我们统一用"完整 step 的时间"来对比：

调度方式	P=4, M=4 气泡率
GPipe	~25%
1F1B	~30%

差距其实不大，但 1F1B 有一个巨大的优势——激活显存小很多。等下我们详细说。

4.3 Interleaved 1F1B：进一步优化

上面说的是"简单版"1F1B。Megatron-LM 中使用的是 Interleaved 1F1B（交错式 1F1B）——每个 GPU 负责更多的"虚拟 stage"，让 forward 和 backward 更加交错，气泡率进一步降低。

这个概念稍微复杂一点，初学者记住结论就行：1F1B 比 GPipe 气泡小、激活显存管理更灵活，是当前大模型训练的主流选择。

---

五、气泡率分析与优化

气泡率直接影响 GPU 利用率，是流水线并行最重要的性能指标之一。下面我们深入分析影响因素和优化策略。

5.1 影响气泡率的三大因素

因素一：Stage 数 P

P 越大，气泡越大。因为 fill 和 flush 阶段的长度都是 P-1。

• P=2：fill 只需 1 个时间单位，气泡小
• P=8：fill 需要 7 个时间单位，气泡大

直觉理解：流水线越长，"热身"和"收尾"的时间就越长。

因素二：Micro-batch 数 M

M 越大，气泡越小。因为满载阶段变长了，气泡占比相对降低。

• M=4：气泡占比高
• M=32：气泡占比低很多

直觉理解：生产线上产品越多，流水线"填满"之后的满载时间就越长，相对的空闲时间占比就越小。

因素三：Forward / Backward 时间比

通常 backward 是 forward 的 1.5~2 倍（因为要算梯度，计算量更大）。这个比例影响 GPipe 的气泡率，但对 1F1B 影响较小。

5.2 优化策略总结

优化策略	原理	代价
增加 micro-batch 数	M 增大，气泡占比降低	激活显存增加，调度开销增加
使用 1F1B / Interleaved	forward/backward 交错，减少等待	实现更复杂，激活显存管理更精细
减少 stage 数	P 减小，气泡直接减小	每卡需要放下更多层，显存压力大
Activation Checkpointing	用计算换显存，支持更多 M	训练速度降低 ~20-30%
各 stage 计算量均衡	避免某个 stage 成为瓶颈	划分更复杂，可能需要 profiling

5.3 一个实际的数字感

以 70B 模型、4 节点（32 卡）训练为例，常见配置：

• TP=8（节点内）
• PP=4（跨节点）
• DP=1（或 2，取决于总卡数）
• Micro-batch 数 M = 16~32

气泡率大约在 10%~20% 之间——这是一个可以接受的范围。

如果 PP=8，M=16，气泡率可能高达 30%+，就不太划算了。

---

六、激活值 vs 气泡的权衡

聊到这里，你可能发现了一个关键的权衡：micro-batch 数越多，气泡越小，但激活值显存越大。

6.1 为什么激活值会变大？

在流水线并行中，每个 stage 都要保存自己负责的那些 micro-batch 的激活值，因为 backward 的时候要用。

• GPipe：forward 阶段保存所有 M 个 micro-batch 的激活，backward 阶段逐步释放
• 1F1B：每个时刻同时在流水线上的 micro-batch 大约是 P 个，所以需要保存的激活值更少

对比：

调度方式	每 stage 保存的 micro-batch 数	激活显存
GPipe	M（所有 micro-batch）	大
1F1B	~P（流水线上的数量）	小

比如 P=4, M=32：

• GPipe 每卡保存 32 份激活
• 1F1B 每卡保存约 4 份激活
• 差距是 8 倍！

这就是为什么 1F1B 更受欢迎——气泡率差不多，但激活显存小太多了。

6.2 Activation Checkpointing 怎么帮上忙？

Activation Checkpointing（激活重计算）的思想是：前向传播时不保存所有中间激活，反向传播时重新计算。

这样做的好处是可以把激活显存降低 30%~50%，代价是多花 20%~30% 的计算时间。

对于流水线并行来说，Activation Checkpointing 尤其重要，因为：

1. PP 本身就要保存多个 micro-batch 的激活值，显存压力大
2. 有了 Checkpointing，就可以用更大的 M 来降低气泡率
3. 相当于用计算时间换来了更低的气泡率和更高的 GPU 利用率

---

七、代码实践：气泡率计算脚本

理论说了这么多，我们写个脚本实际算一算不同配置下的气泡率，找找数字感觉。

"""
流水线并行气泡率计算器
计算 GPipe 和 1F1B 在不同配置下的气泡率
"""


def gpipe_bubble_rate(num_stages, num_micro_batches, fw_bw_ratio=1.0 / 2.0):
    """
    计算 GPipe 的气泡率

    公式: (P-1)/M * (t_fw / (t_fw + t_bw))
    - 因为 forward 全部完成后才开始 backward
    - fw_bw_ratio: t_fw / t_bw，默认 1:2

    推导：
    总时间(以t_fw为单位): M*fw + (P-1)*fw + M*bw + (P-1)*bw = (M+P-1)*(fw+bw)
    气泡时间: (P-1)*fw + (P-1)*bw = (P-1)*(fw+bw)  —— 不对，这是理想情况
    实际 GPipe: forward结束后才backward，有额外的等待

    简化版：气泡率 = (P-1) / M  *  fw/(fw+bw)
    （只有 forward 阶段的 fill/flush 产生气泡，backward 阶段也有，但比例不同）
    """
    # 更准确的计算：考虑 forward 和 backward 各有 (P-1) 的气泡
    fw_time = num_micro_batches + (num_stages - 1)  # 单位: t_fw
    bw_time = num_micro_batches * (1.0 / fw_bw_ratio) + (num_stages - 1) * (1.0 / fw_bw_ratio)

    total_time = fw_time + bw_time
    bubble_time = 2 * (num_stages - 1)  # forward和backward各有P-1的气泡

    return bubble_time * fw_bw_ratio / (total_time * fw_bw_ratio)


def gpipe_bubble_simple(num_stages, num_micro_batches, fw_ratio=0.33):
    """
    简化版 GPipe 气泡率公式
    fw_ratio = t_fw / (t_fw + t_bw)，通常 backward 是 forward 的 2 倍，所以约 0.33
    """
    return (num_stages - 1) / num_micro_batches * fw_ratio


def one_f_one_b_bubble_rate(num_stages, num_micro_batches):
    """
    1F1B 调度的气泡率
    公式: (P-1) / (M + P - 1)

    直觉：流水线有 M+P-1 个"时间槽"，其中 P-1 个是气泡
    """
    return (num_stages - 1) / (num_micro_batches + num_stages - 1)


def interleaved_1f1b_bubble_rate(num_stages, num_micro_batches, num_chunks=2):
    """
    Interleaved 1F1B（交错式）的气泡率估算
    num_chunks: 每个 stage 分成的虚拟块数
    比标准 1F1B 气泡率更低
    """
    # 简化估算：interleaved 让气泡大约减少 num_chunks 分之一
    base_bubble = one_f_one_b_bubble_rate(num_stages, num_micro_batches)
    return base_bubble * (num_stages - 1) / (num_stages * num_chunks - 1)


def print_bubble_table():
    """打印不同配置下的气泡率对比表"""
    print("=" * 70)
    print("不同配置下的流水线气泡率对比 (假设 fw:bw = 1:2)")
    print("=" * 70)
    print(f"{'P(stage)':>8} {'M(mb)':>8} {'GPipe':>10} {'1F1B':>10} {'Interleaved(2x)':>16} {'GPU利用率':>10}")
    print("-" * 70)

    for P in [2, 4, 8, 16]:
        for M in [4, 8, 16, 32, 64]:
            gpipe = gpipe_bubble_simple(P, M)
            f1b1 = one_f_one_b_bubble_rate(P, M)
            interleaved = interleaved_1f1b_bubble_rate(P, M, num_chunks=2)
            util = 1 - f1b1  # 以 1F1B 为准算利用率

            if M == 4:  # 只打印第一行的P标记
                p_str = str(P)
            else:
                p_str = ""

            print(f"{p_str:>8} {M:>8} {gpipe:>9.1%} {f1b1:>9.1%} {interleaved:>15.1%} {util:>9.1%}")
        print()


def recommend_config(num_stages, target_bubble=0.15):
    """
    给定 stage 数和目标气泡率，推荐需要的 micro-batch 数
    基于 1F1B 公式反推: M >= (P-1) / target - (P-1)
    """
    M_min = (num_stages - 1) / target_bubble - (num_stages - 1)
    return int(M_min) + 1


def main():
    # 1. 打印气泡率对比表
    print_bubble_table()

    # 2. 推荐配置
    print("=" * 70)
    print("气泡率 15% 以下所需的最小 Micro-batch 数 (1F1B)")
    print("=" * 70)
    for P in [2, 4, 8, 16]:
        M_min = recommend_config(P, target_bubble=0.15)
        actual_bubble = one_f_one_b_bubble_rate(P, M_min)
        print(f"  P={P:2d}:  M ≥ {M_min:3d}  (实际气泡率: {actual_bubble:.1%})")

    # 3. 一个实际案例
    print("\n" + "=" * 70)
    print("实际案例: 70B 模型训练配置")
    print("=" * 70)
    P = 4  # PP=4
    M = 32  # 32 个 micro-batch
    gpipe = gpipe_bubble_simple(P, M)
    f1b1 = one_f_one_b_bubble_rate(P, M)
    print(f"  配置: PP={P}, M={M}")
    print(f"  GPipe 气泡率: {gpipe:.1%}  → GPU利用率: {1-gpipe:.1%}")
    print(f"  1F1B  气泡率: {f1b1:.1%}  → GPU利用率: {1-f1b1:.1%}")
    print(f"  1F1B 比 GPipe 提升: {(f1b1 - gpipe) / gpipe * 100:.1f}% 气泡减少")

    # 4. 验证一个极端情况
    print("\n" + "=" * 70)
    print("极端情况: M→∞ 时气泡率趋近于多少？")
    print("=" * 70)
    for P in [2, 4, 8]:
        M_large = 10000
        bubble = one_f_one_b_bubble_rate(P, M_large)
        print(f"  P={P}: 气泡率 → {bubble:.4%} (理论极限: 0%, 因M很大P可忽略)")


if __name__ == "__main__":
    main()

运行结果解读：

1. P 越大，气泡率越高——8 stage 的气泡率明显比 2 stage 大很多
2. M 越大，气泡率越低——64 个 micro-batch 的气泡率远小于 4 个
3. 1F1B 比 GPipe 好——尤其是在 M 较大时，1F1B 的激活显存优势更明显
4. Interleaved 进一步降低气泡——通过虚拟 stage 交错，气泡率可以再降一些

自己动手试试：

1. 修改 fw_bw_ratio，看看 forward/backward 时间比例对 GPipe 气泡率的影响
2. 计算 PP=8 时，要达到 10% 气泡率需要多少 micro-batch
3. 如果激活显存有限，只能用 M=8，P=8 时气泡率是多少？能接受吗？

---

八、常见误区与小结

8.1 常见误区

误区一：“PP 通信少就一定比 TP 好”

PP 虽然通信量小，但有气泡问题，GPU 利用率低。而且 PP 切的是层，如果模型不够深（比如只有 12 层），PP 的并行度很有限。实际生产中通常是 TP + PP 组合使用。

误区二：“micro-batch 越多越好，反正气泡率低”

micro-batch 越多，气泡率确实越低，但：

• 激活显存越大（要保存更多 micro-batch 的激活值）
• 调度开销越大（send/recv 次数变多）
• 每个 micro-batch 太小可能导致 GPU 利用率下降

要在气泡率和显存之间找到平衡点。

误区三：“GPipe 已经过时了，没人用了”

GPipe 虽然气泡大，但实现简单，概念清晰，在教学场景和小规模实验中仍然很有价值。而且 GPipe 的激活值管理更简单，在某些特殊场景下可能反而合适。

误区四：“PP 只能跨节点用”

不对。PP 可以在节点内用（比如 8 卡做 PP=8），也可以跨节点用。只是因为 PP 通信量小，跨节点的代价相对低，所以通常跨节点用 PP、节点内用 TP 是更优的组合。

8.2 全文小结

1. 流水线并行（PP） 按层切分模型，数据像流水一样依次经过各个 stage，适合模型很深的场景。
2. GPipe 是最经典的流水线方案，采用 fill-flush 方式，forward 全部完成后再 backward，气泡大但实现简单。
3. 1F1B 调度 让 forward 和 backward 交错执行，减少了气泡，降低了激活显存，是当前主流方案。
4. 气泡率 是 PP 最重要的性能指标，受 stage 数 P、micro-batch 数 M、fw/bw 比例影响。
5. 优化策略：增加 M、用 1F1B/interleaved、减少 P、配合 Activation Checkpointing。
6. 激活值 vs 气泡 是核心权衡——M 大气泡小但显存大，需要根据实际情况取舍。

---

九、思考问题

1. 为什么 PP 的 stage 间通信量比 TP 小？两种并行的通信内容有什么不同？
2. GPipe 和 1F1B 的主要权衡是什么？在显存充裕和显存紧张的情况下，分别应该选哪种？
3. 假设有一个 8-stage 的流水线，使用 1F1B 调度，要让气泡率低于 10%，至少需要多少个 micro-batch？
4. 为什么 1F1B 调度比 GPipe 需要保存的激活值更少？两者的激活值显存分别和什么成正比？
5. 你有一个 64 层的 Transformer 模型，8 个节点（每节点 8 卡），如何设计 PP stage 划分？需要考虑哪些因素？

---

下一篇预告：当模型大到需要 TP、PP、DP 三种并行一起上的时候，就进入了 3D 并行的世界。我们将学习如何把三种并行方式优雅地组合起来，再加上混合精度、梯度累积、Activation Checkpointing 这些"辅助技能"，打造高效的大规模训练系统。

转载自 CSDN-专业IT技术社区

原文链接：https://blog.csdn.net/bumblebee16/article/details/167128921