挑战 TensorFlow、PyTorch，“后浪”OneFlow 有没有机会？

2020 年 8 月 4 日 CSDN

作者 | 一流科技工程师成诚

编者按：7月31日，一流科技在创业1300天后，他们宣布开源自研的深度学习框架OneFlow，此前，CSDN对CEO袁进辉进行了专访。本文中，一流科技工程师成诚将从OneFlow设计思路、特点、开源趋势等角度出发，详细阐释其在深度学习框架的竞争中胜出的可能性。

主要内容如下：

OneFlow的设计思路，它的独特性
OneFlow的特色一：Actor机制
OneFlow的特色二：SBP机制
对人工智能/深度学习未来的看法
对开源的感想和总结

我是OneFlow的一名工程师，在公司创业之初就加入了。研究生就读于清华软院，读研期间在OneFlow实习了两年，毕业之后正式入职一年。我对OneFlow的整体设计是比较熟悉的，对深度学习框架也有一点小小的见解。

时至今日，OneFlow即将开源，我相信OneFlow的设计和实现可以给大家一种眼前一亮的感觉，也相信在众多深度学习框架的竞争中，OneFlow有机会胜出。

OneFlow是独特的

时至今日，一个框架有没有机会成功，要看它有没有差异化的特点。OneFlow是有其独特的设计理念和技术路线的。目前市面上已有的众多开源框架，用户最多的是PyTorch和TensorFlow，除此之外还有各大公司的自研框架：PaddlePaddle、MXNet、MindSpore、MegEngine等等。其中TensorFlow突出的特点是最完备，推理、serving、XLA、tensorboard等一应俱全，工业部署的主流还是TensorFlow；PyTorch的特点是易用性最好，eager执行、动态图、跟python的交互方式等等，非常受科研人员喜欢。

可以说完备性和易用性这两极分别被TF和PyTorch占据，OneFlow作为一个小团队打造的框架，如果单纯的模仿别的框架，跟TF比完备性，跟PyTorch比易用性，那是绝对没戏的。

但深度学习框架还有第三极：性能。OneFlow的特点就是追求极致的性能，而且是分布式多机多卡环境下的横向扩展性。OneFlow的核心设计理念就是从分布式的性能角度出发，打造一个多机多卡体验就像使用单卡一样容易，而且速度最快的深度学习框架。

深度学习是吞没算力的巨兽。多机多卡的理想很丰满，现实很骨感，用户会遇到多机多卡加速比不高，得不偿失，用户会遇到参数梁巨大时，现有框架不支持模型并行而无法训练。为解决此类问题，业界不仅仅改进深度学习框架自身，还研发了多种第三方插件，譬如NCCL, Horovod, BytePS，HugeCTR，Mesh-tensorflow, Gpipe等等。但是，仍只满足一小部分需求。

下面我就分别介绍两点OneFlow相较于其他框架的独特设计，来体现OneFlow是如何看待分布式场景下的深度学习训练的。

Actor——用一套简洁的机制解决所有令人头秃的技术难题

（关键特性：去中心化调度流水线数据搬运是一等公民传输被计算掩盖控制逻辑被执行逻辑掩盖）

在OneFlow的设计中，分为Compiler和Runtime两个时期，Compiler把用户定义的网络、分布式环境信息等编译成一个静态图的中间表示（称之为Plan）；Runtime时期，各个机器根据Plan里的Actor描述信息真实地创建属于自己机器的众多Actor们，整个深度学习训练期间，OneFlow的执行基本单元就是Actor，Actor之间通过消息机制通信，静态执行图上的节点就是Actor，节点之间的连边是Register，Register存储了Actor之间的生产者消费者信息，以及生产者Actor产生的数据。

1、Actor机制实现去中心化调度

OneFlow的运行时去中心化调度就是用Actor机制实现的。在整个由Actor构成的静态图中，没有一个中心的调度节点，每个Actor都只需要关心自己消费的那些Actor，和消费自己生产出的数据的那些Actor即可。这样在超大规模的分布式训练场景下，完全的去中心化调度可以避免中心调度的单点性能瓶颈问题。

