关于深度学习框架的一些自己见解
上次写博客还是10多年前，然后还险些出现了事故，所以就一直没有写了。一晃就10几年过去了，虽然有时候看着一些知乎大咖（前同事等）在消费我们原来做的一些事情，很想回应几句，但是最后也是不了了之。最近团队一位知乎大咖去创业了，所以为了团队的建设，所以决定开始写写自己对于系统领域的一些感悟，也希望通过这个能够让大家更加理解我们团队，以及分享我本人对于系统的一些理解。纯属自己的观点，大家有不同观点，欢迎指正和讨论。
我一直工作在分布式系统的领域，从大数据到AI工程，其实有不少做系统的我们这行的人，很多都有这个路径，所以第一篇就先聊聊自己一个感悟，说说最近做的深度学习框架，算是开一个头。
我觉得做深度学习框架其实有两个派别的人，一派是从分布式系统的人来做的，另外一派是从做算法的人来做的。不同的人的背景不同，所以做这个事情的角度也会不同，从而产生不同门派。tensorflow属于系统派，而pytorch属于算法派。像我们这种做系统，特别是做过超大规模分布式系统的人，往往最担心的就是要对一个已经部署在成千上万台的计算集群上的平台软件需要做重大重构，这个中间困难没有做过这个事情的人可能不会太有体感，这么大一个平台，公司不可能财力让你能够去镜像一个集群去做任务的迁移，并且越大公司的平台上用户数众多，业务都会耦合在一起去完成公司的使命，基本上你不可能有时间点可以让全公司的业务团队放下他们自己手头的优先级来配合你做这种迁移，哪怕你工程能力非常强，这种迁移中间不会出现任何意外。何况很复杂系统要做到这一点基本上是很难的。所以我们做系统的，往往会把系统设计得更加具有可扩展性，从而尽最大可能去避免这种大的重构和推倒重来。当我们在面对需要构建一个深度学习框架的时候，我们第一时间就在设想这个框架需要能够从规模上很好支持分布式，能够很好的扩展到任意大的深度模型的框架，我们希望构建一个系统，能够像人脑一样能够把视觉，语音，语言等多种模型能够一同训练。其实这个就是tensorflow这样系统构造的时候的原始想法，把整个计算构成一个Tensor的Flow图。因为分布式本身就很复杂，需要处理各种节点相互的数据和执行中的各种依赖关系。这些事情由人来写代码，太繁琐且容易出错，所以自然地，我们就会设想由系统来负责这种依赖关系。这也就是为什么我们希望整个分布式执行的计划是一个静态图，然后系统再根据用户指定的或者系统智能的决定的placement进行分图，并在这些分图中添加合适的Send-Recv的Op从而构成一个分布式的执行计划。但是这样的设计理念也会带来一些困恼，我们在模型训练时候有时候有些类似控制图的部分，在这种设计理念下，我们必须要把这些控制流图的代码也op化，然后把这些op也整体串联在Tensor的Flow执行图中，大家有兴趣了解细节的话也可以看看我的老朋友也是前同事Yuan Yu写的文章：Dynamic Control Flow in Large-Scale Machine Learning, Eurosys2018. 但是这种方式会使得一些习惯单机开发的研究人员觉得比较晦涩。同时也是因为分布式的原因，我们做系统的很自然会把模型的开发过程分成构图和执行两个阶段。构图的时候只是生成一个逻辑执行计划，然后通过显式方式的提交（或者execute）过程进行执行。这种方式让研究人员觉得不能一边写代码一边就能够马上看到代码片段的结果，所以这也造成很多人诟病TensorFlow的模式不太容易debug自己的模型程序，其实这也是因为分布式带来负担。但是也是因为TensorFlow是静态图的方式，其可以做到训推一体，在训练出来的模型能够导出模型图，并且在这个图上进行系统化的推理优化从而能够非常方便部署到线上。这个系统性化的方法带来另外一个优势。
框架的另外一派是算法派，特别是感知类模型（图像，语音，语言类）训练，因为这类训练一般都是同步训练，然后“分布式训练”也不像前者那样设想是任意异构的分布式执行图（即每个分布式节点的执行逻辑可以不同），因为是数据并行，这样我们就可以利用MPI的AllReduce的通讯源语来进行梯度的汇集计算。算法同学需要是一种丰富的可扩展的在GPU上能够很好运行的，并且能够很好进行自动梯度的算子库，并且因为面向是数据并行的场景，这样话在神经网络部分其实都是单机程序，从而可以利用任何python的语法糖去构建任何的动态的训练控制逻辑（大家也把这种称作动态图），对于算法研究人员来讲，这种方式写代码比较随性也方便debug，所以在研究界pytorch得到大量的关注和使用。
刚才说过TensorFlow从设计之初就在考虑可以超大的模型分布式训练的场景，但是没有预想到硬件的发展也非常迅速，显存越来越大以及训练技术的发展，还有非常精细化优化显存的工作，比如DeepSpeed等把optimizer所需要的显存sharding化掉，使得除了超大规模稀疏模型训练外，感知类的SOTA模型一直可以利用数据并行的方式来进行训练。从而使得TensorFlow这种设计理念看上去有overdesign的嫌疑。并且就算超大规模稀疏模型训练，因为TensorFlow整体化的设计理念，不把Parameter Server作为游离在Flow图之外，使得他在超大规模场景下的scalability上出现了问题，从而催生一堆自建PS+深度学习框架的（稀疏）模型训练框架。这是另外一个话题，我会在日后写一写在这个领域上我们一些工作。
好在随着transformer的出现，我们终于有方法能够回归到最初那个梦想，使得我们可以把多种数据（视觉的，文字的）合在一起训练多模态的模型，因为问题规模的增大，必然需要更多参数的模型来支持，所以我们迅速将模型大小从几十亿增加到万亿规模，这个时候就必然需要能够支持很好模型并行框架，这也是为什么最近这个领域重新变得火热，比如类似OneFlow，MindSpore，PaddlePaddle，Mesh Tensorflow，GShard，以及我们阿里的Whale框架。
其实从设计理念来看，模型并行正是回归到原来TensorFlow一开始设计时候的设想，只是那个时候因为模型并行的需求不够，没有必要提供比较好高层自动分布式的抽象，写模型的人还是可以自己精细化去构造每个计算节点的子图，整体上TensorFlow的框架只是负责把这些子图能够自动通过Send-Recv进行连接，并且在Runtime能够合法的进行计算。而现在，因为需求增多，算法迭代需求的增多，迫切需要一种高层次的自动分布式框架，从而使得算法同学能够去快速简单构造一个逻辑图的方式去构造自己神经网络，而由系统层来帮助他来进行复杂模型并行的构成。所以其实可以看到TensorFlow的设计理念正好就是为这个考虑的，利用静态图，我们可以逻辑性去描述一个网络训练，然后在执行时候在进行系统化的分图和分布式训练。所以说自动分布式的需求并没有超越原来设计的基本范畴，也是因为这样，我们采取和谷歌Gshard类似技术路线去提供自动分布式的能力。正是站在原有框架基础上去做增量。
不同于GShard更加关注于谷歌TPU集群，我们关注于异构的GPU集群，这里所说异构是因为我们不如谷歌这么有钱，构建非常大的同构化TPU集群，我们集群中有不同年代的GPU和CPU，这些GPU各自算力和显存都大小不一。也正是因为这样，其实给我们系统提出更大挑战，我们在进行自动分布式时候需要在cost model上考虑好这些差异点。这样才能做到比较优化的分布式训练。这也是我们自动分布式框架Whale一种差异性和核心能力之一。
其实系统派的框架和算法派的框架也在进行一定的融合，TensorFlow提出了Eager模式，通过TF.function在eager模式下可能单步执行计算，得到Tensor来提高可调式性；而Pytorch通过Trace或者Parse的方式转化为TorchScript的图描述，从而能够更好支持训练到推理的工程化。但是这种动静结合其实只是在一定层次的，比如如果考虑分布式，Trace的方式去得到TorchScript就不足够。需要进一步去限制构图能够使用的API，这也是像NV的megatron以及微软DeepSpeed在Pytorch上去支持分布式所带来的一些约束，感兴趣的可以读读OneFlow的Blog：《GPT-3难以复现，为什么说PyTorch走上了一条“大弯路”？”》
总结下我的观点：我觉得现在深度学习框架主要流行的两个TensorFlow和Pytorch是有其设计理念的原因的。我们做Whale正是在这种理解的基础上进行路线选择，并且认为应该站在已有的工作基础上去做增量的东西。而不是再去造一个别人做过的轮子。接下来我还陆续展开我们分布式框架Whale，大规模稀疏模型训练工作，编译系统Ansor和DISC，以及我们如何把分布式，编译和调度有机结合方面一系列系统工作，敬请关注。如果大家对于我们PAI团队的工作有兴趣，非常欢迎和我们联系，我的邮箱是weilin.lw@alibaba-inc.com
PS：写Blog写的不多，还需要多训练，感觉有点虎头蛇尾的，呵呵。