转载：
今天在arxiv上放出了一篇AI模型自动分布式的文章，欢迎大家参与交流。
文章地址：https://arxiv.org/pdf/2302.08141.pdf
文章向大家介绍了我们的自动分布式系统Rhino。读论文难免会陷入细节中，我在这里为读者列一下重点：
1. Rhino是一套系统，而不仅仅是分布式算法
这个系统是client/server模式。client可以是任何能够生成XLA HLO的前端。server负责进行分布式策略规划，以及自动分布式任务拉起。client与server解耦的动机是期望将来更方便地对接不同的前端框架。
我们还设计了runtime graph，我们有自己的task scheduler。
2. Rhino选择了HLO作为分布式策略规划的输入IR
HLO的粒度刚刚好，目前看到的超大模型基本上有几万条HLO指令。这个规模是我们的算法可以handle的。大家都知道AI模型自动分布式问题是NP难的，但是离开规模谈NP难就是耍流氓。在HLO这个级别，指令之间的连接非常稀疏。纵然有数万条指令，大部分指令仅仅读入1～2条其它指令。完美诠释了弱水三千只取一瓢饮。对于SPMD策略探索来说，分布式通讯的cost来源就是这些指令之间的连接。稀疏的连接是我们敢在HLO上进行策略探索的底气所在。
此外，在HLO这个层级做分布式策略规划，可以避免依赖灵活多变的模型高层次抽象。各种高层次AI模型的抽象概念，比如双向语言模型，自回归语言模型，LSTM, RNN, encoder, decoder, attention，multi head，batch norm, layer norm，token sequence。这些东西Rhino全不懂，Rhino也不care这些抽象概念。这个我后面写文章细说。
3. Rhino的分布式策略探索是全自动的
当然，Rhino同时也允许用户加入annotation部分干预策略探索。
4. Rhino定义了不同的优化级别
目前Rhino有O2, O3两个优化级别，这两个优化级别都采用cost-based的策略探索算法。O2级别而外采用了一些heuristic，超大模型的策略探索耗时更短，但牺牲了策略的质量。O3级别追求更优的策略质量，但策略探索耗时也长一些。我们曾经还有一个O1的优化级别，它采用完全基于规则的策略探索算法。目前已经弃用了。
对千亿模型的话，O3的策略探索耗时长则二十多分钟，短则数秒钟。目前来看无需使用O2。但是随着通用人工智能(比如ChatGPT)的兴起，一个明显的趋势是模型越来越大。必要的编译耗时管理还是很有意义的。我们判断O2优化级别将来会有用武之地。
5. Rhino对策略探索问题进行了合理拆解
我们采用了多种手段，把策略探索问题拆解成优化子问题，运用多种算法，分别对子问题进行求解。有效管理了问题的复杂性。概要说的话，Rhino基于critical nodes（什么是critical nodes？看看论文就知道了）对全图进行了子图划分。子图内部我们进一步划分cone(锥体)。子图之间我们采用动态规划算法，子图内部的cone之间采用ILP算法，而cone内部采用动态规划算法。
6. 我们的流水线stage自动划分比较有特色
划分stage之前，无需把DAG排成拓扑线性序，我们把stage划分建模成ILP问题，使用ILP Solver自动寻找通讯量最小的切分方案
上面没有涉及细节。等我有时间，再写一下Rhino技术的底层逻辑，技术选择背后的原因，我们的一些思考。
内容中包含的图片若涉及版权问题，请及时与我们联系删除
评论
沙发等你来抢