For: PipeDream: generalized pipeline parallelism for DNN training

简介

当前许多DNN (Deep Neural Network, 深度神经网络)模型被设计的越来越大、越来越复杂，其对应训练所需数据量也越来越多。这些模型的训练需要大规模的并行来缩短训练时间。一些方法将许多个batches同时给多个工作节点进行并行计算，这就导致了这些工作节点间需要进行频繁大量的数据通信。随着工作节点的增多，数据通信会很快成为瓶颈，以致并行效率大大降低。本文提出一种新的方法PipeDream，此方法使用流水线方式​优化单个batch的训练（分割成更小的batch），从而提高数据的吞吐量。这不但降低了通信量，而且可以将通信开销均摊到计算过程中，从而极大的提高了训练效率。

背景

现代DNN模型一般被设计为由许多层组成，在设计并行训练的方法时，这些层常常会被分配给不同的工作节点来进行计算。本节，我们主要讨论两类DNN训练时常用的并行方法：batch间并行 与batch 内并行 。

Batch 间并行

数据并行

数据并行是最常用、最简单的batch间并行方法。输入数据被分割成不同的部分，分别发送给不同的计算节点 (如上图所示)。这些计算节点在本地都保存了模型最新的拷贝。它们使用本地数据进行计算，节点间通过定期同步的方式（比如通过 all reduce操作）来进行模型同步，以防不同节点间模型的差异太大而带来收敛问题。

数据并行方法中，模型同步又分为同步和异步两种。

• 在同步方法中，训练过程被分成许多轮，每一轮开始的时候，所有计算节点需要拉取同样的最新模型到本地，一轮训练在所有计算节点计算完成后，并将结果传送给模型（或梯度）聚合节点进行同步后，该轮才结束。因此同步训练的速度受到速度最慢节点的制约

• 在异步方法中，梯度/模型的聚合不必等待所有节点都计算完成，而是只要有一个节点结果计算出来就发送给聚合节点进行训练参数的聚合，合并后会立即拉取更新后的最新模型进行下一轮训练。在一些方法中甚至不用等待一个节点把当前batch数据的反向传播计算全部完成，就可将先计算得到的梯度先行发送以进一步提升并行效率。这种方法虽然看上去训练数据吞吐量提高了，但由于计算节点间的不同步性，有可能会导致无法收敛的问题，因此在使用时要注意超参设置。

此类方法有一些一些问题：

1. 在模型较大时，通信很有可能成为瓶颈
2. 在多GPU服务器上，GPU间通信（如PCI-E总线）带宽很有可能成为瓶颈
3. 随着工作节点的增加，通信量会快速增长，进一步降低训练速度
4. 如果进行GPU升级（1080 -> V100），通信延时越发明显

在数据并行方面，还有很多降低通信量的方法，比如：量化、梯度过滤、采用更大的batch size等，此处不再赘述。

模型并行

此类方法，将模型分割成不同部分给不同的计算节点，每个计算节点负责计算模型中的一部分。每个计算节点计算完成之后，需要将中间的计算结果发送给负责后续层的节点，模型并行可以让训练非常大的模型成为可能。简单的模型并行一般情况下不会被使用，因为其运行效率极低。下图展示了一个由4个工作节点组成的计算集群的计算情况。其中模型被分成四部分，分别交给四个计算节点，并且反向传播计算的时间是前向计算的两倍（一般反向计算的时间都会被前向计算慢几倍）。可以看出，图中有大量空闲（Idle）块，各个计算节点计算资源的使用效率极低 （图中数字表示不同batch的数据）。

混合并行

在一些论文中，混合使用了模型并行和数据并行的方法，这种混合方法的根据模型的特点对模型进行分割，感兴趣可查看相关论文。

Batch 内并行

Batch 间并行主要关注如何更快的完成多个batch的训练，而batch 内并行关注如何更快的完成一个batch的训练。此处简要介绍一种流水线方法GPipe（后面，我们有时间再专门介绍，本文也可以看作是GPipe的一个轻量级改进）。

GPipe将一个batch的数据再次分割成一些更小的micro-batch，其计算资源利用示意图如下图所示：

图中一个batch被分成了四个micro-batch，用1、2、3、4表示。可以看到，将batch分割后的资源利用率相比我们之前简单的模型并行要提高很多。GPipe中使用的也是同步训练的方法，此方法中需要等所有的micro-batch计算完成之后，由聚合节点将所有计算结果汇总后才可以进行下一轮。可以看出worker 4在反向计算过程结束后轮空了6个时间段，worker 3、2分别轮空了4、2个时间段，因此GPipe在某些情况下还是可以改进的。

