(OSDI 2022) Unity来了！结合Parallelism和Algebraic ...
大家好，好久不见。很久没有扫论文了，今天较为快速地分享一下OSDI 2022的一篇关于自动分布式引擎的文章——Unity。OSDI 2022还有另一篇自动分布式文章Alpa，想查看这篇解析笔记的同学可以点入这个链接：https://zhuanlan.zhihu.com/p/487588274
一些背景
文章的作者是Zhihao Jia，代表作有大家熟悉的FlexFlow，TASO等。本文提出的Unity是在TASO的工作基础上又往前延续了一步，将自动分布式并行化也融入到TASO的方法论中，最终实现了Algebraic transformation和Parallelism的联合优化。
因为Unity文章中介绍的方法大量引用TASO，所以建议提前了解下TASO的idea和具体做法，避免理解上过度跳跃。
这回我们先看Propsal和Motivation，再看结果，最后翻过来深究Unity Approach。
总览
文章开头就提出，Unity是首个将algebraic transformation和parallelization做联合优化的分布式DNN训练系统。
Algebraic transformation一般是针对单机优化而设计，而分布式Parallelization优化也有概率和单机优化相互冲突。所以将这两个搜索空间糅合，一定能够产生更好的hybrid方案，但搜索空间必然爆炸，这是非常直观的，也是论文核心要解决的问题。
组合空间爆炸怎么办？
要么剪枝，要么改良搜索算法，要么both。在搜索算法改良方面，Unity使用Hierarchical search algorithm做联合优化方案的搜索。好巧不巧的是，在另一篇文章Alpa中，面对inter-op和intra-op组合的巨大搜索空间，也是使用的Hierachical search。分层解决问题，让人爱不释手。
Motivation
思考一个问题，为什么一定要将两种优化同时考虑，即做Joint optimization？如果按顺序一个一个应用行不行？
显然，作者一定会比较Sequential optimization和joint optimization的优劣的。如果按顺序将优化逐个应用到模型上，那么无论是先Algebraic transformation，还是先Parallelization都有各自的问题。
[*]先Parallelization：会有实现上的工程问题。因为后发生的Algebraric transformation优化，会引入或替换产生新的operator，导致这些新的operator没有被parallelization分析过。
[*]先Algebraic transformation：不会出现上面的问题，但这会有概率产生bad case。毕竟二者各做各的，利益很可能有冲突。
所以，Joint optimization是必然选择，搜索空间爆炸是必然面临的问题。
下图对比了Sequential optimization和Joint optimization的图变换效果。Sequential optimization先应用fusion将Matmul和ReLU合并成一个op，然后再执行并行化（data parallel），而Joint optimization能够产生某些Op被fuse，另一些op不被fuse，并且同时拥有Parallelization能力的策略。这说明，Joint optimization确实会产生Hybrid 策略。
性能优化效果速览
在32 nodes共192 GPUs的集群上（注意单机6卡，网络通信也是100GB的IB），对7个real-world DNNs做了测评。实验表明Unity可在20min内搜索出优于当前framework 3.6x的优化方案
Benchmark models
具体是哪些模型呢？参看下面的表。
乍一看ResNeXt-50，Inception-v3，BERT-Large和MLP都不是很大的Dense模型，至少和流行打榜的GPT-3，M6等在规模上差了很多，但喜人的是增加了DLRM和XDL这两个推荐模型的benchmark，至少可以说明Unity提出的方案并没有“过拟合“到仅仅是Transformer-base结构，还是挺好的。
E2E Performance
性能上，虽然Unity不一定outperforms 100%的case，但它在全自动化的基本前提之上，做到下面的加速效果，至少在产品上已经是比较理想的状态了。
Joint optimization效果
联合优化的威力如何？实验表明，叠加的并行方案越多，效果越好。同时将Algebraic transformation和parallelization联合考虑，效果越好。
策略搜索时间
我标记了下面一段话。因为搜索时间和operator数量规模有关，所以十分关注这个数字。在323个operators的情况下，搜索时间大概是20min以内。但如果是比GPT-3更大的模型，不知道时间会不会超过1小时。众所周知，如果在tf上实现一个DNN，其operator的数量可能动不动就上千个。
Unity方案
Unity是在TASO之上的工作，所以必须在新的feature——并行化方面，做一些改造和适配。Unity方法如下：
[*]提出Unified graph表示——Parallel Computation Graph （PCG）。它将DAG、parallelization和数据搬运都考虑在内
[*]可以对PCG方便的生成混合策略，不必在为每种优化单独设计策略生成的工程方法
[*]Joint optimization：使用Hierarchical search方法，在PCG替换和device placement的空间中做探索。在这里，cost model考虑了计算、通信甚至是网络拓扑和设备的异构性。
此图是Unity的workflow，关注点在：
[*]输入是PCG
[*]Joint optimization而不是sequential optimization
[*]TASO中的子图替换和生成模块：Substitution generator
Parallel Computation Graph (PCG)
PCG能够同时表达三种东西：computation，communication和parallelism。为此，需要设定三种描述，以增强表达能力。
Tensor Representation
将Tensor建模成a set of dimensions，每个维度都有两个信息：size和degree。size好理解，degree就是该维度被切分的份数。每个tensor也都有一个replica dimension，表达该tensor被复制的份数。
Machine Mappings
每个Tensor被partition成不同的分片，会有自己的id。machine mappings维护了每个task id映射到的device。除次之外，machine mapping还具有一些扩展能力，如面对计算能力不对等的节点时，可以通过添加一个维度说明该信息来指导search algorithm。
这个图展示了Machine mapping的过程。我们设想，某个n维Tensor在每个维度上分别被切分，如标记degree为(d_1,d_2,...,d_n)，那么一共就有d_1*d_2*...*d_n个sub Tensor。每个Sub Tensor都有自己的index，比如记录为(i_1,i_2,..,i_n)，machine mapping就是要解决这个index为(i_1,i_2,..,i_n)映射到哪个计算资源的问题。
Parallelization Operators
回想下TASO做法，先对原Computation graph生成多个substitutions graph，也就是变换过的新子图，这些新子图在文献中被称为candidate substitutions，然后再将他们应用到原Graph上做部分替换，然后评估cost，不断决策。
那么如果想让变换的图具备分布式的能力，就得引入新的专门表达分布式模式的算子，于是就有了这一小节的工作。抽象了共六种，共三对的operator。之所以是看成pair的形式是因为某个operator一定是另一个operator的inverse版本。
[*]Partition and combine
[*]Replicate and Reduce
[*]Pipeline and Batch
下面是Paralellization Operator在PCG上的表示示例。黄色的节点就是引入的Parallelization Operator。有了这些Operatior，就可以让PCG具备并行化的表示能力了。
Graph Substitutions
这部分基本可以沿用TASO的思路了。在定义了上面PCG后，我们有了Parallelism Operator，因此可以借用TASO的方法做Graph Substitutions。回顾一下TASO的Graph Substitutions方法，包括两个步骤：
[*]Substitution generation：这里使用的启发式generation方法和TASO几乎相同，只不过这次是在PCG上，并且要考虑新引入Parallelism operator算子。同时，在TASO中使用的fingerprint也需要扩展（fingerprint是substituion graph的hash key，相同的key可以粗略认为是等价可替换的关系）
[*]Substitution verification：使用Z3 prover，和TASO相同。
因此，在Graph Substitutions阶段，已经能够将algebraric transformations和parallelization合并考虑产生新的hybrid变种。
在使用图替换时，我们可以借助joint optimization发现一些新的优化可能性，如下面的(b)。如Add->Concat+Reduce，进一步可以优化为AllReduce。
策略搜索
搜索的目的是决定：
[*]选取用哪些candidate substitutions
[*]对其做何种machine mapping方案，以达到优化e2e training time的目的。
Unity的Hierarchical search分成三个层级，每个层级解决的问题并不一样。
[*]candidate子图替换——Middel level
[*]为candidate子图做machine mapping映射——Lowest level
[*]大图分割——Highest level
整体的层次如下所示。
Substitution Selection——Middel level
这里的搜索方法依然使用了TASO中cost-based回溯算法。具体可以参考TASO论文，这里只简单描述下在Unity上的应用：首先，维护一个装有candidate PCG的有序queue（需要按照execution time排序），算法不断迭代取出best candidate。然后，用它不断通过为每个location应用每个available substitution，来生成新的candidates。
此算法有个必须的依赖项，就是cost model。而cost的预估又离不开machine mapping，因此引入lowest level的部分。
Finding Optimized Machine Mappings——Lowest level
这一层次解决的就是PCG candidate的machine mapping问题。Unity认为一个PCG candidate依然是一个比较复杂的子图，所以需要将结构按属性做分图后才可做machine mapping。这里提出了两种结构：
[*]Sequential graph
[*]Parallel graph
图作为一种具有递归性质的结构，可以借用dynamic programming方法将图进行分离肢解。然后，就可以将计算资源分配给每个部分了。细节可能得看code才能知道。
Scaling to Large Graphs
这是一个解决规模爆炸的方法。可以将大的图分离成若干子图后进行Graph substitution和machine mapping。但要考虑的是各个sub
graph之间的串联问题。串联时要尽可能保证切分兼容，如果万不得已，需要插入类似reshard的parallelism operator才可实现等价变换。
这里的每个子图的优化是不是要考虑全局性，论文中的描述和讨论并不是很多，可能还得跟进code才能看出。Paper中提到的唯一一个指导跨sub graphs的优化条件是：尽可能保证上下游的串联兼容性。
页:
[1]