【INFOCOM 2022】支持任意网络拓扑的同步流水线并行训练算法，有效减少大规模神经网络的训练时间

今天要介绍的文章是来自阿里云机器学习平台 PAI 与香港大学吴川教授团队等人在 INFOCOM 2022 发表的论文，论文名是"Efficient Pipeline Planning for Expedited Distributed DNN Training"。论文提出了一个支持任意网络拓扑的同步流水线并行训练算法，有效减少大规模神经网络的训练时间。

原有技术问题

• 同步流水线并行（pipeline parallelism）在训练轮间插入同步墙，用于累计梯度和更新模型保证模型收敛性。但同步墙会阻塞流水线，导致同步流水线并行的 GPU 利用率低，训练时间长。

• 异步流水线并行去掉同步墙，允许提前往流水线中加入未来训练轮的 microbatch 来提高 GPU 利用率，导致 microbatch 在过时的模型参数上训练，会影响模型的收敛速度甚至导致模型无法收敛。

新技术创新点

本文提出了一个支持任意网络拓扑的同步流水线并行训练算法，保证模型的收敛性能不改变，同时相比之前的同步流水线并行策略，有效减少大规模神经网络的训练时间。论文主要的贡献包括：

1. 一个基于递归最小割的 GPU 排序算法，通过分析 GPU 间网络拓扑确定 GPU 的模型部署顺序，保证最大化利用 GPU 间带宽；

2. 一个基于动态规划的模型切分部署算法，高效率找到最优的模型分割与部署方案，平衡模型在每个 GPU 上的运算时间与模型切片间的通信时间；

3. 一个近似最优的列表排序算法，决策每个 microbatch 在各个 GPU 上的执行顺序，最小化模型的训练时间；

新技术主要框架以及关键技术点

作为分布式机器学习的一种主流训练方式，流水线并行通过同时进行神经网络计算与中间数据通信，减少训练时间。我们考虑混合的同步流水线并行，通过结合数据并行（data parallelism）进一步减少训练时间。一个同步流水线并行方案包含模型切分设备部署与 microbatch 执行调度两个部分。上图给出了一个 6 层神经网络模型在 4 块 GPU 上进行同步流水线并行训练的示例。由图表（a）所示，模型被切分成三个 stage，其中第二个 stage 由于其计算量较大，被复制到两个 GPU 上通过数据并行（data parallelism）的方式训练。图表（b）表示模型的三个 microbatch 的具体训练过程，包括每一个 microbatch 在每一个 stage 上的前向和后向计算。其中，由于第二个片段以数据并行方法在 GPU2 和 GPU3 上训练，在全部 microbatch 训练完成后通过 AllReduce 算子同步模型片段参数。

microbatch 执行调度

先考虑在给定了模型切分设备部署方案下的执行调度问题，这个问题属于著名的 NP-hard 的 job shop 问题的一个特殊情况。针对此，我们设计了一个基于列表排序的近似最优的调度算法，包括 1）一个执行顺序排序算法，决策每一个 stage 上每一个 microbatch 的前向和后向计算的执行顺序；和 2）一个基于上述排序的调度算法。

我们定义每个 microbatch 在一个 stage 上的一个前向或者后向计算为一个 computation block，定义每个 microbatch 在两个连续的 stage 间的前向通信（传递 activations）或者后向通信（传递 gradients）为一个 communication block。由于一个 microbatch 在最后一个 stage 上的后向计算可以紧接在其前向计算之后，我们将这两个计算合并成一个 computation block。假设模型被切分成了 $S$个 stage，每个 microbatch 在每轮训练中包括$2S-1$个 computation block 和$2S-2$个 communication block。我们将这$4S-3$个 block 按照执行顺序依次排列，组织成一个有序表$J=\left\{1，2，...，4S-3\right\}$。对每个 block$j\in J$，我们维护一个队列$Q_j$，用来存储所有的其前序 block$J-1$已经被处理了而当前 block$J$尚未被处理的 microbatch index。换言之，$Q_j$存储全部当前准备执行 block$s$，我们维护一个执行顺序队列$U_s$，用来表示每个 microbatch 在当前 stage 的前向和后向计算的顺序，每个元素用 (microbatch index, block number)表示。

我们的排序算法具体如下。从 1 到$4S-3$依次遍历$J$中，我们从每一个非空的$Q_j$，$j\in J$队列中取出一个 microbatch m，将其加入到$Q_j$中。如果 j 是一个 computation block，我们就将$\left(m，j\right)$加入到$U_s$中。在遍历 J 一轮的过程中，对于每个 block，我们最多将一个 microbatch 加入到该 block 对应的$U_s$中。我们重复上述遍历过程，直到全部的$Q_j$变为空，表示我们已经将全部的 microbatch 的 computation block 都加入到了执行顺序队列$U_s$中。当排序算法完成后，我们严格按照每个 stage 上的执行顺序队列$U_s$指定的顺序对每个 microbatch 的 computation block（即该 microbatch 的前向和后向计算）调度。对于存在数据并行的 stage（即该 stage 被复制到了多个 GPU 上），我们在该 stage 全部 microbatch 计算结束完成后加入 AllReduce 算子同步模式参数。

通过理论分析，证明了该调度算法下的单轮训练时间不会超过最优单轮训练时间的$\left(2+\left(4S-4\right)/M\right)W$倍，其中$M$为一轮训练中的 microbatch 数量，而$W$为模型在单一 stage 上的包括 AllReduce 的全部计算时间，和两个连续 stage 间的全部通讯时间的最大值。因此，在模型切分与部署阶段，我们需要最小化$W$来取得最优的流水线训练效率。

