Beyond Data and Model Parallelism for Deep Neural Networks

对训练深度神经网络的计算需求很高，分布式训练已经是常规操作。目前的深度学习系统通常用数据并行或模型并行。但是这些策略在并行程度上通常无法达到最优。这篇文章提出了一个复杂的深度神经网络的并行策略：SOAP。Sample，Operation，Attribute，Parameter。提出了FlexFlow：一个在SOAP维度为特定的并行机器随机搜索快速并行策略的深度学习框架。为了加速这个搜索的过程，FlexFlow用了一个创新性的、可以准确预测一个并行策略的表现、比原有方法更快速的执行模拟器。实验结果表明FlexFlow可以很大程度上增大训练的吞吐率。本文由斯坦福大学团队发表于MLSys 2019。

背景

模型越来越复杂，训练集越来越大，训练模型的计算需求也越来越大。因此分布式训练成为了常规操作。深度学习框架中对并行的应用很简单，最常用的是数据并行。数据并行对计算密集、参数较少的DNN很有效。另外一种常见的是模型并行。已有工作对池化卷积层用数据并行，对全联接层用模型并行。（利用experts’ knowledge）。但是仍然不是最优的。有工作做出了自动化并行策略。已有的自动化的框架只探索不同op之间的并行或者单个op的并行。其实两者结合可以得到一个更快的策略。

这篇文章提出FlexFlow，一个可以自动在更大的范围内找出快速并行策略的深度学习框架。为了形式化这一问题，我们首先定义了SOAP。Operation维度描述了一个DNN中不同的operation是如何并行的。另外，对于一个单独的DNN operation来说，sample和parameter维度描述训练样例和模型参数如何在不同设备之间分布。最终，attribute维度定义一个sample中不同的attribute是如何划分的。已有的系统都是在SOAP的子集中划分的。
在SOAP这个更大的范围内搜索的一个主要的挑战是快速评估候选的并行方案已找到一个高效的方案。已有的工作依赖于在硬件上执行一轮训练来评估不同方案的执行时间。在SOAP的范围内，这样的方法代价太高。

为了解决这样的问题，FlexFlow提出了一个创新性的执行模拟器，可以准确预测并行策略的表现，比profile真实的运行快了三个数量级。设计模拟器的挑战在于如何准确估计不同DNN op的执行时间（非线性，取决于硬件）。模拟器依赖于两个事实：（1）很多DNN模型只用少数几个不同的op（2）op的执行时间通常差异不大，很大程度上取决于输入数据。FlexFlow的模拟器对于每种数据大小，用一个op的计算时间来衡量同种类op的计算时间。然后，这些估算被用于预测各种各样的并行策略。另外，模拟器使用了一种delta simulate算法，这种算法基于对之前的模拟的更新对新的策略作出模拟。何以有的方法比，这样的方法有两个优势：更快、所需资源更少。
模拟器的预测准确率很高。FlexFlow的execution optimizer使用一种马尔可夫蒙特卡洛搜索算法探索SOAP的搜索空间，并给予对之前的候选策略的模拟表现选出候选策略。搜索过程结束后，optimizer返回最佳的策略。

概览

程序接口

与多数深度学习框架不同，FlexFlow用设备拓扑结构描述所有可用的应尽设备和他们之间的关联。拓扑结构中的“边”有带宽和延迟的标签。
FlexFlow可以自动为一个计算图和一个设备拓扑结构找到合适的并行策略。主要有两大优势：提供了易于编程的接口；可移植性（为不同的硬件自动选择搞笑的策略）。

FlexFlow架构

Execution optimizer为计算图和设备拓扑图选择高效的并行策略。

图1：FlexFlow Overview

限制

执行模拟器假设每一个op的执行时间是可预测的且与input tensor的内容无关。所以该方法不适合执行时间data dependent的模型。

SOAP搜索空间

图2：1D卷积的并行策略样例

对于一个op o_i，用Pi表示它可以被分割的维度。Pi总会包括sample维度。如果一次分割需要把模型中的不同参数分割开，则就被称为parameter维度；否则，被称为attribute维度（例如在卷积中，把二维图片按照长宽分开）。o_i分割的config c_i定义了这个o_i在不同的设备之间是如何并行的。对于P_i中每一个可并行的维度，c_i包括一个正整数表示那个维度并行的“程度”。|c_i|是ci_中所有可并行维度的degree的乘积。我们对每一个维度采用同样大小的分割来保证分布式负载的良好的平衡。ci把oi分成|ci|个独立的任务，分别用ti:k表示。

图3：模拟并行策略的样例

策略S描述了一个application的可能的并行方案。S包括对于每一个oi的ci。

执行模拟器 Execution Simulator

输入：算子计算图G，设备拓扑结构D，并行策略S
输出：执行时间
simulator的重要假设：（1）每个task的执行时间都是可预测的，波动小，与input tensor的内容无关。（2）不同设备之间传输数据的时间为数据大小/带宽。（3）每个设备按照FIFO的顺序执行任务（GPU就是这样的）。（4）每个设备在完成一个任务后，只要下一个任务的数据准备就绪就立刻开始执行下一个任务，overhead可忽略不计。
为了模拟一次执行，模拟器首先建立一个task graph，然后运行模拟算法。

Task Graph

把设备之间的物理连接视为communication device，每一次数据传输被视为communication task。T = (T_N , T_E )。

全模拟算法

是后续delta模拟算法的baseline。
首先懂Dijkstra算法遍历。所有任务都被放到一个队列里，出队列的顺序是按照ready time的增序。该算法最终返回所有任务重最慢的一个执行完所需时间。

delta模拟算法

使用一种MCMC搜索算法，每次只改变一个op的划分方式。这种情况下，前后两个策略的时间通常没有改变。delta simulation算法只重新模拟改变最终结果了的。对于同样的task graph，full和delta的模拟算法会给出同样的结果。

执行优化器 Execution Optimizer

输入：算子计算图G，设备拓扑结构D
输出：最有效的并行策略
问题抽象为最小化总执行时间。这个方法避免了平衡执行时间和通信时间二者的问题。
这是一个NP难的问题，但有方法可以简化。可能的策略数量是op数量的指数，所以不可能穷尽整个搜索空间。为了找到一个开销低的策略，采用开销最小化搜索。

MCMC采样

MCMC是从概率分布中获取样本的一个采样方法。公式1：把cost function转化为概率分布的常见方法。
MCMC开始与整个搜索空间的任意一点。然后产生一系列的点，这一系列的点达到 p(·)给出的分布。
使用某种算法生成马尔可夫链，维护一个当前的策略S和一个修改后的策略S。若S被接受，则S替换S。这一过程无限重复，直到某一规定的时间被耗尽。公式2：S被接受的标准（可能性？）。
若S比S的cost更高，则S也有可能被接受（？）。MCMC趋向于贪心搜索，更倾向于选择cost小的策略。

搜索算法

在当前的策略中随机选一个op o_i，把ci替换成一个随机的config。
用已有策略/随机生成的策略作为candidate。对于每一个初始的策略，如果满足下面两个条件之一，搜索算法就会提出新的candidate：（1）当前处世策略的搜索时间budget被耗尽；（2）一般的搜索时间里都不能找出更好的策略。