文献阅读-Beyond Data and Model Parallelism for Deep Neural Networks

博主 2889 2021-10-27

文献：BEYOND DATA AND MODEL PARALLELISM FOR DEEP NEURAL NETWORKS
发表于2019 ACM Machine Learning System（SysML）
github开源：https://github.com/flexflow/FlexFlow

此文将关于分布式训练的并行策略问题，变为一个空间搜索问题。作者提出SOAP搜索空间，用以在训练设备和DNN算子上建立可并行维度构成的空间，通过MCMC（蒙特卡诺马尔科夫链）算法在SOAP空间中搜索最优并行策略，提出了FlexFlow框架，用于自动化的在有限时间内寻找最优并行策略并完成DNN训练任务。

背景和工作介绍

DNN模型越来越大，参数越来越多，单机训练已经out，分布式训练才是未来趋势。事实上的DNN网络并行训练策略无非两种，数据并行和模型并行。数据并行是在多个设备上部署相同的模型，把数据拆分到不同设备进行训练。数据并行适用于参数少，计算多（eg：CNN），但不适用于参数多的计算（eg：RNN的embedding）。模型并行是将模型拆分部署到不同的设备，让数据在设备之间移动，这减少了不同设备的参数同步开销（对比数据并行），但需要数据在不同算子（设备）迁移。

但是，针对不同具体NN model部署到device进行训练的过程，也就是如何在已有的设备、DNN结构等限制条件下寻找最优的并行策略，是一个难题（NP-hard）。

在本文中，作者介绍了一个全面的SOAP (Sample-Operator-Attribute-Parameter)并行化策略的搜索空间，它概括并超越了以前的方法。Operator算子维描述了DNN中不同的算子是如何并行化的。对于单个Operator，Sample和Parameter维度表示训练样本和模型参数如何在设备上分布。最后，Attribute维度定义了如何在Sample中划分不同的Attribute(例如，图像的高度和宽度维度)。

作者提出了FlexFlow深度学习引擎，可以在SOAP搜索空间中自动查找任意DNN的快速并行策略，FlexFlow考虑在所有SOAP维度中并行化任何DNN(线性或非线性)，并探索了一个更全面的搜索空间，其中包括作为特殊情况的现有方法。因此，FlexFlow能够找到明显优于现有方法的并行化策略。为了在SOAP空间中搜得更快，FlexFlow使用了两个主要组件:一个快速的增量执行模拟器来评估不同的并行化策略，以及Markov Chain Monte Carlo (MCMC)搜索算法，该算法利用增量模拟器来快速探索大的搜索空间。

FlexFlow execution optimizer执行优化器使用执行模拟器作为oracle，使用MCMC搜索算法来探索SOAP搜索空间，并基于之前候选策略的模拟性能迭代地提出候选策略。执行模拟器还可以与其他搜索策略一起工作，如基于学习的搜索算法。当搜索过程结束时，执行优化器返回它发现的最佳策略。

image.png

FlexFlow 概述图

结构概述

使用Operator Graph描述算子，节点o_i表示算子（卷积、矩阵乘），边（o_i, o_j）表示张量，代表o_i的输出和o_j的输入。用device topology表示所有可用硬件设备和他们的连接方式。每个d_i∈D_N 代表了一个物理设备，比如CPU或者GPU，每条节点之间的边代表他们存在硬件链接，比如NVLink，PCI-e，或者网络链接。

FlexFlow使用上述的算子图和设备拓扑图作为输入，自动的在SOAP搜索空间中寻找并行策略。

SOAP搜索空间

对于任意一个算子，作者通过定义算子的输出张量如何进行分割来对算子的并行化操作进行建模。具体来说，定义算子的可并行维度维P_i，P_i是所有可分割维度的集合。P_i一定包含一个Sampling维度。对于其他可并行维度，如果分割该维度需要对模型参数进行拆分，那么此维度称为Parameter维度，否则称为Attribute维度。最后，还考虑不同算子的算子并行性，称为Operator维度。

对于P_i中的每一个可并行维度，c_i都包含一个正整数，这个正整数表示该维度的并行度。|c_i|是所有可并行维度的并行度c_i的乘积。一个并行化配置c_i将操作符o_i划分为|c_i|个独立的任务，表示为ti:1;...;ti:|c_i|，同时c_i还包括每个任务Ti:k (1≤k≤|ci|)的设备分配。

以下是一个矩阵乘法的例子：
image.png

输出的并行策略记为S，并行策略S描述了应用的一种可行并行方法。S包括对于算子o_i的并行化配置c_i。

执行模拟器simulator

simulator的输入是算子图G，设备拓扑图D，并行策略S，然后在D上以策略S运行G，输出预测的执行时间。

模拟器的执行依赖于一些假设：

  1. task的执行时间是可预测的，并且task多次运行的时间的方差低，且与输入的张量内容无关。
  2. 对于设备d_i和d_j时间的连接，如果带宽是b，那么传输s大小的tensor的时间就是s/b。即带宽被充分利用且只用于传输数据的任务。
  3. 设备处理task使用FIFO调度策略。
  4. 设备的处理时间忽略不计。就是说一旦输入张量到达，之前的task已经处理完毕了，那么就可以立马开始处理输入张量。

任务图task graph