模型切分与部署

递归 GPU 排序

实际系统里 GPU 间的网络拓扑复杂，本文设计了一个递归 GPU 排序算法（recursive device ordering，RDO）确定 GPU 的模型部署顺序，保证最大化利用 GPU 间带宽。GPU 间的网络拓扑可以建模为一个图$G=\left(V，E\right)$，其中图节点$V$表示每个 GPU ，边$E$的权重表示链接两个 GPU 间最小带宽大小。排序算法 RDO 将拓扑$G$作为输入，利用最小割算法 [1] 找到图 G 上权重和最小的割，并基于此图将一分为二为$G_1$和$G_2$，继续作为输入调用 RDO 算法递归切割，直到我们可以对全部的 GPU 排序，记为$\left(v_1，v_2，...，v_{|V|}\right)$。

基于动态规划的模型切分与部署

基于$\left(v_1，v_2，...，v_{|V|}\right)$，本文采用动态规划对模型进行切分和部署。用$W\left(l，\xi ，r，i\right)$表示当我们将模型的前$l$层切分成$\xi$个 stage 部署在 GPU$v_1$到$v_i$上，且最后一个 stage 被复制到$r$个 GPU 上时，最优的（即最小的）单个 stage 全部计算时间和任意连续 stage 间全部通信时间的最大值。我们可以按如下公式递归计算出$W\left(l，\xi ，r，i\right)$：

第一个部分表示将前$l^{\prime }$层切分成$\xi -1$个 stage 部署在 GPU$v_1$到$v_{i-r}$上，且最后一个 stage 被复制到$r^{\prime }$个 GPU 上时的最优单一 stage 计算或者任意连续 stage 间全部通信时间的最大值。第二个部分表示前$l’$层和后$l-l^{\prime }$层间通信的时间，而第三个部分表示$l-l^{\prime }$层的包括 AllReduce 的全部计算时间。

对于一个 L 层的深度学习模型，我们有$W=mi{n}_{1\le r\le V}W\left(L,S,r,V\right)$。因此，给定模型切分的 stage 数量$S$，我们可以基于动态规划算法得到最优的$W$以及对应的切分方案。遍历全部可能的模型切分 stage 数量$S\left(1\le S\le V\right)$，我们可以找到最优的模型切分和部署方案。

实验

我们基于 Tensorflow 1.14.1 实现了我们的同步流水线规划算法（synchronous pipeline planning，SPP），并在两个测试平台上验证了我们的算法。一个测试平台由四个 GPU 服务器构成，每个服务器陪伴有两块 1080Ti GPU，通过一个 50Gbps 交换机相连。第二个测试平台为一个单机四卡服务器，配备有四块 Tesla V100 GPU，通过 NVLink 相连。我们选取了七个具有代表性的深度学习模型用于测试，模型具体信息如下表：

本文将 SPP 与四个具有代表性的分布式机器学习训练方式进行比较验证，包括：1）数据并行；2）GPipe[2]，一个同步流水线并行框架；3）PipeDream[3]，一个异步流水线并行框架；和 4）HetPipe[4]，一个针对异构集群的异步流水线并行框架。针对异步流水线算法 PipeDream 和 HetPipe，作者通过插入同步墙的方式将他们改为同步流水线并行算法用于比较。具体结果如下，其中 speed-up 的计算方式为$\frac{baselinetime-SPPtime}{SPPtime}$。

可以看到 SPP 在每个模型上都取得了超越基线算法的表现。进一步，作者将 SPP 与原版的 PipeDream 进行比较，在 V100 GPU server 上训练 VGG19 模型，采用 ImageNet 训练集，将 90% 验证集准确率作为最终的训练目标。实验结果如下：

可以看到，尽管 SPP 与使用我们实现的同步流水线模式 PipeDream 相比，每个训练轮训练时间只有很小的加速比，但在端到端的训练过程中，SPP 比 PipeDream（采用原始的异步流水线设计）训练速度高出了 9.05%。这是因为 PipeDream 的异步流水线训练减缓了模型收敛的进度，进一步延长了模型训练时间。

结论

本文提出了一个支持任意网络拓扑的同步流水线并行训练算法，有效减少大规模神经网络的训练时间。算法对一个给定的深度学习模型进行模型切分与部署，充分利用 GPU 间通信带宽，并设计一个高效的调度器，在不同的 GPU 上的 stage 进行 microbatch 执行调度，最大限度地减少训练时间。我们在两个 GPU 测试平台上进行的比较实验证明，我们的设计优于最先进的方法。

参考文献：

[1] M. Stoer and F. Wagner, “A Simple Min-Cut Algorithm,” Journal of the ACM, vol. 44, no. 4, pp. 585–591, 1997.

[2] Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu et al., “GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism,” in Proc. of NeurIPS, 2019.

[3] D. Narayanan, A. Harlap, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, P. B. Gibbons, and M. Zaharia, “PipeDream: Generalized Pipeline Parallelism for DNN Training,” in Proc. of ACM SOSP, 2019.

[4] J. H. Park, G. Yun, C. M. Yi, N. T. Nguyen, S. Lee, J. Choi, S. H. Noh, and Y. ri Choi, “HetPipe: Enabling Large DNN Training on (Whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism,” in Proc. of USENIX ATC, 2020.

论文 PDF 链接：https://arxiv.org/pdf/2204.10562.pdf