大模型训练（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$块GPU，而单块GPU上做一次forward和backward的时间为： $T_{FB}=(T_F+T_B)$
• 图中整体面积为： $K个GPU\ast (K\ast T_{FB})$（宽：$K$，长：$K\ast T_{FB}$）
• 图中实际在做forward和backward的面积为： $K\ast T_{FB}$
• 图中阴影部分的面积为： $K\ast K\ast T_{FB}-K{T}_{FB}$