7月31日，一流科技在创业1300天后，宣布开源自研深度学习框架 OneFlow。与目前业界主流的开源 AI 框架，以及国内已经开源的三款 AI 框架相比，OneFlow 似乎开发的不太早、开源的更晚，它要凭借什么在铁板一块的市场格局情况下杀出重围？此前，InfoQ专访了OneFlow 创建者、一流科技创始人袁进辉，深入解读 OneFlow 背后的故事以及他对于 AI 框架技术、产品和商业化的思考。本文中，一流科技工程师成诚将详细解释OneFlow的设计思路及特色，从技术角度阐释OneFlow在深度学习框架竞争中胜出的可能性。

本文主要内容如下：

OneFlow的设计目标
OneFlow的特色一：Actor机制
OneFlow的特色二：SBP机制
总结

OneFlow的设计目标

OneFlow的设计目标是追求极致的性能，特别是 分布式 多机多卡环境下的横向扩展性，希望能让用户使用多机多卡就像使用单机单卡一样容易，且享受线性加速的运行效率。

为什么OneFlow要聚焦于分布式场景的性能和易用性呢？随着深度学习的发展，模型越来越大，训练深度学习模型所需的算力会越来越高，同时 模型增大的速度要大于GPU单卡显存扩容的速度；训练对算力的增长要求要大于GPU单卡算力增长的速度。 单卡的计算能力和显存远远不能满足深度学习模型训练的需求，必须借助多机多卡并行加速。

若深度学习框架可以让互联的多个GPU协同工作好，实现线性加速比，即使每块GPU性能“一般”，也可满足任意规模的算力需求，这就是所谓的横向扩展性，我们坚信这是算力增长的解决之道。

但是，已有框架都是首先聚焦于单卡的用户体验，仅对适合数据并行的多机多卡场景处理的较好，即把单卡的计算图镜像复制到多机多卡上，各个卡和机器之间辅助于模型同步的模块。

对于BERT/GPT-3等参数量巨大的模型，用户在使用已有深度学习框架时常常会遇到多机多卡不好用、效率低下或无法实现等问题。用户做分布式训练常常需要较高的学习成本，还需要关心多机多卡之间模型的同步问题。业界为解决分布式深度学习的痛点，不仅改进深度学习框架自身，还研发了多种第三方插件，譬如NCCL，Horovod，BytePS，HugeCTR，Mesh-tensorflow，Gpipe等等，但仍不能满足用户极致的性能需求。

OneFlow的核心设计理念是，让多机多卡分布式训练高效地协同运转，同时要让用户在多机多卡的训练体验上就像单卡一样简单容易。下面我们来介绍OneFlow实现此目标最核心的两点想法，以说明OneFlow是如何看待分布式场景下的深度学习训练的。

Actor ：用一套简洁的机制解决几乎所有技术难题

在OneFlow的设计中，分为Compile和Runtime两个时期。Compile时期把用户定义的神经网络、分布式环境信息等编译成一个静态图的执行计划Plan，Plan由执行单元Actor的描述信息组成；Runtime时期，各个机器根据Plan里的Actor描述信息真实地创建属于自己机器的众多Actor实例，然后启动Actor运行系统。

整个深度学习训练期间，OneFlow的执行基本单元就是Actor，Actor对应静态执行图上的节点，Actor之间生产、消费的数据存储在Register中，Actor之间通过消息传递来协作运行。

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

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

每个Actor内部都有一个状态机，Actor收发的消息、执行的情况都会改变自己的状态。需要注意的是，Register是存储块，存放了Actor生产出来的数据，而消息是包含了Register存储块的内存地址的轻量级数据，Actor之间传递的是消息，而不是Register，这样就实现了zero-copy。

当Actor收到了新消息，判断它执行所需要消费的Register已经就绪，且它生产的数据有空闲的Register可以写入时，这个Actor就执行（Act）一次，生产出一个Register。生产完以后，该Actor就向需要消费这个Register的那些消费者Actor们发消息，表示你们可以来读取我生产的数据了；同时该Actor还需要把它消费完的那些Register还给这些Regsiter的生产者Actor们，表示我用完了你们的数据，你可以回收了。Actor内部的状态机如图1所示。

图1 Actor内部状态机

在Actor启动之后，会根据与其他Actor之间收发消息来切换自己的两个状态：等待状态和执行状态。

一个Actor收到的消息一般分为几个类型：a) 上游的生产者Actor发来消息说你可以来读我生产的数据了；b) 下游的消费者Actor发来消息说我用完你生产的数据了。当这个数据被所有消费者都用完以后，就可以回收成为空闲块等待下一次被该Actor重新生产一份新的数据。

一个Actor收到消息以后都会去尝试判断当前是否满足执行条件，执行条件一般有两个：a) 需要读取的数据是否都到齐了；b) 是否有空闲块可以拿来被生产。当满足执行状态以后Actor就开始调用自己内部的Kernel真实的去读写数据。

执行完毕后Actor会向上下游发消息：a) 向下游的消费者Actor发消息说我刚生产了一块数据，你可以来读了；b) 向上游的生产者Actor发消息说我刚用完了你之前发给我的数据了。

Actor只需要关心上下游的消息就能判断自己能不能执行。每个Actor都通过自己内部的状态机和收发消息机制实现了完全去中心化的分布式协同工作。

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。

图 2 Actor生产消费关系和执行时序图

在上面的例子中，到了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机制就实现了流水并行。

在这里我抛出一个更进一步深入的问题：整个数据流的执行像一个网络，数据在网络中的流动就完成了计算，如何避免生产者生产太快，消费者消费不及，以及如何避免生产者生产太慢，消费者感到饥饿的问题，这涉及到对计算、内存、传输带宽的规划，尽可能使系统的瓶颈之处最宽，需要解决流控（flow control）的问题以及资源分配问题（如每个Actor的Register到底分配几个内存块配额），这是非常关键的问题，也是OneFlow系统已解决的。

3 数据搬运是一等公民

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

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

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

图 3 数据是如何从一个设备搬运到另一个设备上的

图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示例图

图4展示了一种可能的Placement例子。用户定义了一个由3个Op组成的网络：Op_0 -> Op_1 -> Op_2。其中Op_0和Op_1的Placement是Device 0，Op_2的Placement是Device 1，这就是一个流水并行的例子，Oneflow会自动在Op_1和Op_2之间插入需要的数据搬运的Copy Op。

2 SBP

SBP是OneFlow独有的概念，他是三个单词的首字母组合：Split、Broadcast、PartiaSum（以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的简单示例。

图 5 几种SbpParallel的简单情形

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：

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

2）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。

图6 MatMul的两种合法SbpSignature

图7 混合并行示例

在图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的代码仓库：
https://github.com/Oneflow-Inc/oneflow

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

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

创作场景

都 2020 年了，开源 AI 框架 OneFlow 还有机会吗？| 技术深度解读