每个Actor内部都有一个状态机，Actor之间的收发消息和Actor自己内部执行都会改变自己的状态。当Actor收到了它执行所需要读取的Register，且它当前有空闲的Register可以生产的时候，这个Actor就执行（Act）一次，生产出一个Register。生产完以后，该Actor就向需要消费这个Register的那些消费者Actor们发消息，表示你们可以来读取我产生的数据了；同时该Actor还需要把它消费完的那些Register还给这些Regsiter的生产者Actor们，表示我用完了你们的数据，你可以回收了。Actor内部的状态机如图4所示。

2、Actor机制实现流水线

上面我们介绍了Actor的内部状态机，Actor之间的消息传递和数据传递是依赖Register实现的。一个Actor是否能执行，只跟两个条件相关：1）自己消费的那些Register是否可读；2）自己生产的那些Register是否有空闲块可写。对于一个Register，如果我们运行时给它分配多个空闲块，那么相邻的两个Actor就可以同时工作，重叠起来，这样就实现了各个Actor之间的流水线。理想状态下整个静态执行图的执行时间就是整个系统中是性能瓶颈的那个Actor运行的总时间，其余Actor的执行时间就被流水线掩盖起来了。

我们举一个例子来解释Actor机制下的流水线是如何运转起来的。图2是一个由3个Actor（a,b,c）组成的静态图的时序图。其中深绿色的Regst方块表示正在被使用的Register块，白色的Regst方块表示同一个Register的备用空闲块。

1）在Time0时刻，Actor a产出了一个Regst_a_0，Actor b 和 Actor c由于没有可读的Register，所以处在等待状态。假设每个Actor的执行时间都是单位时间。

2）到Time1时刻，Actor a给Actor b发消息说你可以来读我产出的Regst_a_0了，Actor b收到了消息，并检查自己生产的Register b是否有空闲Regst块可用，发现有可用的Regst_b_0，于是Time1时刻Actor b执行，读取Regst_a_0，写Regst_b_0；同时Actor a还会去看自己是否有空闲块可写，发现有，Time1时刻Actor a也在执行，写Regst_a_1。（这里需要说明的是，Regst_a_0和Regst_a_1逻辑上是属于同一个Register，只是空间上分成了不同的空闲块备份而已。在深度学习训练任务中，Regst_a_0和Regst_a_1里存放的是同一个operator产出的不同batch的数据。）于是Actor a和Actor b就并行工作起来了。Actor c由于没有数据可读，仍在等待。

3）到Time2时刻，Actor b 生产出了Regst_b_0，于是给下游的消费者Actor c发消息说你可以来读我生产的Regst_b_0，同时给上游的生产者Actor a发消息说我用完了你的Regst_a_0。此时Actor a 已经把刚刚生产的Regst_a_1又发给了Actor b，Actor b检查自己仍有Regst_b_1空闲，于是Actor b开始读Regst_a_1，写Regst_b_1；Actor c收到Regst_b_0，发现自己有Regst_c_0空闲，于是Actor c开始执行，读Regst_b_0, 写Regst_c_0；Actor a收到了Actor b用完还回来的Regst_a_0，检查Regst_a_0所有的消费者都用完了，于是将Regst_a_0回收，标记为空闲块，同时Actor a还可以继续执行，写Regst_a_2。

在上面的例子中，到了Time2时刻，其实Actor a、b、c都在工作，在深度学习训练任务中，Time2时刻Regst_b_0、Regst_c_0存放的是Batch 0的数据，Regst_a_1、Regst_b_1存放的是Batch 1的数据，Regst_a_2存放的是Batch 2的数据。通过一个Register有多个空闲块的设计，Actor机制就实现了流水并行。

3、数据搬运是一等公民

在多机多卡的分布式环境中，各个机器和各个设备之间的数据传输往往是影响系统横向扩展性的最重要因素，如果传输开销可以被计算开销掩盖，那么分布式深度学习训练就可以达到理想的线性加速比。相较于其他的框架，OneFlow把数据搬运视为跟数据计算同等地位的操作，提出“ 数据搬运是一等公民 ”的思想。

已有框架在编译期的关注焦点是数据计算，认为数据搬运是背后隐式发生的，因此在静态分析计算图时略过计算和搬运的重叠编排，OneFlow在计算图中显式表达了数据搬运，而且在静态分析时同等对待数据搬运和数据计算，以最大化重叠搬运和计算。