PipeDream

PipeDream是一种流水线并行的方法，此法中混合使用了batch内优化和batch间的优化方法。PipeDream将DNN模型分割成不同的阶段(stage)，每个阶段可以包含一个或多个层，不同的计算节点（这里主要指一个GPU）分别计算不同的阶段(stage)（包括前向计算和反向计算）。在简单的模型并行中，单一时刻只有一个GPU在工作。为了可以让流水线高效运行，PipeDream使用了和GPipe一样的策略，将大的batch分为许多小的micro-batch，这些小的micro-batch被一个个的按序发送给计算节点。与GPipe不同的是：PipeDream的计算节点在反向穿过计算过程中得到某一阶段（层）的梯度之后，立即更新该节点所负责阶段（stage）的参数，而不必等待整个反向计算过程完成 （异步性），我们下面会看到这样做的好处。

计算节点完成某个(micro-)batch的计算后，会启用异步进程将结果数据发送给后续的处理节点，在此通信过程中计算节点可以同时进行下一个micro-batch的计算，也就是通信和计算是并行的。

与简单的batch间并行方法相比，PipeDream的优势来自于以下两点：

• PipeDream的流水线的通信量大大减少了。在传统的方法中（如简单数据并行），需要发送整个模型大小的数据。而PipeDream在不同阶段（stages）间，只需要传输一个阶段所输出的中间结果即可；并且，这些结果只需要发送给负责下一个阶段的计算节点，而无需发送给所有其它节点。这带来了通信量上的大幅降低
• PipeDream中计算和通信在时间上是重叠运行的，这极大的提高了效率

PipeDream的基本工作流程如下图所示：

PipeDream工作流可以分为以下几步：

1. 使用DNN模型的结构作为输入，得到模型的基本参数信息(通过一些数据实时跑一下或者计算等方式)，比如每一层的参数量大小，计算时间，激活值大小等
2. 根据第一步计算得到的结果，结合一些额外的限制信息，比如带宽大小、内存限制、硬件图谱等，计算得到一个最优模型分割方法 (模型被分割成多个 阶段（stages）)
3. 使用PipeDream 运行时(runtime)结合输入数据对于分割优化后的DNN模型进行模型的训练

上图的工作流中，还有三个问题待解决。

问题及其解决方案

分割方法

在前面的例子中，为了便于说明，我们假定了每一个阶段的计算时间都是相同的。在实际计算时，这显然是不可能的。实际上，各层的计算时间几乎都是不同的。而根据木桶原理，计算资源的利用效率会受到最慢阶段（stage）的制约。因此，我们的分割方法应该尽可能保证所有阶段计算得都比较快。

文中对于整个计算模型进行了建模。首先，计算资源被进行了层次化建模，每一层的节点的通信带宽应该是相同的。以下图举例来说：

这个计算集群被分成两层，其中绿色的小矩形方块表示GPU计算节点，为第一层。每一个虚线框住的四个GPU表示它们在同一个服务器上（一个服务器4个GPU），为第二层。$B_1, B_2$​表示带宽，显然$B_1 > B_2$​。

我们从顶层来看:

对于下层（第一层，也就是GPU层）来说，我们也可以使用类似的方法得到总计算时间。这样，每一层只需要得到下层的计算结果即可优化当前层。

文中提出使用动态规划的方法来解决模型的划分问题。

在使用算法得到最优的分割前，需要得到一些原始的配置信息，这些配置信息包括：

1. DNN模型第$l$层的计算时间$T_l$ （包含前向计算和反向传播计算）
2. DNN模型第$l$层输出的激活值(activations)和反向传播的梯度的大小 $a_l$
3. DNN模型第$l$层参数(权重，weights)大小$w_l$

假设一个阶段（stage）由$i\to j$​​​层组成，那么该阶段的时间消耗由中间结果的传输时间和计算时间中的最大值决定。最终优化的目标是让所有阶段消耗时间最短。

假设我们的DNN模型共有$N$层，我们的计算集群有$L$层，我们用符号$A^k(i\to j, m_k)$​表示我们对模型进行分割后，最慢的那个阶段所消耗的时间。在动态规划算法中，我们可以将其分为两个子问题:

1. 将DNN模型中的$s+1\to j$层分割出去作为一个新的阶段(stage)，并将其交给$m^\prime$​个计算节点来计算，那么根据前面的配置信息，我们很容易得到该阶段的时间消耗
2. 计算子问题$A^k(i\to s, m_k - m^\prime)$

