仅此一文让您掌握OneFlow框架的系统设计(中篇)

本文是OneFlow系统设计分享系列文章的中篇，主要介绍OneFlow的编译期Compiler如何将Job编译为Plan的。其中最精华的部分是OneFlow的Boxing模块，负责构建两个逻辑上的Op对应的两组物理上的Op在任意情形下的物理子图，完成了分布式训练中各个机器各个设备之间的数据拷贝、切分、传输、通信的子图搭建。值得一提的是，Boxing模块的代码实现是非常直观且易扩展的，使用了设计模式中的责任链模式（Chain of Responsibility），未来我们会结合OneFlow的代码实现分享一些C++编程技巧的文章，以及为什么OneFlow要使用这些编程技巧，解决了哪些问题，敬请期待~

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

如果您对OneFlow这套致简致快的框架设计感兴趣，或者对深度学习框架、分布式系统感兴趣的话，本文就会让您完全掌握OneFlow的系统设计。相信读完这篇文章，您就会理解我们是如何看待分布式深度学习训练的，我们为什么要这样设计，这样设计的好处是什么，以及我们为什么相信OneFlow这套设计是分布式深度学习训练框架的最优设计。全文的主要内容如下：

• 深度学习框架原理
• OneFlow系统架构设计（简略版）
• OneFlow完整运行流程 与 各模块的交互方式
  • \1. 分布式集群环境初始化
  • \2. Python端搭建计算图
  • \3. 编译期： OneFlow(JobSet) -> MergedPlan
  • \4. 编译期： Compiler(Job)->Plan
  • \5. 运行时： Runtime(Plan)

全文分上、中、下三篇。本文是中篇。

这三篇系列文章会偏重于工程实现细节，所以会引用大量的OneFlow源码，适合对OneFlow底层代码实现感兴趣的读者阅读，尤其是在想要读一个框架的代码但是却不知道该按照什么顺序阅读时，本文可以给您一些参考。在这三篇着重于代码解读的文章之后，我会再写一篇梳理OneFlow核心概念的文章，作为精简版，可以快速掌握OneFlow的系统设计。

上篇请见：

https://blog.csdn.net/OneFlow_Official/article/details/111868930

OneFlow完整运行流程 & 各个模块之间交互方式

我们通过介绍一次OneFlow完整运行的流程来了解系统中的各个主要模块是如何协同工作的。

4. 编译期： Compiler(Job)->Plan

为了方便理解，我们再简要描述一些重要的概念和抽象：

• Job：用户定义的逻辑上的计算图，由逻辑上的Op组成。
• Plan：编译生成的物理上的计算图，由物理上的Task组成。
• Task：运行时Actor的配置描述，一个Actor与一个Task一一对应，Task内部有Op的运行时描述Kernel的配置。Task并不一定关联某个用户计算图Job中的逻辑上的Op，因为并编译期会增加很多物理上的Op用于数据搬运、网络传输、切分/拼接等操作。Task中标记了自己是在哪台机器哪个设备上，并使用哪个线程工作等。Task还需要指定自己产出Regst的regst num、内存类型、所属内存块的偏移量等等信息。
• Regst：运行时的数据存储、Actor之间通信的基本单元。存储某个具体的Tensor。

编译期（Compiler）的设计体现了OneFlow作为一个分布式深度学习训练框架的很多重要的设计原则：

1）一致性视角（Consistent View）

OneFlow把整个分布式集群抽象成为一个超级设备，用户使用OneFlow做分布式训练跟做单机单卡的训练没有任何区别。体现在：

• 编译期Compiler仅在Master机器上编译整个Plan，其他Worker机器等待获取Plan启动运行时即可。Master上编译的Plan就包含了所有机器所有设备上的基础单元——Task（Actor）。
• 所有的分布式训练过程中各个机器各个设备之间的数据通信、同步操作均被Compiler自动生成，无需用户关心和编写分布式训练中的数据同步。

2）数据搬运是一等公民

OneFlow将所有的数据加载、预处理、拷贝、网络传输、Tensor的自动切分/拼接/广播/求和操作都抽象成了跟计算Op一样的运行时执行体——Actor，即在分布式的物理计算图上显式表示了数据搬运的操作。这样做的好处是OneFlow可以感知到所有的数据搬运、同步操作，因此编译期Compiler可以更好的在整个物理计算图上做全局调度，使得这些数据搬运操作尽可能被计算操作所掩盖，对数据搬运操作的性能优化转而变成了图分析与图优化。

3）编译期全局调度，运行时去中心化调度

OneFlow的运行时是一个极其简单的抽象——Actor，每个Actor仅需要关心和自己相关的上下游Actor的消息就可以知道自己能否工作，这样做的好处是运行时系统不会因为有中心调度节点导致性能瓶颈（在计算图非常大的情况下）。为了做到这一点，OneFlow的调度工作大多都是在编译期完成的，Compiler会做好全局的内存调度、Op执行调度、通信调度等工作，使得运行时的调度开销尽可能的低，从而达到更快的训练速度。

4）天然支持流水线，解决流控问题

编译期Compiler通过推导和设置Task产出Regst的regst num，可以使得运行时相邻Actor之间可以流水并行起来。同时还可以通过背压机制（Back Pressure）解决流控问题（Control Flow）。具体的Actor机制如何解决流水线和流控问题的讨论我放在下篇中介绍。

在原生的OneFlow设计中，Compiler输入是一个Job（用户定义的op list），经过编译生成OneFlow的中间表示IR（Intermediate Representation）——Plan，Plan是一个被Runtime直接读取就能生成运行时执行图的描述。而上面介绍的OneFlow(JobSet)->MergedPlan是为了支持Python前端交互 + 多Job（Train/Eval同时做）而后设计出来的。我们下面介绍OneFlow的Compiler做了哪些事。

4.1 JobCompleter().Complete(job)

第一步，经过JobCompleter将Job不断重写。经过多个Pass以生成最终的Job。中间借助OpGraph抽象不断推导新的Job对应的逻辑图。这些Pass包括一些优化如插入KeepHeaderOnly节点；增加Source/Sink的Tick节点使得图成为一个单源节点和单汇节点；增加控制边；计算临界区；以及使用XRT框架重新构建Job。

XRT框架会将Job中的OpGraph进行有选择的合并，并选取使用XLA或者TensorRT来进行编译生成优化后的Kernel。对于OneFlow而言，这些都是