挑战TensorFlow和PyTorch，“后浪”OneFlow胜算几何？

验“金”室

深度学习算法、算力和数据被看作是当前人工智能浪潮的“三驾马车”。其中深度学习算法通过封装至软件框架的方式供开发者使用，深度学习框架是整个技术体系的核心，实现对人工智能算法的封装、数据处理以及底层硬件的调度使用。正如Windows、Linux在PC时代和iOS、Android在移动互联网时代的角色那样，深度学习框架上承算法应用，下接底层硬件，所以也被称作“人工智能的操作系统”。

目前市面上已有的众多开源框架，用户最多的是PyTorch和TensorFlow，除此之外还有各大公司的自研框架：PaddlePaddle、MXNet、MindSpore、MegEngine等。但是由于业内对于框架的共识越来越广泛，框架的发展越来越相似，导致AI框架领域的竞争日趋激烈：PyTorch和TensorFlow逐渐占据越来越多的用户份额，建立了强大的开源生态；最早推出的Caffe、Theano已经由于技术无法跟上需求的变化，逐渐退出历史舞台；微软的CNTK已经停止了后续维护；MXNet也只在少数领域还有应用……AI框架领域正在从百家争鸣的“春秋时代”走向诸侯割据的“战国时代”。

即便在这么困难的背景下，深度学习框架家族还是迎来了一位备受瞩目的新成员——OneFlow框架。2020年7月31日，在经过1300多个日夜奋斗后，中国一流（OneFlow）科技公司的深度学习框架OneFlow正式在Github上开源，采用Apache 2.0开源协议，这是国内首个由初创公司团队自研并开源的AI框架，立刻登上了深度学习开源社区的话题榜。“后浪”OneFlow会挑战“前浪”TensorFlow和PyTorch吗？胜算几何？下面将从OneFlow的自我定位、技术优势以及未来前景三个方面为大家进行粗浅分析。

一、OneFlow的自我定位

一个框架有没有机会成功，要看它有没有差异化的特点。TensorFlow突出的特点是完备，推理、serving、XLA、tensorboard等工具一应俱全，成为了产业界的首选；PyTorch的特点是易用性好，eager执行、动态图、与python的交互方式等，非常受学术界科研人员喜欢；可以说完备性和易用性这两极已经分别被TensorFlow和PyTorch占据，如果OneFlow单纯地模仿别的框架，与TensorFlow比完备性，与PyTorch比易用性，那么将毫无胜算。所以OneFlow将自己定位为深度学习框架的第三极——性能的领跑者。

深度学习是吞没算力的巨兽，当使用多机多卡集群环境扩展算力时用户会发现当前并没有一种框架完全满足需求，经常会遇到多机多卡加速比不高，得不偿失或者参数量巨大但是现有框架仅支持数据并行导致无法训练的尴尬场景。为解决该问题，业界除了不断改进深度学习框架自身外，还研发了多种第三方插件，例如NCCL、Horovod、BytePS、HugeCTR、Mesh-tensorflow、Gpipe等，但是也只能满足一小部分需求。例如TensorFlow虽然被大众认为在分布式领域是强项，但也仅是在数据并行领域，在小模型的情况下表现得比较好，对于较复杂的大模型来说并不能满足需求；PyTorch易用性好但是为此牺牲了部分效率，在分布式训练任务中表现一般。

OneFlow有其独特的设计理念和技术路线，就是追求极致的性能，而且是分布式多机多卡环境下的横向扩展性。OneFlow从研发之初就主打高效性，致力于解决大规模分布式AI训练问题，其核心设计理念是从分布式的性能角度出发，打造一个多机多卡体验就像使用单卡一样容易而且速度最快的深度学习框架。

为此，OneFlow团队开发了多项诸如编译器自动编排并行模式和流水线、静态调度、去中心化协议、流式执行引擎等技术，并取得了显著效果。有权威第三方机构对OneFlow进行了深度学习框架的基准测试，初步结果如下图显示，对于参数量不大的模型的分布式训练，例如CNN和BERT等，OneFlow与现有框架相比会有10％以上的性能提升；对于超大参数量的大模型，仅支持数据并行的现有框架无法支持，需要引入模型并行和流水并行技术进行深度改造，与深度改造后的现有框架相比，OneFlow在性能上具有数量级的优势。