为了模拟执行的过程，对执行过程使用task graph任务图进行建模。任务图包括从算子派生出来的任务以及任务之间的依赖关系。模拟器能够在任务图上运行模拟算法，生成执行时间线。

任务图首先对单独的从算子中继承出来的task进行建模。任务图对硬件设备之间的连接单独建模，称为communication device，这种抽象出来的设备只能进行communication task（例如数据传输）。也就是设备拓扑图中的设备和硬件之间的连接被建模成了不同的设备。

对给定的G，D，S，通过下列步骤生成任务图T = (T_N， T_E)。对于图中每一个node t ∈ T_N，结点代表一个task（普通的task或者communication task），对于每一条node之间的边（t_i, t_j）∈ T_E，边表示了任务t_j的开始依赖于任务t_i的完成（即t_i到t_j是串行的）。值得注意的是，边不代表数据依赖，只指示顺序限制关系，因为数据传输在任务图作为单独的结点t存在。

任务图生成过程：

  1. 对了算子o_i ∈ G，其并行配置是c_i，将任务ti:1,...,ti:c_i添加到T_N中。
  2. 对于G中的张量(o_i，o_j)，通过任务ti:k_i（1≤k_i≤|c_i|）运行输出子tensor的write任务，通过任务tj:k_j（1≤k_j≤|c_j|）运行输入子tensor的read任务。如果ti和tj任务处理的tensor是一样的，且两个任务被分配给了同一个设备，那么给T_E对应的两个结点添加一条边（t_i, t_j）。如果t_i和t_j处理相同tensor，但任务不在同一个设备，那么在t_i和t_j之间添加一个结点t_c，表示t_i和t_j之间需要进行数据传输任务t_c，然后添加两条边t_i->t_c，t_c->t_j。

image.png
图4a为一个标准3层RNN并行策略，代表经典的模型并行策略。图4b为对应的任务图，方形和六边形分别代表normal task和communication task，edge代表任务的依赖。

任务图中每个任务的属性：
image.png

Delta Simulation Algorithm vs Full Simulation Algorithm

全模拟算法和增量模拟算法。
全模拟算法一开始使用上述方法进行任务图生成，然后使用Dijkstra’s变种算法进行task的属性设置。当任务ready后，其被添加到全局优先队列中，队列中的任务以准备时间升序排列。下图展示了经全模拟算法后任务图的生成结果。
image.png

当改变单个算子的并行配置的时候，FlexFlow使用MCMC搜索算法来处理生成新的并行策略。也因此，执行时间线中的大部分位置是不会改变的，作者提出了一种增量模拟算法，算法只重新执行原来改变过的任务图中的任务以及其影响的时间线上的任务。在模拟一个新策略的时候，增量模拟算法首先从存在的任务图中更新任务和依赖，然后将更新的任务加入一个全局优先队列。类似于Bellman-Ford最短路径算法，增量模拟算法迭代的从队列取出更新的任务并向后续的任务传播。

下图展示了将o3算子的c_i设为1之后使用增量模拟算法的结果。
image.png

执行优化机制

执行优化器的输入是算子图G和设备拓扑图D，输入策略S。以模拟器作为准则，FlexFlow把并行策略优化问题转换为损失最小问题，因为其目的是减小预测的执行时间。

FlexFlow使用Metropolis-Hastings algorithm对MCMC采样进行分析，该算法维护当前策略S，并随机生成一个新策略S*，S*以以下式子代表的概率进行接受
image.png
概率受β影响，MCMC类似于一个贪心算法，总是朝损失最小的方向前进。

策略的生成原则：一个当前并行策略中的算子被随机选择，然后用一个随机配置替代该算子的配置，然后使用模拟器的预测执行时间作为MCMC的计算式子中的损失函数。将现存策略（数据并行，专家设计策略）和随机生成策略作为初始候选策略提供给搜索算法，算法迭代的生成候选策略。

评测

使用2个集群和6个DNN进行实验。集群配置：image.png
实验结果：
image.png

并行性能

下图展示了每次迭代中6个benchmark测试的训练性能：
image.png

ResNet-101的策略和数据并行策略相似，结果也相似。在其他DNN上取得了1.3-3.3x speed up。

下图展示了FlexFlow在NMT model上的并行性能
image.png
FlexFlow减少了数据传输的时间，减少了每次迭代的总时间。

Inception-v3上Tensorflow和Flewflow的损失函数随时间变化；
image.png

模拟器性能

模拟器预测时间和实际时间的比较：
image.png
可以看出相差不超过30%，模拟器预测的时间较准。

全模拟算法和增量模拟算法的时间比较：
image.png

案例

image.png
Inception-v3 在4个 P100 GPU上的并行策略。对于每个Operator，竖直和水平方向分别表示样本并行化和参数并行化，图中每个GPU用不同的颜色代表。

image.png
NMT 在4个P100 GPU上的并行策略。对于参数多计算少的层，使用单个GPU计算，避免参数同步。其次，对于参数巨大，计算也密集的layer（softmax layers），FlexFlow使用参数维并行，同时切分参数子集并将计算分配给各个任务。对于多个循环层(如LSTM和attention layer)， FlexFlow利用了不同层之间的并发性以及每个Operator内部的并行性，在平衡负载的同时降低了参数同步成本。

文献