在OneFlow最终的执行图中，数据搬运操作也是一个个Actor。除了在设备上做数据计算用的Actor以外，还有计算机内存到GPU显存之间的数据拷贝Actor，机器之间做网络通信的网络Actor，负责数据的切分、合并、复制的Actor，负责读取磁盘数据的Actor，负责加载保存模型的Actor等等。很多其他框架都把数据加载、多卡模型梯度的同步、网络、模型加载更新等分别做成一个单独的模块，而OneFlow的设计是所有的功能都在一张由Actor组成的静态执行图里实现了。OneFlow这样的设计不仅简洁、优雅，还非常高效。

图3表示了没有GPU-Direct的情况下，在OneFlow的Runtime阶段，一个设备上的计算节点如果消费了另一个设备的计算节点，数据是如何搬运过去的。

4、尽可能并行

在OneFlow的设计中，所有的出发点都是希望可以尽可能并行，从而达到最优的分布式性能。比如考虑到分布式训练模型梯度同步时，显存到内存的传输带宽高于机器之间的网络传输带宽，OneFlow会做两级的scatter和gather操作（本机的和各个机器之间的），用于增加locality，提高整体性能；又比如在异步启动深度学习训练时，python端用户的控制逻辑跟OneFlow运行时的执行图是并行执行的，同时OneFlow有一套互斥临界区的设计保证执行的高效性和正确性；数据加载部分无论是从磁盘读数据还是从python端喂数据，OneFlow都能保证尽可能并行，使得计算设备不会因为要等数据而导致性能下降。

已有框架要尽可能重叠数据搬运和计算，一般借助多层回调（callback）函数，当嵌套层次过多时，会遇到所谓的callback hell麻烦，正确性和可读性都可能下降。但在OneFlow中，以上的这些并行并发特性，都是在这一套简洁的Actor机制下实现的，解决了令人头秃的callback hell问题。

此外，在多机的网络通信部分，OneFlow底层的网络通信库原生支持RDMA的高性能通信，也有一套基于epoll的高效通信设计。而相比于最流行的Pytorch，多机还需要通过RPC来做数据同步。

Placement+SBP——OneFlow如何做到分布式最易用

OneFlow是目前分布式场景中支持数据并行、模型并行、流水并行等最易用的深度学习框架。用户只需要像单卡一样去搭建网络模型，并告诉OneFlow有哪些机器哪些卡，OneFlow就会用最高效的方式把这些机器和设备使用起来。

这源于OneFlow的一套独特的设计：ConsistentView（一致性视角）。对于多机多卡，OneFlow会把它抽象成一个超级大的设备，我们称之为逻辑上的设备，这个逻辑设备的显存是实际多个物理设备的显存之和，这个逻辑设备的算力也是实际多个物理设备的算力之和。用户只需要定义在这个逻辑上的超级设备里，深度学习模型是如何构建的，其余的便不需要用户来操作，由OneFlow来完成逻辑上的设备到物理上的设备的映射。

这里先明确两个概念：“逻辑上的”和“物理上的”。“逻辑上的”表示OneFlow把分布式集群抽象成一个超级计算机之后的计算和数据，“物理上的”表示那些真实的部署到各个机器和设备上的计算和数据。深度学习网络是由Op构成的计算图，Op之间生产消费Tensor数据。在多机多卡的环境下，一个逻辑上的Op会对应多个真实的物理上的Op，每个物理上的Op实际执行的计算都是这个逻辑Op计算的一部分，一个逻辑上的Tensor也会对应多个物理上的Tensor，每个物理上的Tensor都是逻辑Tensor的一部分。

对于其他的框架定义的分布式训练，每张卡是一个“world”，多卡之间根据暴露出来的接口来同步模型梯度；而对于OneFlow而言，多机多卡也都是一个“world”，我们使用一套Placement+SBP的方式做全局的统筹管理。

1、Placement

在OneFlow的计算图搭建过程中，每个计算Op都有一个属性叫做Placement，表示了该逻辑上的Op，是要部署到哪些机器哪些设备上的。对于常见的数据并行，就是所有的Op都部署到所有的设备上。但OneFlow也支持用户指定Op的Placement，比如当网络过大单卡根本放不下的时候，在OneFlow可以让网络的前一部分在一张卡上，后一部分在另一张卡上，用一种“接力”的方式工作，实现流水并行。