OneFlow深度学习框架的基准测试结果

二、OneFlow的技术优势

OneFlow与现有开源框架相比最大的优势就是不仅仅支持数据并行还支持模型并行和流水线并行，从而在分布式训练任务中，尤其是超大参数量的大模型场景中保持着绝对领先优势。其核心技术优势得益于两个自研特色机制：Actor机制和SBP机制。

1.Actor机制

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

Actor内部状态机

该机制具有以下特点:

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

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

• 数据搬运与数据计算同等重要。在多机多卡的分布式环境中，各个机器和各个设备之间的数据传输往往是影响系统横向扩展性的最重要因素，如果传输开销可以被计算开销掩盖，那么分布式深度学习训练就可以达到理想的线性加速比。现有的框架在编译静态计算图时主要关注于数据计算而忽视了数据计算和数据搬运的时间窗口重叠安排，OneFlow在计算图中显式表达了数据搬运，而且在静态分析时同等对待数据搬运和数据计算，以最大化重叠搬运和计算时间窗口。

• 尽可能并行。在OneFlow的设计中，所有的出发点都是希望可以尽可能并行，从而达到最优的分布式性能，例如在分布式训练模型梯度同步时，如果显存到内存的传输带宽高于机器之间的网络传输带宽，那么OneFlow将尽量减少网络通信而专注提升单机运算效率，从而升级整体性能。

2.SBP机制

SBP机制是OneFlow独有的概念，即表示一种逻辑上的Tensor与物理上的多个Tensor的映射关系，该机制可以使得用户定义的计算图中的op、Tensor以任意的方式分布在各个机器和各个设备上。例如对于多机多卡环境，OneFlow会把它抽象成一个巨大的设备，即逻辑上的设备，这个设备的资源是实际环境多个物理设备的资源之和，用户只需要在这个逻辑上的超级设备里定义深度学习模型的构建，其余即可交由OneFlow来完成逻辑上的设备到实际物理设备的映射。无论是数据并行、模型并行还是流水并行，OneFlow都可以做自动处理，同时也支持用户自定义配置。

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

三、OneFlow的未来前景

未来随着深度学习的发展，模型会越来越大，训练深度学习模型所需的算力会越来越高，同时模型增大的速度要大于GPU单卡显存扩容的速度，所以分布式深度学习训练会越来越常见，越来越重要，BERT、GPT-3等模型的出现佐证了这一点。但是单个芯片的能力总是受到物理约束，算力增长的解决之道在于横向扩展，而这必须从系统软件的角度去寻找答案，一旦实现就能够把各类芯片连接起来提供源源不断的算力，即使是性能一般的芯片，只要系统软件可以让它们协同工作好，就可以满足任何算力需求,从而真正解决降低分布式编程的门槛和实现逼近线性加速的运行效率两大难题，即通过深度学习框架让用户在多机多卡的训练体验上就像单卡一样简单容易，同时让多机多卡分布式训练实现线性加速比，投入多少倍的计算资源，就可以得到多少倍的加速效果。

OneFlow区别于其他深度学习框架的核心就是对于分布式训练的思考和设计，致力于分布式训练最快且和单卡一样简单易用，真正面向未来的深度学习系统发展趋势，从根源上避免了像TensorFlow这样庞大不灵活、对深度学习系统方向的最新研究成果吸收速度慢、转向调整困难等问题。OneFlow团队也在不断提升框架的完备性和易用性，希望建立完整良好的开源生态。

目前OneFlow框架已经应用于国内头部互联网公司、安防公司、AI芯片公司、云计算公司、政务云以及科研院所等各类机构，涵盖了图像处理、自然语言理解、广告推荐等多种场景。而且OneFlow框架是国内团队开源的深度学习框架，对于未来推进技术国产化具有重要意义。

综上，“后浪”OneFlow挑战“前浪”TensorFlow和PyTorch并非无谋之举，而且相信随着深度学习系统的不断发展，最终会成为深度学习框架家族的重要成员，让我们拭目以待吧!

附录

OneFlow代码仓库：

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

OneFlow官网：

https://www.oneflow.org/

OneFlow文档：

https://docs.oneflow.org/

更多精彩内容

FCC30+

长按左边二维码

关注我们不迷路