我们易用动态规划算法来得到一个最优划分。

文中给出了该算法的时间复杂度，为$\sum_{k=1}^L O(N^3 m_k^2)$​，在论文的所有实验中，计算最优分割的时间不超过8秒。

计算调度

神经网络的流水线与传统的流水线不同的是，神经网路训练过程分为两个部分：前向计算与反向传播。由于流水线的存在，节点在计算时可能面临两个选择：

1. 进行前向计算
2. 进行前一个batch的方向传播计算

文中提出了一个非常简单而又实用的调度方案： 1F1B （one forward one backward）。也就是在流水线稳定阶段，每进行一次前向计算，就会进行一次反向传播计算，如此间隔进行。看下图示例：

可以看出，和前面的GPIpe流水线不同的是：

1. 编号较小的反向传播计算被插入到后续micro-batch的前向计算之前，如上图中worker 4 的计算中：batch 1的反向传播计算被插入到batch 2的前向计算之前
2. 一轮(一个batch分割而成的4个micro-batch)计算完成之后，不再等待其它后续micro-batch计算完成后同步所有计算结果，而是反向传播计算出一个梯度之后，直接更新模型（异步性）

把以上两点结合之后，可以明显看出计算资源利用率的提高。在上图中，流水线稳定阶段，理论上计算资源利用率达到了100%。

笔注：在使用流水线时，可以根据之前使用的 batch size 大小来指导前micro-batch(文中也有使用术语minibatch)的大小。实际使用PipeDream时，这个micro-batch的大小可以根据实际实验效果调节。

另外，文中提出的分割方案中，一个阶段中是可以使用多个计算节点的，也就是多个计算节点都来计算同一个阶段（这种典型的应用场景是某个阶段计算特别耗时，然后该阶段内的分层有类似残差网那样的分层结构，无法分割）。

对于一个阶段多个计算节点情况下，文中给出的调度方案是：轮流计算(round-robin)，也就是使用最简单的轮询方案来调度计算。

正确性保证

前面我们提到，PipeDream中使用了异步计算，如果不做任何处理，就会导致一些 batch 在前向传播和反向传播时使用的模型参数不一样，这样的计算显然是错误的，以下图为例：

我们以batch 5为例，batch 5 在worker 1上进行第一阶段的前向计算时，batch 1 的反向传播计算已经完成，也就是batch 1对应的模型更新已经同步到模型中去。而batch 5 在worker 1上进行第一阶段的反向传播计算时，batch 2/3/4的更新都同步到了模型中，此时的模型与batch 5进行前向计算时的不一致，因此直接使用最新模型计算会导致错误。因此，文中提出对于每次更新前都缓存一下之前版本的模型，并且给这些缓存的模型进行版本编号。这样，就可以让每个batch的数据在反向传播计算时，使用与前向传播同样的版本的模型进行计算，文中称此方法为 Weight stashing。

垂直同步

从上图中我们可以看到，上文提到的方式只能保证同一个阶段，前向和反向计算使用同一个模型。而不同阶段使用的模型版本是不一致的。比如上图中 batch 5 进行计算 阶段1 计算时只有batch 1计算结果被更新了，而进行 阶段2 的更新时 batch 2的训练结果也被更新到了模型中。

这个问题显然也很容易通过版本编号来解决。在实际使用时，此法可以不用，文中的所有实验也都没有使用垂直同步方法，得到的结果也很好。

其它细节

文中提到了一些实现细节，其中一个值得注意的是，文中提到NCCL在数据并行方面效率较高，而Gloo在pipeline间小的tensor交换上效率高，因此在本文的实验中，在进行GPU间信息交换时都是用的Gloo，大家在实现时可以参考。

实验

文中进行了大量充分的实验，分别对比了不同的数据集、不同的计算集群、不同的模型结构之间的差异。实验给出的结果当然都很好，这里给出几个主要发现：

1. PipeDream在加速比方面表现非常优异
2. PipeDream方法由于其它现有的batch内优化方法
3. PipeDream大大降低了通信量，并且在使用数据并行时，没有显著增加内存的使用
4. 将PipeDream和模型并行、数据并行同时使用的效果优于单独使用模型并行、数据并行的效果

对于这些点，文中都有实验支撑，对此感兴趣的可以下载论文仔细查看结果。

引用

[1] Narayanan, Deepak, et al. "PipeDream: generalized pipeline parallelism for DNN training." Proceedings of the 27th ACM Symposium on Operating Systems Principles. 2019.