现有的深度学习框架通常使用模型并行或数据并行来解决深度神经网络（DNN）的并行计算问题，但这两种策略往往会导致模型无法达到最优结果。今年ACM的SysML大会上，斯坦福大学Matei团队在论文《Beyond Data and Model Parallelism for Deep Neural Networks》中提出了一种新的DNN并行化策略搜索空间——SOAP。SOAP从样本（Sample）、运算符（Operator）、属性（Attribute）以及参数（Parameter）四个维度度寻找更好的DNN并行化策略。为了加速这个搜索过程，该论文同时提出了一个新的深度学习引擎——FlexFlow。**与目前最先进的方法相比，FlexFlow能够将训练吞吐量提高3.3倍并且可以达到更高的精度。**作者还比较了FlexFlow和TensorFlow在Inception-V3上进行端到端训练时的性能，FlexFlow比TensorFlow减少了38%的训练时间。我们对该方法做了梳理与总结，本文是AI前线第82篇论文导读。

1 背景介绍

随着深度学方法的发展，DNN模型变得越来越大，这导致了DNN的计算成本也变得越来越高。因此，对于DNN的实际应用，分布式异构DNN训练集群成为了标准化的做法。尽管深度学习模型以及硬件集群变得越来越复杂，但是用于并行化训练的框架（例如，TensorFlow、Caffe2以及MXNet）却仅用了一些简单的并行化方法，并且其优化结果往往并不是最理想的。

最常用的神经网络并行化可分为数据并行和模型并行。数据并行的基本思想是使用多个设备上的模型副本同时对数据子集进行训练，并在迭代结束时同步跨副本的模型参数。数据并行对于计算密集型且参数量少的运算具有高效性，例如卷积操作。但对于参数量大的运算操作则不能得到最优的结果，例如嵌入操作。而模型并行则是让分布式系统种不同的设备负责网络模型的不同部分，例如将卷积操作分给一部分GPU而嵌入操作分给另一部分GPU或CPU设备。

图1：左半部分为模型并行，不同的机器负责模型中的不同运算模块；右图为数据并行，在不同的设备上创建了完整的模型副本对数据子集进行处理。

