图解分布式训练（一）—— 流水线并行（Pipeline Parallelism）
来自：AiGC面试宝典
宁静致远 2023年09月29日 11:02
为什么需要流水线并行（Pipeline Parallelism）？
回顾ChatGPT的发展历程，我们可以总结出大语言模型（LLM）取得惊艳效果的要点（重要性从高到低排序）：
- 愿意烧钱，且接受"烧钱 ≠ 好模型"的现实
- 高质量的训练语料
- 高效的分布式训练框架和充沛优质的硬件资源
- 算法的迭代创新
在大模型训练这个系列里，我们将一起探索学习几种经典的分布式并行范式，包括流水线并行（Pipeline Parallelism），数据并行（Data Parallelism）和张量并行（Tensor Parallelism）。
📝 通俗解释：分布式训练就是让多个GPU同时工作，共同完成模型训练任务。就像一个工程项目，如果只让一个人做会很慢，但如果把工作分给多个人同时做，就会快很多。
微软开源的分布式训练框架DeepSpeed，融合了这三种并行范式，开发出3D并行的框架，实现了千亿级别模型参数的训练。
📝 通俗解释：DeepSpeed是微软开发的"神器"，它把三种不同的并行方法组合在一起，就像把汽车、飞机、轮船的优点结合在一起一样强大，能够训练超大规模的AI模型。
本篇文章将探索流水线并行，经典的流水线并行范式有Google推出的GPipe，和微软推出的PipeDream。两者的推出时间都在2019年左右，大体设计框架一致。主要差别为：在梯度更新上，GPipe是同步的，PipeDream是异步的。异步方法更进一步降低了GPU的空转时间比。虽然PipeDream设计更精妙些，但是GPipe因为其"够用"和浅显易懂，更受大众欢迎（PyTorch的PP接口就基于GPipe）。因此本文以GPipe作为流水线并行的范例进行介绍。
📝 通俗解释：GPipe和PipeDream是两种不同的流水线并行方案。简单理解：同步就像所有人等所有人都完成后才继续，异步就像各干各的、做完就继续。PipeDream更高效但实现更复杂，GPipe虽然性能稍差但更容易理解和使用。
本文内容包括：
- 优化目标
- 模型并行
- 流水线并行
- 切分micro-batch
- Re-materialization（active checkpoint）
- 实验效果
- GPU数量 VS 模型大小
- GPU数量 VS 训练速度
- GPipe下GPU的时间消耗分析
一、流水线并行（Pipeline Parallelism）优化目标是什么？
当你从单卡"穷人"变成多卡"富翁"时，你做分布式训练的总体目标是什么呢？ （虽然手握一张A100怎么可能是"穷"呢）
- 能训练更大的模型。理想状况下，模型的大小和GPU的数量成线性关系。即GPU量提升x倍，模型大小也能提升x倍。
- 能更快地训练模型。理想状况下，训练的速度和GPU的数量成线性关系。即GPU量提升x倍，训练速度也能提升x倍。
📝 通俗解释：分布式训练有两个核心目标：一是让更多GPU能训练更大的模型（像扩大仓库容量），二是让更多GPU能更快完成训练（像增加员工提高效率）。理想情况下，GPU翻倍，模型容量和训练速度都应该翻倍。
这是目标，也是难点，难在于：
- 内存压力：训练更大的模型时，每块GPU里不仅要存模型参数，还要存中间结果（用来做Backward）。而更大的模型意味着需要更多的训练数据，进一步提高了中间结果的大小。加重了每块GPU的内存压力。我们将在下文详细分析这一点。（对应着GPU中的内存限制）
- 通讯开销：数据在卡之间进行传输，是需要通讯时间的。不做设计的话，这个通讯时间可能会抹平多卡本身带来的训练速度提升。（对应着GPU间的带宽限制）
📝 通俗解释：GPU内存就像手机的运行内存，内存不够就会崩溃。而GPU之间的数据传输就像发微信消息，需要时间，如果设计不好，等待消息的时间反而比真正做事的时间还多，就亏大了。
明确这两个训练目标后，我们来看并行范式的设计者，是如何在现有硬件限制的条件下，完成这两个目标的。
二、图解模型并行（Model Parallelism）的必要性
- 动机：当你有一个单卡装不下的大模型时
- 解决办法：把模型切成不同的层，每一层都放到一块GPU上，如下图：
![图表描述：图片展示了一个垂直堆叠的模型结构被切分到了不同的GPU上。最下方是"训练数据X"，箭头向上指向GPU0，GPU0上方有箭头指向GPU1。GPU0和GPU1分别承载了模型的不同层（用矩形框表示），展示了模型层在不同GPU间的分配。]
📝 通俗解释：想象一个超大的书架（模型），一本书（一层）根本放不下。那么解决方案就是：把书架分成几部分，每部分放在不同的房间里（GPU），大家分工合作存放书籍。
此时，模型做一轮forward和backward的过程如下：
![图表描述：流水线并行示意图。这是一个甘特图，横轴表示时间Time。纵轴隐含表示不同的GPU层级。图中有蓝色的F0块（Forward）呈阶梯状上升，红色的B0块（Backward）呈阶梯状下降。最右侧有一列黄色的update块。具体结构为：
- 第一列（最左下）：F0
- 第二列：F0
- 第三列：F0, B0 (顶部)
- 第四列：B0
- 第五列：B0
- 第六列：B0
- 最右侧：四个update块垂直排列]
其中下标表示batch编号，这里只有一个batch，因此下标都是0。每一行表示一个GPU。每一列表示timestep。
📝 通俗解释：这张图展示了GPU的工作流程。蓝色是"前向传播"（计算输出），红色是"反向传播"（计算梯度），黄色是"更新参数"。可以看到GPU0先开始工作，然后GPU1、GPU2、GPU3依次跟上，全部完成后统一更新参数。就像工厂的流水线，但早期设计比较简陋。
这张图的含义是：我在GPU0上做完一次forward，然后将GPU0上最后一层的输入传给GPU1，继续做forward，直到四块GPU都做完forward后，我再依次做backward。等把四块GPU上的backward全部做完后，最后一个时刻我统一更新每一层的梯度。
这样做确实能训更大的模型了，但也带来了两个问题：
问题1：GPU利用率不够
![图表描述：Bubble示意图。这是一个带有阴影区域的甘特图。横轴是Time。图中有F₀（Forward）和B₀（Backward）块分散在不同位置。大部分背景区域被画上了斜线阴影，表示bubble（空转时间）。右侧有一列Update块。]
📝 通俗解释：图中的阴影区域就是"Bubble"（气泡），表示GPU在"等待"的空闲时间。可以看到大部分时间都有GPU在无所事事地等着别人完成工作，这就像工厂里有人很忙，有人却在发呆。
如图，阴影部分所表示的时间段里，总有GPU在空转。在GPipe中，将阴影部分定义为bubble。我们来计算一下bubble。假设有 $K$ 块GPU，而单块GPU上做一次forward和backward的时间为：
$$ t_{fb} = (t_f + t_b) $$
那么：
- 图中灰色长方形的整体面积为：$K \times K \times t_{fb}$ （宽 = K，长 = $K \times t_{fb}$）
- 图中实际在做forward和backward的面积为：$K \times t_{fb}$
- 图中阴影部分的面积为：$K \times K \times t_{fb} - K \times t_{fb} = (K-1) \times K \times t_{fb}$
- 图像阴影部分的占比为：$(K-1) \times K \times t_{fb} / (K \times K \times t_{fb}) = (K-1) / K$
📝 通俗解释：这个公式算的是GPU有多少时间在"摸鱼"。如果只有2块GPU，50%的时间在摸鱼；如果有8块GPU，87.5%的时间在摸鱼！GPU越多，浪费越严重。
则我们定义出bubble部分的时间复杂度为：
$$ O\left(\frac{K-1}{K}\right) $$
当K越大，即GPU的数量越多时，空置的比例接近1，即GPU的资源都被浪费掉了。因此这个问题肯定需要解决。
📝 通俗解释：这就像一个旅游团，人越多，等齐所有人的时间就越长，大部分时间都浪费在等待上了。所以必须改进！
问题2：中间结果占据大量内存
![图表描述：显存占用示意图。展示了Forward和Backward过程中的显存变化。左侧有一个向上的粉色箭头标记为Forward，右侧有一个向下的粉色箭头标记为Backward。中间有两个GPU层级：GPU0和GPU1。底部有一个绿色块标记为"训练数据X"。GPU0和GPU1之间有浅蓝色方块表示中间结果，箭头指示数据流向，标记为Z。]
📝 通俗解释：Forward过程中，每一层的计算结果（Z）都需要保存下来，因为Backward时需要用到这些中间结果来计算梯度。这就像解题时需要保留每一步的计算过程，不然最后没法回溯检查。
在做backward计算梯度的过程中，我们需要用到每一层的中间结果z。假设我们的模型有L层，每一层的宽度为d，则对于每块GPU，不考虑其参数本身的存储，额外的空间复杂度为：
$$ O\left(N \times \frac{L}{K} \times d\right) $$
📝 通俗解释：这个公式表示每块GPU需要保存的中间结果大小。N是数据量，L是模型层数，d是每层宽度，K是GPU数量。模型越大、数据越多，需要的显存就越大。
从这个复杂度可以看出，随着模型的增大，N，L，d三者的增加可能会抵消掉K增加带来的GPU内存收益。因此，这也是需要优化的地方。
📝 通俗解释：虽然GPU数量增加了，但模型也在变大，中间结果也在变大，就像你买了更大的仓库，但货物也在增多，仓库压力并没有减轻多少。
三、流水线并行（Pipeline Parallelism）图解
朴素的模型并行存在GPU利用率不足，中间结果消耗内存大的问题。 而GPipe提出的流水线并行，就是用来解决这两个主要问题的。
3.1 切分micro-batch
核心思想：在模型并行的基础上，进一步引入数据并行的办法，即把原先的数据再划分成若干个batch，送入GPU进行训练。 未划分前的数据，叫mini-batch。在mini-batch上再划分的数据，叫micro-batch。
📝 通俗解释：原来是把一个超大的任务（整个mini-batch）一次性给所有GPU。现在换一种思路：把这个大任务切成许多小块（micro-batch），像流水线一样依次送入GPU。这样GPU就不会长时间空闲了。
图例如下：
![图表描述：Gpipe流水线并行执行过程时序图]
- 纵轴：表示不同的计算阶段（stage1, stage2, stage3, stage4），通常对应不同的GPU。
- 横轴：表示时间流逝。
- 方块：代表计算任务。蓝色方块标记为F（Forward，前向传播），橙色方块标记为B（Backward，反向传播）。例如，F11表示stage1处理第1个micro-batch的前向计算，B44表示stage4处理第4个micro-batch的反向计算。
- 箭头：表示数据依赖和传递方向。Forward阶段数据从stage1流向stage4；Backward阶段梯度从stage4流回stage1。
- Bubble：图中间有一个黄色区域标记为"bubble"，表示GPU在等待依赖数据完成时的空闲时间（气泡）。
📝 通俗解释：这张图展示了"流水线"的运作方式。就像工厂的装配线，不同的GPU处理不同的任务批次。GPU0正在处理第1个任务时，GPU1已经在处理第2个了。这样GPU就不会闲着，大家都有活干。
其中，第一个下标表示GPU编号，第二个下标表示micro-batch编号。假设我们将mini-batch划分为M个，则流水线并行下，bubble的时间复杂度为：
$$ O\left(\frac{K-1}{K+M-1}\right) $$
📝 通俗解释：这个公式是改进后的气泡占比。K是GPU数量，M是micro-batch数量。当M很大时，分母(K+M-1)会远大于(K-1)，气泡占比就变得很小了。
注：（推导过程略，可参照第二部分的bubble推导流程）
GPipe通过实验证明，当 M≥4K 时，bubble产生的空转时间占比对最终训练时长影响是微小的，可以忽略不计。
📝 通俗解释：经验法则：把任务分成足够多的小块（比如GPU数量的4倍以上），就能让GPU几乎一直有活干，空闲时间可以忽略不计。
将batch切好，并逐一送入GPU的过程，就像一个流水生产线一样（类似于CPU里的流水线），因此也被称为Pipeline Parallelism。
📝 通俗解释：这就是"流水线并行"名字的由来——像工厂流水线一样，把工作分成多个工序，每个GPU负责一道工序，数据像产品一样在流水线上流动。
3.2 Re-materialization（Active Checkpoint）
动机：虽然切分micro-batch方法能够解决GPU的空置问题，提升了GPU计算的整体效率。但是如何解决GPU的内存问题呢？前文说过，随着模型的增加，每块GPU中存储的中间结果也会越大。
📝 通俗解释：解决了"人闲"的问题，现在要解决"仓库满"的问题。中间结果占用太多内存，GPU也会"爆仓"。
核心方法：GPipe采用了一种非常简单粗暴但有效的办法：用时间换空间，在论文里，这种方法被命名为re-materialization，后人也称其为active checkpoint。
📝 通俗解释：这个方法很聪明：与其保存所有中间结果（占内存），不如不保存，等到需要用的时候再重新算一遍。这就像你不记笔记，考试时现推导公式——虽然多花了时间，但省了记笔记的纸张。
具体来说，就是几乎不存中间结果，等到backward的时候，再重新算一遍forward，图例如下：
![图表描述：Re-materialization机制示意图]
- 整体流程：左侧粉色向上箭头表示Forward（前向）过程，右侧粉色向下箭头表示Backward（反向）过程。
- 底部输入：有三个绿色的框，标记为"micro-batch0"，表示输入数据。
- 中间计算层：展示了GPU0和GPU1上的计算块（浅蓝色矩形）。
- 内存策略标注：
- Z（算完即废）：在Forward过程中，大部分中间层的输出结果（Z）被标记为"算完即废"，意味着不保存以节省显存。
- Z（保存）：只有特定的中间结果（通常是每一段/每一块的边界输入）被标记为"保存"。
- 含义：在Backward阶段，如果需要之前被丢弃的中间结果Z，系统会利用保存下来的Z重新进行一次局部的Forward计算来恢复它，从而用计算时间换取了显存空间。
📝 通俗解释：图中展示了"选择性记忆"的策略。只保存每个阶段开始的输入（关键checkpoint），中间的计算结果用完就扔。需要用的时候，再用保存的输入重新算一遍。这就是用CPU计算时间换取GPU内存空间。
每块GPU上，我们只保存来自上一块的最后一层输入z，其余的中间结果我们算完就废。等到backward的时候再由保存下来的z重新进行forward来算出。
📝 通俗解释：类似于"关键中间点"的思路。比如你要从北京去广州开会，沿途经过石家庄、郑州、武汉等站点。如果你只记住每个站点的位置（checkpoint），而不是记住整条路的所有细节（中间结果），就能大大减轻记忆负担。需要回溯时，从最近的关键点重新走一遍就行。
现在我们来计算每块GPU峰值时刻的内存：
每块GPU峰值时刻存储大小 = 每块GPU上的输入数据大小 + 每块GPU在forward过程中的中间结果大小
每块GPU上固定需要保存它的起始输入，我们记起始输入为N（即mini-batch的大小）。
每个micro-batch是流水线形式进来的，算完一个micro-batch才算下一个。在计算一个micro-batch的过程中，我们会产生中间变量，它的大小为：
$$ \frac{N}{M} \times \frac{L}{K} \times d $$
📝 通俗解释：N/M就是每个micro-batch的数据量。把大batch分成M份后，每个micro-batch产生的中间结果只有原来的1/M，显存压力大大减轻。
注：（其中M为micro-batch个数）
因此，每块GPU峰值时刻的空间复杂度为：
$$ O\left(N + \frac{N}{M} \times \frac{L}{K} \times d\right) $$
与朴素模型并行中的GPU空间复杂度：
$$ O\left(N \times \frac{L}{K} \times d\right) $$
相比，可以发现，由于采用了micro-batch的方法，当L变大时，流水线并行相比于朴素模型并行，对GPU内存的压力显著减小。
📝 通俗解释：对比两个公式：原来的内存需求是N×L/K×d，现在是N + N/M×L/K×d。当模型很大（L很大）时，第二项虽然也会大，但比第一项小很多。特别是N这个固定开销被分离出来了，不会随着模型增大而无限增长。
如果你使用PyTorch提供的pipeline接口，其中有一个参数叫checkpoint，就是用来做这一项的。
![图表描述：PyTorch Pipe API文档截图] 图片展示了PyTorch中Pipe APIs in PyTorch的文档截图。 代码片段显示： CLASS torch.distributed.pipeline.sync.Pipe(module, chunks=1, checkpoint='except_last', deferred_batch_norm=False) 其中checkpoint='except_last'被红框圈出。
📝 通俗解释：PyTorch已经内置了这个功能！只需要设置
checkpoint='except_last'，就会自动使用re-materialization技术，省去手动实现的麻烦。
文档说明Pipe结合了pipeline parallelism与checkpointing以减少训练所需的峰值内存，同时最小化设备利用率不足。
最后，再提一点，在micro-batch的划分下，我们在计算Batch Normalization时会有影响。GPipe的方法是，在训练时计算和运用的是micro-batch里的均值和方差，但同时持续追踪全部mini-batch的移动平均和方差，以便在测试阶段进行使用。Layer Normalization则不受影响。
📝 通俗解释：Batch Normalization需要计算整个batch的统计信息，但被切成多个micro-batch后，每个micro-batch的统计量不一样。解决方案是：训练时用micro-batch的统计量，同时维护一个全局移动平均供测试时使用。这就像你每节课都考试（micro-batch），同时也记录期中期末成绩（全局平均）用于最终评级。
四、流水线并行（Pipeline Parallelism）优缺点
回顾第二部分的两个目标，GPipe真的实现了吗？如果实现不了，又是因为什么原因呢？我们来看下实验效果。
4.1 GPU数量 VS 模型大小
| NVIDIA GPUs (8GB each) | Naive-1 | Pipeline-1 | Pipeline-2 | Pipeline-4 | Pipeline-8 |
|---|---|---|---|---|---|
| AmoebaNet-D (L, D) | (18, 208) | (18, 416) | (18, 544) | (36, 544) | (72, 512) |
| # of Model Parameters | 82M | 318M | 542M | 1.05B | 1.8B |
| Total Model Parameter Memory | 1.05GB | 3.8GB | 6.45GB | 12.53GB | 24.62GB |
| Peak Activation Memory | 6.26GB | 3.46GB | 8.11GB | 15.21GB | 26.24GB |
| Cloud TPUv3 (16GB each) | Naive-1 | Pipeline-1 | Pipeline-8 | Pipeline-32 | Pipeline-128 |
|---|---|---|---|---|---|
| Transformer-L | 3 | 13 | 103 | 415 | 1663 |
| # of Model Parameters | 282.2M | 785.8M | 5.3B | 21.0B | 83.9B |
| Total Model Parameter Memory | 11.7GB | 8.8GB | 59.5GB | 235.1GB | 937.9GB |
| Peak Activation Memory | 3.15GB | 6.4GB | 50.9GB | 199.9GB | 796.2GB |
GPipe分别在AmoebaNet（图像）和Transformer（自然语言）两个大模型上做了实验。
- Naive表示单卡
- Pipeline-N表示re-materialization + N卡
- AmoebaNet-D和Transformer-L一行表示超参数的量
of Model Parameters表示模型的参数量
- Total Model Parameter Memory表示模型参数所占内存大小
- Peak Activation Memory表示峰值时中间结果大小。可以发现，中间结果占据的内存大小是相当可观的。
📝 通俗解释：这个表格展示了用更多GPU能训练多大的模型。可以看到，随着GPU数量增加，能训练的模型参数量也在增加。特别是Transformer模型，从32卡到128卡，模型从210亿参数增加到839亿参数，几乎是4倍关系！
从实验结果里，我们可以发现：
- 在Transformer上，GPipe基本实现了模型大小（参数量）和GPU个数之间的线性关系。例如从32卡增到128卡时，模型的大小也从21.0B增加到82.9B，约扩4倍
- 对AmoebaNet而言，却没有完全实现线性增长。例如从4卡到8卡，模型大小从1.05B到1.8B，不满足2倍的关系。本质原因是AmoebaNet模型在切割时，没有办法像Transformer一样切得匀称，保证每一块GPU上的内存使用率是差不多的。因此对于AmoebaNet，当GPU个数上升时，某一块GPU可能成为木桶的短板。
📝 通俗解释：Transformer效果好是因为它的结构很规整，像切豆腐一样可以均匀分成多份。AmoebaNet结构不规则，切完后有的GPU很轻松，有的GPU很吃力，整体效率就上不去。这就像让不同体质的人搬砖，有人已经累趴了，有人还很轻松。
4.2 GPU数量 VS 训练速度
4.2.1 关掉NVLinks
为了验证GPipe框架带来的收益，实验中关掉了NVLinks（GPU间快速通信的桥梁。估计是通过强迫GPU先连CPU然后再连别的GPU做到的）。关掉的意义在于说明，不靠硬件本身的高效通讯带来的收益，GPipe一样能做的很好。实验效果如下：
| AmoebaNet | Transformer | |||||
|---|---|---|---|---|---|---|
| GPU | K: 2 | 4 | 8 | K: 2 | 4 | 8 |
| M=32 | 1 | 1.7 | 2.7 | 1 | 1.8 | 3.3 |
M=32表示micro-batch的数量为32，K表示GPU数量。从实验结果可知，在关掉NVLinks的情况下，GPipe一样也能实现随着GPU数量的增加，训练速度也增加的效果。虽然这两者间不是线性的。同样，因为模型切割不均的原因，AmoebaNet的表现不如Transformer。
📝 通俗解释：关掉NVLinks就像把高速公路换成普通公路，看看在普通硬件条件下GPipe表现如何。结果发现，即使通讯变慢了，GPipe依然能让训练速度随GPU数量增加而提升，说明这个方法确实有效。
4.2.2 开启NVLinks，并寻找最佳M
| AmoebaNet | Transformer | |||||
|---|---|---|---|---|---|---|
| TPU | K: 2 | 4 | 8 | K: 2 | 4 | 8 |
| M=1 | 1 | 1.13 | 1.38 | 1 | 1.07 | 1.3 |
| M=4 | 1.07 | 1.26 | 1.72 | 1.7 | 3.2 | 4.8 |
| M=32 | 1.21 | 1.84 | 3.48 | 1.8 | 3.4 | 6.3 |
当重新开启NVLinks后，我们来看M的大小（即流水线的核心）对训练速度的影响。
📝 通俗解释：这个表格展示了micro-batch数量M对训练速度的影响。M太小，GPU会空闲；M太大，虽然效率高但管理开销也大。需要找到最佳平衡点。
- 当M=1的时候，如前文所说，GPU的空置率太高，因此两个模型都没有实现训练速度和GPU个数间的线性关系
- 当M=4时，表现明显好转
- 当M=32时，表现最佳，且Transformer基本实现了训练速度和GPU个数的线性关系
📝 通俗解释：M: 1就是原始的朴素方法，效率很低；M: 32时效果最好，训练速度几乎和GPU数量成正比。这验证了之前的理论分析：M越大，bubble越小，效率越高。
4.2.3 GPipe下时间消耗分布
![饼图描述：时间消耗分布饼图]
- Computation（计算）: 65.6% （蓝色区域，占比最大）
- Weight Update（权重更新）: 22.5% （黄色区域）
- Recompute（重计算）: 3.2% （绿色区域）
- Load Imbalance（负载不均衡）: 4.1% （粉色区域）
- Bubble Overhead（气泡开销）: 2.9% （青色区域）
- Setup Overhead（设置开销）: 1.7% （深蓝色点）
📝 通俗解释：这张饼图展示了GPU时间都花在哪了。将近2/3的时间（65.6%）用在真正的计算上，这是正常的。重计算（re-materialization）只占3.2%，说明用时间换空间的方法开销不大。气泡开销只有2.9%，说明流水线已经很好地解决了GPU空闲问题。
对每块GPU来说，约2/3的时间，是真正花在计算上的。 其余1/3的时间，大部分花在re-materialization策略下的重计算上。因为采用流水线的方法，bubble的时间也被压缩到很短，可以忽略不计。
总结
本文详细介绍了流水线并行（Pipeline Parallelism）的基本概念和经典实现GPipe。通过以下两种核心方法：
- 切分micro-batch：将大batch切成多个小batch，让GPU始终有活干，解决GPU空闲问题
- Re-materialization（Active Checkpoint）：用时间换空间，不保存中间结果，需要时重新计算，解决内存压力问题
GPipe基本实现了模型大小和训练速度与GPU数量的线性关系，是大模型分布式训练的重要基石。
📝 通俗解释：流水线并行就像现代化工厂的装配线：把大任务拆成小任务（micro-batch），每个工人（GPU）只负责一部分；同时只保留关键节点的信息（checkpoint），省下空间存储。这样既提高了效率，又节省了资源。