流水线并行原理与Gpipe详解

内容整理自学习笔记，仅供面试备考参考；不构成录用、培训或考试承诺。

1. 为什么需要流水线并行？

大语言模型取得惊艳效果的关键要素（按重要性排序）：

1. 愿意烧钱，且接受"烧钱 ≠ 好模型"的现实
2. 高质量的训练语料
3. 高效的分布式训练框架和充沛优质的硬件资源
4. 算法的迭代创新

分布式训练的两个核心目标：

目标	理想效果	实际难点
训练更大的模型	模型大小与GPU数线性增长	每卡不仅存参数还需存中间结果，大模型需要更多数据进一步加大中间结果
更快训练模型	训练速度与GPU数线性增长	网络通讯开销可能抹平多卡带来的速度提升

2. 模型并行及其必要性

动机： 当单卡装不下大模型时

方案： 把模型隔成不同的层，每层放到一块GPU上：

GPU0: Layer 0~2 (前3层)
GPU1: Layer 3~5 (后3层)

模型并行下的前向+反向过程：GPU0做forward → 传给GPU1继续forward → 4块GPU依次forward → 再依次backward → 统一更新梯度。

3. 朴素模型并行的两个问题

3.1 GPU利用度不足（Bubble问题）

阴影时间段里总有GPU空转，称为bubble。

假设K块GPU，单块GPU一次forward+backward时间为 $t_{fd}+t_{bd}$：

• bubble面积占比随K增大接近1，GPU资源被浪费
• 必须解决空转问题

3.2 中间结果占据大量内存

backward计算梯度需要用到每层中间结果z。假设模型L层、宽度d，每块GPU额外空间复杂度为 $O(L\cdot d\cdot N/K)$，其中N为mini-batch大小。

随着模型增大（N、L、d增加），K增加带来的GPU内存收益被平滑掉。

4. Gpipe流水线并行方案

4.1 切分micro-batch

核心思想： 在模型并行基础上引入数据并行，将mini-batch划分为M个micro-batch，逐一送入GPU。

GPU0_0 → GPU1_0 → GPU2_0 → GPU3_0   (micro-batch 0)
GPU0_1 → GPU1_1 → GPU2_1 → GPU3_1   (micro-batch 1)
GPU0_2 → GPU1_2 → GPU2_2 → GPU3_2   (micro-batch 2)
...

Bubble优化效果：

• 划分M个micro-batch后，bubble时间复杂度降低
• 当 $M\ge 4K$ 时，bubble产生的空转时间占比可忽略不计

将batch切好并逐一送入GPU，类似CPU流水线，因此称为Pipeline Parallelism。

4.2 Re-materialization（Activation Checkpoint）

动机： micro-batch解决了GPU空转问题，但中间结果仍然占用大量内存。

核心方法： 用时间换空间——几乎不存中间结果，backward时重新计算forward。

具体做法：

• 每块GPU只保存来自上一块的最后一层输入z
• 其余中间结果算完即丢弃
• backward时由保存的z重新forward恢复所需结果

空间复杂度对比：

方案	每块GPU峰值空间复杂度
朴素模型并行	$O(L\cdot d\cdot N/K)$
流水线并行+checkpoint	$O(d\cdot N/M+d\cdot N/M)$

L变大时，流水线并行对GPU内存的压力显著小于朴素模型并行。

# PyTorch中使用pipeline checkpoint参数
# pipeline接口中的checkpoint参数即用于此功能

BatchNorm注意： micro-batch划分下，训练时计算micro-batch的均值和方差，同时追踪mini-batch的移动平均和方差用于测试阶段。LayerNorm不受影响。

5. 流水线并行实验效果

5.1 GPU数量 vs 模型大小

Gpipe在AmoebaNet（图像）和Transformer（NLP）上实验：

模型	线性增长效果
Transformer	基本实现线性增长（32→128卡，模型21.08B→82.9B约4倍）
AmoebaNet	未完全线性（4→8卡，1.05B→1.8B，不满足2倍）

原因： AmoebaNet切割不均匀，某块GPU可能成为木桶短板。

5.2 GPU数量 vs 训练速度

M值	训练速度与GPU数的线性关系
M=1	GPU空置率太高，未实现线性
M=4	明显好转
M=32	表现最佳，Transformer基本实现线性

5.3 Gpipe下单GPU时间消耗分布

时间占比	内容
~2/3	真正的计算
~1/3	re-materialization的重计算
极短	bubble空转时间（可忽略）

6. 流水线并行优缺点总结

维度	优点	缺点
显存	相比朴素模型并行大幅降低中间结果占用	仍需存储检查点输入
速度	micro-batch大幅减少bubble空转	re-materialization引入额外计算时间
通信	使用P2P通信，成本最低	切割不均匀时某些GPU成为瓶颈
通用性	模型切割需均匀才能最大化效率	不适用于层结构不均匀的模型