7.3. 流水线并行

在数据并行和模型并行以外，流水线并行是另一种常用的并行加速方法。流水线并行往往被应用在大型模型并行系统中。这种系统通过算子内并行和算子间并行解决单设备内存不足的问题。然而，当这类系统的运行中，计算图中的下游设备需要长期持续处于空闲状态，等待上游设备的计算完成，才可以开始计算，这极大降低了设备的平均使用率。这种现象被称为模型并行空洞（Model Parallelism Bubble）。

图7.3.1 流水线并行系统。\(F_{i,j}\)表示第\(j\)个微批量的第\(i\)个前向stage，\(B_{i,j}\)表示第\(j\)个微批量的第\(i\)个反向stage。

为了减少空洞，提升设备使用率，我们可以在模型并行系统中构建流水线。这种做法的核心想法是将一个数据小批量（Data Mini-batch）划分为多个微批量（Micro-batch）。假设一个数据小批量有\(D\)个训练数据，这个小批量可以被划分为\(M\)个微批量，那么微批量的大小就是\(D/M\)。每个微批量相应进入训练系统，完成前向传播（Forwards propagation）和反向传播（Backwards propagation），计算出梯度。每个微批量对应的梯度将会缓存，等到全部微批量完成，缓存的梯度会被加和，算出平均梯度，更新模型参数。

图7.3.1 进一步给出了一个流水线并行的执行例子。在本例中，模型参数需要切分给4个设备存储。为了充分利用起来这4个设备，我们将小批量切分为2个微批量。当设备1完成第一个微批量的前向传播后（表示为\(F_{0,0}\)）后，他会将中间结果发送给设备2，触发响应的前向传播任务（表示为\(F_{1,0}\)）。与此同时，设备1也可以开始第二个微批量的前向传播任务（表示为\(F_{0,1}\)）。前向传播会在流水线的最后一个设备–设备3–完成。系统于是开始反向传播。设备4开始第1个微批量的反向传播任务（表示为\(B_{3,0}\)）。该任务完成后的中间结果会被发送给设备3，触发响应的反向传播任务（表示为\(B_{2,0}\)）。与此同时，设备4会缓存好对应第1个微批量的梯度，接下来开始第2个微批量计算（表示为\(B_{3,1}\)）。当设备4完成了全部的反向传播计算后，他会将本地缓存的梯度进行相加，并且除以微批量数量，计算出平均梯度，该梯度用于更新模型参数。

流水线并行的关键因素是流水线泡沫（Bubble）。当设备完成前向传播后，必须等到全部反向传播开始，在此期间设备会处于空闲状态。在 图7.3.1 中，我们可以看到设备1在完成2个前向传播任务后，要等很多时间才能开始2个反向传播任务。这其中的等待时间即被称为泡沫。为了减少设备的等待时间，一种常见的做法是尽可能的增加微批量的数量，从而让反向传播尽可能早的开始。然而，使用非常小的微批量大小，可能会造成加速器无法被充分利用。因此最优的微批量大小是多种因素的折中。其中最核心的因素是流水线泡沫的大小和加速器的计算能力。