大模型训练（1）：流水线并行
背景
本文学习探索流水线并行，经典的流水线并行范式有：
- Google推出的Gpipe（Gpipe是同步的梯度更新，因为其“够用”和浅显易懂，更受大众欢迎，torch的pp接口就基于Gpipe）
- 微软推出的PipeDream（PipeDream是异步的梯度更新，设计更精妙些，更进一步降低了GPU的空转时间比）
两者的推出时间都在2019年左右，因此本文以Gpipe作为流水线并行的范例进行介绍。
1 优化目标
当你从单卡变成多卡时，如何使用多卡进行分布式训练，做分布式训练的总体目标是什么呢？
- 能训练更大的模型：理想状况下，训练速度不变的情况下，模型的大小和GPU的数量成线性关系。即GPU量提升x倍，模型大小也能扩大x倍
- 能更快地训练模型：理想状况下，训练的速度和GPU的数量成线性关系。即GPU数量提升x倍，训练速度也能提升x倍
这是目标，也是难点，难在于：
- 内存压力： 训练更大的模型时，每块GPU里不仅要存模型参数，还要存中间结果（用来做Backward）。而更大的模型意味着需要更多的训练数据，进一步提高了中间结果的大小。加重了每块GPU的内存压力。我们将在下文详细分析这一点。（对应着GPU中的内存限制）
- 网络通讯开销： 数据在卡之间进行传输，需要通讯时间和足够的通信带宽。不做设计的话，这个通讯时间可能会抹平多卡本身带来的训练速度提升。（对应着GPU间的带宽限制）
明确这两个训练目标和两个训练难点后，我们来看并行范式的设计者，是如何在现有硬件限制的条件下，完成这两个目标的。
2 模型并行
当你有一个单卡装不下的大模型时，一个直接的解决办法是，把模型隔成不同的层，每一层都放到一块GPU上，如下图：
此时，模型做一轮forward和backward的过程如下：
- 每一行表示一个GPU
- 每一列表示timestep 即时间戳
这张图的含义是：在GPU0上做完一次forward，然后将GPU0上最后一层的输入传给GPU1，继续做forward，直到4块GPU都做完forward后，再依次做backward。等把四块GPU上的backward全部做完后，最后一个时刻统一更新每一层的梯度。这样做确实能训更大的模型了，但也带来了两个问题：
2.1 GPU利用度不够
- 箭头部分所表示的时间段里，代表该GPU在空转。在Gpipe中，将阴影部分定义为bubble。我们来计算一下bubble。假设有 K K K块GPU，而单块GPU上做一次forward和backward的时间为： T F B = ( T F + T B ) T_{FB}=(T_F+T_B) TFB=(TF+TB)
- 图中整体面积为： K 个 G P U ∗ ( K ∗ T F B ) K个GPU*(K∗T_{FB}) K个GPU∗(K∗TFB)（宽： K K K，长： K ∗ T F B K*T_{FB} K∗TFB）
- 图中实际在做forward和backward的面积为： K ∗ T F B K*T_{FB} K∗TFB
- 图中阴影部分的面积为： K ∗ K ∗ T F B − K T F B = ( K − 1 ) K ∗ T F B K∗K*T_{FB}−KT_{FB}=(K−1)K*T_{FB} K∗K∗TFB−KTFB=(K−1)K∗TFB
- 图像阴影部分的占比为： ( K − 1 ) K T F B / K K T F B = K − 1 K (K−1)KT_{FB}/KKT_{FB}=\frac{K−1}{K} (K−1)KTFB/KKTFB=KK−1
则我们定义出bubble比例： K − 1 K \frac{K-1}{K} KK−1，有效计算比例： 1 K \frac{1}{K} K1，从这里可以发现，有效计算比例*GPU数量=1，即为增加了更多GPU后，训练时间没有变化，因为结合利用率后的等效GPU数量为1。当K越大，即GPU的数量越多时，空置的比例接近1，即GPU的资源都被浪费掉了。因此这个问题肯定需要解决。
2.2 中间结果占据大量内存
在做backward计算梯度的过程中，
- 需要用到每一层的中间结果 Z Z Z
- 模型有 L L L层，每一个GPU上需要完成 L K \frac{L}{K} KL层的计算
- 每一层的宽度为 D D D，每一个GPU上都完成 N N N的micro-batchsize
- 对于每块GPU，不考虑其参数本身的存储，额外的空间复杂度为 O ( N ∗ L K ∗ D ) O(N∗\frac{L}{K}∗D) O(N∗KL∗D)
- 从收益上看，模型被切分到 K K K个GPU，模型需要的内存变小
- 但是，从上面复杂度可以看出，随着模型的增大，
N
N
N，
L
L
L，
D
D
D三者的增加可能会平滑掉
K
K
K增加带来的GPU内存收益。因此，这也是需要优化的地方。
3 流水线并行
朴素的模型并行存在GPU利用度不足，中间结果消耗内存大的问题。而Gpipe提出的流水线并行，就是用来解决这两个主要问题的。
3.1 切分micro-batch
流水线并行的核心思想是：在模型并行的基础上，进一步引入数据并行的办法，即把原先的数据再划分成若干个更小的batch，分批送入GPU进行训练，算一部分就让下一个GPU动起来。未划分前的数据，叫mini-batch。在mini-batch上再划分的数据，叫micro-batch。
图例如下：
- 不同颜色代表不同的GPU1-4，所完成的任务分别是Layer1-4
- 第二个数字BX表示micro-batch编号
- 假设我们将mini-batch划分为
M
M
M份（注意这里是图中为4个），则流水线并行下
- 总面积： K ∗ ( M + K − 1 ) ∗ T F B K*(M+K-1)*T_{FB} K∗(M+K−1)∗TFB
- 有效计算面积： K ∗ M ∗ T F B K*M*T_{FB} K∗M∗TFB
- bubble面积： ( K − 1 ) ∗ T F B (K-1)*T_{FB} (K−1)∗TFB
- bubble比例： O ( K − 1 M + K − 1 ) O(\frac{K−1}{M+K−1}) O(M+K−1K−1) 。
- Gpipe通过实验证明，当 M ≥ 4 K M\geq4K M≥4K 时，bubble产生的空转时间占比对最终训练时长影响是微小的，可以忽略不计
将batch切好，并逐一送入GPU的过程，就像一个流水生产线一样（类似于CPU里的流水线），因此也被称为Pipeline Parallelism
3.2 re-materialization（active checkpoint）
解决了GPU的利用率问题，提升GPU计算的整体效率。接下来，就要解决GPU的内存问题了。前文说过，随着模型的增加，每块GPU中存储的中间结果也会越大。对此，Gpipe采用了一种非常简单粗暴但有效的办法：用时间换空间，在论文里，这种方法被命名为re-materalization，后人也称其为active checkpoint。
具体来说，就是几乎不存中间结果，等到backward的时候，再重新算一遍forward，图例如下：
每块GPU上，我们只保存来自上一块的最后一层输入 Z Z Z，其余的中间结果我们算完就废。等到backward的时候再由保存下来的 Z Z Z重新进行forward来算出。
现在我们来计算每块GPU峰值时刻的内存：
每块GPU峰值时刻存储空间 = 每块GPU上的输入数据 + 每块GPU在forward过程中的中间结果
- mini-batch的大小 N N N（注意这里是batchsize的大小）
- mini-batch切分为 M M M个micro-batch（注意这里是份数），每个micro-batch是流水线形式进来的，单卡上算完一个micro-batch才算下一个。
- 在计算一份完整micro-batch的过程中，产生中间变量，大小为 N M ∗ L K ∗ D \frac{N}{M}∗\frac{L}{K}∗D MN∗KL∗D
- 每块GPU上固定需要保存它的起始输入 N ∗ D N*D N∗D，每块GPU峰值时刻的空间复杂度为 N ∗ D + N M ∗ L K ∗ D N*D+\frac{N}{M}∗\frac{L}{K}∗D N∗D+MN∗KL∗D
将其与朴素模型并行（不使用时间换空间）的GPU空间复杂度
N
∗
L
K
∗
D
N∗\frac{L}{K}∗D
N∗KL∗D 比较
N
∗
D
+
N
M
∗
L
K
∗
D
N
∗
L
K
∗
D
=
K
+
1
M
∗
L
L
=
K
L
+
1
M
\frac{N*D+\frac{N}{M}∗\frac{L}{K}∗D}{N∗\frac{L}{K}∗D} = \frac{K+\frac{1}{M}∗L}{L} =\frac{K}{L}+\frac{1}{M}
N∗KL∗DN∗D+MN∗KL∗D=LK+M1∗L=LK+M1
可以发现，由于采用了micro-batch的方法，当
L
L
L变大时，流水线并行相比于朴素模型并行，对GPU内存的压力显著减小
如果你使用Pytorch提供的pipeline接口，其中有一个参数叫checkpoint，就是用来做这一项的
最后，再提一点，在micro-batch的划分下，我们在计算Batch Normalization时会有影响。Gpipe的方法是，在训练时计算和运用的是micro-batch里的均值和方差，但同时持续追踪全部mini-batch的移动平均和方差，以便在测试阶段进行使用。Layer Normalization则不受影响。
4 实验效果
回顾第二部分的两个目标，Gpipe真的实现了吗？如果实现不了，又是因为什么原因呢？我们来看下实验效果。
4.1 GPU数量 VS 模型大小
Gpipe分别在AmoebaNet（图像）和Transformer（自然语言）两个大模型上做了实验。
- Naive 表示单卡
- Pipeline-N 表示re-materalization + N卡
- AmeobaNet-D和Trasformer-L 一行表示超参数的量
- # of Model Parameter 表示模型的参数量
- Total Model Parameter Memory 表示模型参数所占内存大小
- Peak Activation Memory 表示峰值时中间结果大小。可以发现，中间结果占据的内存大小是相当可观的。
从实验结果里，我们可以发现：
- 在Transformer上，Gpipe基本实现了模型大小（参数量）和GPU个数之间的线性关系。例如从32卡增到128卡时，模型的大小也从21.08B增加到82.9B，约扩4倍
- 对AmoebaNet而言，却没有完全实现线性增长。例如从4卡到8卡，模型大小从1.05B到1.8B，不满足2倍的关系。本质原因是AmoebaNet模型在切割时，没有办法像Transformer一样切得匀称，保证每一块GPU上的内存使用率是差不多的。因此对于AmoebaNet，当GPU个数上升时，某一块GPU可能成为木桶的短板。
4.2 GPU数量 VS 训练速度
4.2.1 关掉NVlinks
为了验证Gpipe框架带来的收益，实验中关掉了NVlinks（GPU间快速通信的桥梁。实现GPU之间的直接数据通信，而不是通过CPU进行转接）。关掉的意义在于说明，不靠硬件本身的高效通讯带来的收益，Gpipe一样能做的很好。实验效果如下：
M = 32 M=32 M=32表示micro-batch的数量为32， K K K表示GPU数量。从实验结果可知，在关掉NVlinks的情况下，Gpipe一样也能实现随着GPU数量的增加，训练速度也增加的效果。虽然这两者间不是线性的。同样，因为模型切割不均的原因，AmoebaNet的表现不如Transformer。
4.3.2 开启NVlinks，并寻找最佳M
当重新开启NVlinks后，我们来看M的大小（即流水线的核心）对训练速度的影响。
- 当M=1的时候，如前文所说，GPU的空置率太高，因此两个模型都没有实现训练速度和GPU个数间的线性关系
- 当M=4时，表现明显好转。
- 当M=32时，表现最佳，且Transformer基本实现了训练速度和GPU个数的线性关系。
4.3 Gpipe下时间消耗分布
- 对每块GPU来说，约2/3的时间，是真正花在计算上的。
- 其余1/3的时间，大部分花在re-materalization策略下的重计算上。因为采用流水线的方法，bubble的时间也被压缩到很短，可以忽略不计。
参考
- [1811.06965] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
- 对大规模 model training 感兴趣，请问有相关推荐的文章吗？ - 知乎
- 流水线并行技术与飞桨优化实现详解-腾讯云开发者社区-腾讯云
- 图解大模型训练之：流水线并行（Pipeline Parallelism），以Gpipe为例 - 知乎