图4展示了一种可能的Placement例子：

2、SBP

SBP是OneFlow独有的概念，他是三个单词的首字母组合：Split、Broadcast、PartialSum（以PartialSum为例，实际上还可以是PartialMin, PartialMax等reduce操作），全称叫SbpParallel，表示一种逻辑上的Tensor跟物理上的多个Tensor的映射关系。

其中Split表示物理上的Tensor是逻辑Tensor按照某一维度切分后得到的， Split有个参数axis，表示切分的维度，如果把多个物理上的Tensor按照Split的维度进行拼接，就能还原出逻辑Tensor；Broadcast表示物理上的Tensor是跟逻辑上的Tensor完全相同的；PartialSum表示物理上的Tensor虽然跟逻辑上的Tensor形状一致，但是物理上的Tensor里的值是逻辑Tensor里对应位置的一部分，如果把物理上的多个Tensor按照对应位置相加，即可还原出逻辑上的Tensor。图5展示了SBP的简单示例。

SbpSignature是一个SbpParallel的集合，在OneFlow的设计里是Op的属性，它描绘了一个逻辑上的Op被映射成各个设备上的多个物理上的Op以后，这些物理上的Op是如何看待他们输入输出Tensor在逻辑上和物理上的映射关系的。一个Op会有多个合法的SbpSignature，一个最简单的合法signature就是输入输出都是Broadcast，这表示了这个Op需要整个逻辑上的Tensor数据。

当用户构建的逻辑上的计算图确定以后，OneFlow在Compiler生成分布式的物理上的执行图时，会考虑每个Op的Placement和该Op允许的合法SbpSignature列表，在其中选择一个传输开销最小的SbpSignature作为本次训练的SbpSignature，用于指导Compiler生成最高效的执行图。

关于Op的合法SbpSignature的列表，我举一个矩阵乘法（matmul）的Op的例子。定义: Y = matmul(A,B) , A, B, Y都是Tensor，表示Y = AB。那么至少存在两种合法的SbpSignature：

Y: Split(0), A:Split(0), B:Broadcast
Y: Split(1), A:Broadcast, B: Split(1)

两种合法的signature在两个设备上的示意图如图6所示。假设逻辑上的MatMul的输入输出Tensor的形状是：

A(64, 10) X B(10, 50) -> Y(64, 50)

且该Op分布在两个设备上。在第一种SbpSignature下，0号设备上的A是逻辑上A的前一半，1号设备上的A是逻辑A的后一半（按照第0维切分），而两个设备上的B跟逻辑上的B完全一致，两个设备输出的Y分别是逻辑上的Y的前一半和后一半。同样可以分析第二种SbpSignature。

值得一提的是，当A是数据，B是模型的时候，第一种SbpSignature就是数据并行，第二种SbpSignature就是模型并行。如果两个相邻的MatMul op，前一个使用第一种SbpSignature，后一个使用第二种SbpSignature，整个网络就实现了混合并行。图7是一个混合并行的示例，定义了Y0 = MatMul_0(A0, B0) , Y1 = MatMul_1(Y0, B1)这样一个由两个op组成的计算图，其中A0, Y0, Y1是数据Tensor，B0, B1是模型Tensor。

在图7中MatMul_0产出的Y0被MatMul_1消费，但是这两个op对同一个Tensor的SBP看待方式是不一样的，MatMul_0认为Y0是Split(axis=0)切分，但是MatMul_1需要一个Broadcast的Y0输入。这时候OneFlow会自动插入一个“万能”的Boxing Op做必要的数据裁剪、拼接、搬运和求和等等操作，使得所有的Op都可以在分布式环境下拿到自己想要的数据。

另外在数据并行的时候，训练的前向模型Tensor的是Broadcast，对应反向传播的梯度就是PartialSum，当Optimizer需要全部的梯度来更新模型时，就会触发OneFlow的Boxing机制进行高效的梯度同步工作。

3、最易用的分布式并行框架

OneFlow的这套Placement + SBP + Boxing的机制，可以使得用户定义的计算图中的Op、Tensor以任意的方式分布在各个机器和各个设备上，无论是数据并行、模型并行还是流水并行，对于OneFlow而言，都只是一个特定Placement下的特定SbpSignature的组合而已，用户可以方便的配置，也可以交给OneFlow来做自动的处理。