去年的PMLR上，这篇论文的作者Zhihao Jia就提出了一种基于线性计算图的CNN逐层并行化方法——OptCNN（相关文献：(https://arxiv.org/pdf/1802.04924.pdf)。OptCNN使用动态编程的方法联合优化每种运算，但是没有考虑不同运算之间的并行性。此外，OptCNN不适用于大部分的语言模型，包括针对机器翻译、推荐系统等趋向于RNN或其他非线性的深度神经网络。

针对上述问题，SOAP从扩展搜索空间和模拟运行两方面出发，与现有的Metropolitan-Hastings 马尔可夫蒙特卡洛方法（MCMC）相结合寻找模型在并行系统下的最优解。其中，作者通过扩展可行解集合的方法将模型的最优化搜索空间扩展为SOAP搜索空间，以期在更多可能的潜在并行策略中寻找最优的一种情况。此外，使用FlexFlow构建的执行模拟器（execution simulator）可以在SOAP搜索空间中更精确的评估并行化策略的性能。

2 方法介绍

2.1 方法总览

FlexFlow是一种使用运算符图G来描述DNN中所有运算以及状态的深度学习框架。图G中的每个节点o_i表示一个运算，例如矩阵乘法、卷积等；而图G中的每个边(o_i,o_j )则表示一个张量，这个张量是o_i的输出，同时也是o_j的输入。此外，FlexFlow还使用了设备拓扑图D=(D_N,D_E)来描述所有的可用硬件设备以及他们之间的连接关系。

图2：FlexFlow总览图,左上角为运算符图，右上角为设备拓扑图，下半部分为执行模拟器与搜索算法迭代求解过程的示意图。

如上图所示，设备拓扑图中D_E中的一个节点d_i表示一个设备，可能是一个GPU或一个CPU；图D_E中的每个边(d_i,d_j)表示两个设备之间的一个硬件连接，例如NVLink或PCI-e。图中的每个边会使用连接的带宽和延迟对它们进行标记。FlexFlow会将运算符图和设备拓扑图作为输入，在SOAP搜索空间中自动搜索有效的策略。搜索空间中的所有策略都会使用DNN中定义好的计算方式进行实现，因此可以保持与模型设计时相同的准确率。在图2中，优化器使用MCMC搜索算法在SOAP搜索空间中进行搜索，模拟器则用增量模拟算法（delta simulation algorithm）来产生候选的子策略，这两个步骤迭代进行从而找到最优解。

2.2 SOAP搜索空间

对于一个运算符o_i，作者通过研究其输出的分割情况来对它的并行化操作进行建模。作者将o_i的可并行化维度（即哪些方面可以进行并行化）定义为P_i，P_i是o_i的输出的所有可分维度的集合，因此o_i一定包含一个样本（Sample）维度。表1中列出了一些运算符的可并行化维度。

表1：几种不同运算符的并行化维度

对于P_i中的所有其他维度，如果在该维度上进行分割操作需要拆分模型的参数则将其称为参数维度（Parameters），其他的维度则称为属性维度（Attribute）。此外，作者还考虑了运算符维度（Operator）中不同运算符的并行性。图3展示了一些1D卷积运算符在单维度以及混合维度上的并行化配置的例子。

图3：1D卷积的并行化配置示例。虚线表示对张量的分割。

对于一个运算符o_i，它的一个并行化配置（parallelization configuration）c_i表示在跨设备时该运算符会被如何并行化处理。对于每个P_i，c_i包含一个正整数，这个正整数表示在该维度中并行化的等级。|c_i|是c_i的所有并行化维度的并行化等级的乘积。作者在每个维度中使用相同大小的分区来保证均衡的工作负载。并行化配置c_i将运算符o_i分为|c_i|个独立的任务t_{i:1},…,t_{i:|c_i|}，同时c_i也包含了分配给每个任务t_{i:k}的设备。经过上述建模，给定一个任务的输出张量以及它的运算符类型，我们可以推理出执行每个任务所必须的输入张量。图4是对于一个矩阵乘法运算符（U=VW）的并行化示例。该运算被划分为四个独立的任务并被分配到不同的GPU设备上。该运算的输入输出如图中所示。

图4：矩阵乘法并行化示例。

最终，作者使用一个并行化策略S表示应用中的一个可能的并行化结果。S包含了每个运算o_i的所有并行化配置c_i，并且o_i的配置可以从所有可能的配置中被独立地选择。

2.3执行模拟器

执行模拟器将运算符图G、设备拓扑图D以及一个并行化策略S作为输入，通过在设备拓扑图D上使用策略S运行运算符图G来预测执行时间。执行模拟器的设计依赖于以下这些假设：

A1：每个任务的执行时间是可预测的，具有低方差的性质，并且与输入张量的内容无关。
A2：对于d_i和d_j间的具有带宽b的每个连接，将一个大小为s的张量从d_i传递到d_j的时间是s/b，即通信带宽可以被完全利用。
A3：每个设备使用FIFO（先进先出）调度策略处理分配的任务。这是GPU等现代设备使用的策略。
A4：运行时的等待时间可以忽略不计。一旦设备的输入张量可用且设备已完成先前的任务，设备就开始进行处理任务。

为了模拟执行过程，作者借用了OptCNN中的想法，为每个配置测量每个不同运算符的执行时间，并将这些测量值包含在任务图（task graph）中，任务图包含了所有从运算符中派生出的任务以及任务之间的依赖关系。通过在任务图上运行模拟算法，模拟器可以生成一个执行时间线。

任务图模拟了从运算符派生的各个任务之间的依赖关系。为了方便抽象化，作者将设备之间的硬件连接模式转化为只能执行通信任务的通信设备。注意，设备和硬件连接被分别建模为单独的设备，如果可能的话，计算（即，常规任务）和通信（即，通信任务）是可以重叠的。对于一个操作符图G、一个设备拓扑图D以及一个并行化策略S，作者对其构建一个任务图T=(T_N,T_E)。其中任务图中的每个节点表示一个任务（即常规任务或通信任务），任务图中的每个边(t_i,t_j)表示一种任务依赖关系，即直到任务t_i结束之前，任务t_j不能开始。

图5：一个并行化策略示例

图5展示了一个标准3层RNN网络的一个并行化策略示例，该网络包含了一个嵌入层、一个循环层和一个线性层。它表示了一种常用的模型并行化操作，即将每层中的不同操作符分配给一个专用的GPU。图6则是相应的任务图，每个正方形和六边形分别表示常规任务和通信任务，并且每个有向边表示任务之间的依赖关系。

图6：图5中并行化策略相应的任务图

对于任务图中每个任务的属性，作者给出了下表作为参考：

表2：任务图中每个任务的相关属性

作者实现了一个全模拟算法作为基线，从而与他们提出的增量模拟算法进行比较。作者使用了Dijkstra’s 最短距离算法的变体对任务图中的每个任务的属性进行配置。对于Flex Flow，作者使用了MCMC搜索算法来产生新的并行策略。在通常情况下，从一个策略转变为另一个策略时，大部分任务的执行时间并不会产生变化。基于这个观察，作者提出了一种增量模拟算法，直接从前一个任务图推断当前更新的任务图的时间线。增量模拟算法极大的提高了执行模拟器的运行速度，尤其是对于大的分布式机器。图7是全模拟算法与增量模拟算法的任务图比较。

图7：左半部分为全模拟算法的任务图，右半部分为增量模拟算法的任务图。

2.4 优化器

为了降低搜索策略地成本，FlexFlow使用成本最小化搜索来启发式地探索搜索空间并返回最优策略：根据给定的初始策略S，通过随机选择一个运算符并用随机选择的新运算符替换其并行化配置来创建新的候选策略S^*，然后使用模拟器估计新候选策略的成本。新的候选策略被接受并变为新的当前策略地概率服从以下分布（Metropolis-Hastings MCMC）：

换句话说，只要新策略具有相同或更低的成本，就可以保证被接受。但同时系统也可以接受具有较高成本的策略，其可能性由参数β控制。

3 实验

作者使用FlexFlow在两个不同的GPU集群上对六个真实DNN基准进行评估。两个集群结构如下图所示：

图8：NVIDIA Tesla P100集群和K80集群。

图9展示了所有6个DNN基准测试地每次迭代训练效果，对于ResNet-101，FlexFlow搜索到了类似数据并行地策略，除了在最后一个全连接层的单个节点上使用了模型并行外。相比于其他的DNN基准测试，FlexFlow比基线算法找到了更多有效的策略，并且速度提高了1.3-3.3倍。

图9：六个不同的DNN模型训练效果。括号中的数字是实验中使用的计算节点的数量，虚线表示理想的训练过程吞吐量。

与其他并行化策略相比，FlexFlow发现的策略将每次迭代过程中的数据传输减少了2-5.5倍，并且还减少了任务的计算时间。此外，对于一个DNN模型，FlexFlow与其他深度学习系统中的计算方法相同，因此可以得到相同的准确率。表3验证了FlexFlow在实验中可以使DNN基准模型达到了最好的精度。

表4：实验中DNN模型的实现细节、所用数据库以及准确率。

这部分实验中，作者还比较了FlexFlow和TensorFlow在Inception-V3上进行端到端的训练时的性能。作者在ImageNet数据库上使用两种不同的框架训练Inception-V3模型，当模型在验证集上single-crop的条件下top-1精确率达到72%时停止训练。对于FlexFlow和TensorFlow这两种框架，作者均使用了随机梯度下降（SGD）进行优化，其中学习率为0.045，权重衰减为0.0001。图10展示了两种系统的训练曲线，可以看出FlexFlow比TensorFlow少了38%的训练时间。

图10：不同框架下Inception-V3的训练曲线，硬件系统为16个P100 GPU（4节点）。

同时，作者也对模拟器的性能进行了评估。在图11中，作者比较了不同模拟算法进行策略搜索性能。全模拟算法和增量模拟算法分别需要16分钟和6分钟以完成搜索。如果将搜索时间限制在8分钟内，全模拟算法只能得到很差的搜索结果。

图11：全模拟算法和增量模拟算法在16个P100GPU（4节点）对NMT模型进行搜索的性能比较。

此外，作者对本文使用的搜索算法与针对小规模计算的全局最优化搜索方法的性能进行了比较。为了获得合理的搜索空间大小，作者将设备数限制为4，并使用了两种DNN模型。其中一个时6层的卷积神经网络LeNet，另一个是RNNLM的变体。经过实验，作者发现对于全局最优化搜索方法，LeNet和RNNLM分别需要经过0.8和18个小时的搜索才能得到最优策略。而FlexFlow在1秒内就能找到全局最优策略。

4 实例探究

图12给出了在四个P100GPU上Inception-v3的最优并行化策略，该策略在关键路径上使用运算符内并行化，在不同的分支上组合使用运算符内的运算符间并行化。这种并行化策略可以实现工作负载地均衡化，同时可以减少参数同步时的数据传递。与数据并行方法相比，该策略将参数同步时的成本降低了75%，并将每次迭代的执行时间减少了12%。

图12：在四个P100GPU上Inception-v3的最优并行化策略。对于每个运算符，竖直和水平维度分别表示样本并行化和参数并行化。每个GPU由不同的颜色表示。

图13是NMT模型在四个P100GPU上搜索到的最优化并行策略。首先，对于网络中参数量较大的且运算较少的层（例如embedding层），该策略使用单GPU进行运算以减少参数同步的成本。其次，对于参数量大且运算量大的层（例如softmax层），FlexFlow使用参数维度的并行化策略，并为每个任务的参数子集分配计算资源。这一方式在保持工作负载平衡的同时，减少了参数同步的成本。最终，对于多层循环层（例如LSTM和Attention层），FlexFlow使用不同层之间的并发以及每个运算符内的并行性来降低参数同步成本，同时平衡负载。

图13：在4个P100 GPU上并行化NMT的最佳策略。对于每个运算符，竖直和水平维度分别表示样本并行化和参数并行化。每个灰色框表示一个层，其运算符共享相同的网络参数。每个GPU由不同的颜色表示。

论文原文：https://www.sysml.cc/doc/2019/16.pdf

Matei团队其他研究成果

Spark团队开源新作：全流程机器学习平台MLflow

性能翻倍！斯坦福Matei团队推出机器学习模型优化新方法

创作场景

训练速度提升超 3 倍！斯坦福推出全新深度学习并行计算框架 FlexFlow