另外，早在微软推出ZeRO-2框架之前，OneFlow就已经支持了类似的功能，多机多卡情况下，每个模型Tensor都只保存在其中一个设备上，降低梯度计算中的内存占用。

综上，在编译期，OneFlow通过在一套数学上严谨的形式系统来表示所有合法的并行模式，并支持编译器较方便地自动搜索最优并行方案；在运行期，则通过Actor系统最优地、灵活地支持并行、并发执行。OneFlow的内核具有简洁、高效和高扩展性的优点。

浅谈人工智能和深度学习的未来

OneFlow区别于其他深度学习框架的核心就是对于分布式训练的思考和设计，OneFlow致力于分布式训练最快，且分布式训练跟单卡一样简单易用。为什么OneFlow这么看重分布式的性能和易用性呢？这源于我们团队对人工智能和深度学习的未来的看法。

我们认为随着深度学习的发展，模型会越来越大，训练深度学习模型所需的算力会越来越高，同时模型增大的速度要大于GPU单卡显存扩容的速度；训练对算力的增长要求要大于GPU单卡算力增长的速度。所以分布式深度学习训练会越来越常见、越来越重要。BERT、GPT-3等模型的出现印证了我们的判断。单个芯片的能力总是受到物理约束，我们相信算力增长的解决之道在于横向扩展，而这必须从系统软件的角度去求解，一旦解决，即使是把性能“一般”的芯片互联起来，只要系统软件可以让它们协同工作好，就可以满足任何算力需求。

在分布式深度学习训练中，性能至关重要。我们认为性能和横向扩展性是深度学习框架的核心竞争力。我们在2017年就是这样认为的，也是这样做的，这就是OneFlow这个小团队能存活至今的原因。

私以为，如果深度学习的未来是单机单卡就能搞得定的规模的话，那么深度学习也就只能做做目前已知的一些简单任务，深度学习的浪潮也会褪去。

开源，OneFlow还有很长的路要走

OneFlow团队是一个小团队，OneFlow作为一个开源框架还有很多不足。我们框架的易用性和完善性不如Pytorch和Tensorflow。我们团队还在努力补上OneFlow的短板，同时也非常需要开源社区的支持。如果以爬山来打比方，OneFlow相当于从先难后易的南坡爬珠峰，其它框架相当于从先易后难的北坡爬珠峰。

OneFlow是独特的，有区别于其他框架的技术路线，如果你也认同OneFlow的设计理念的话，欢迎您来帮助我们一起完善OneFlow。

我们相信，从技术上看，OneFlow 是更接近完美的那一个框架。

以上就是我对OneFlow这个深度学习框架的一点粗浅的理解，肯定有很多理解不到位的地方，希望大家批评指正。

OneFlow代码仓库：

https://github.com/Oneflow-Inc/oneflow

OneFlow官网：https://www.oneflow.org/

OneFlow文档：https://docs.oneflow.org/

原文链接：

https://www.zhihu.com/question/409036335/answer/1373468192

    
    
      
     
     
       更多精彩推荐
    
    
      
     
     
       
      
      
        ☞30名工程师，历时1300天打造，又一“国产”AI框架开源了
     
     
       
     
     
       
      
      
        ☞
      
      
        钢铁侠马斯克的野望：实现载人航天，开源特斯拉自动驾驶
     
     
       
     
     
       
      
      
        ☞
      
      
        科技股疯狂造富的背后，“泡沫”离我们到底有多远？
     
     
       
     
     
       
      
      
        ☞
      
      
        AI不止能美颜，美妆迁移这样做 | 赠书
     
     
       
     
     
       
      
      
        ☞
      
      
        程序员必备基础：Git 命令全方位学习
     
     
       
     
     
       
      
      
        ☞
      
      
        公链还能这样玩？二次元、出圈与社区自治
     
     
       
     
     
       
      
      
        
       
       
         
        
        
          
         
         
           
          
          
             
              
            
         
         
           
         
         
           
          
          
            点分享
         
         
           
        
        
          
        
        
          
         
         
           
          
          
             
              
               
              
            
          
          
             
             点点赞 
            
         
         
           
         
         
           
          
          
             
              
               
              
            
          
          
             
             点在看