OneFlow源码解析：静态图与运行时

作者｜郑建华
更新｜许啸宇、张文骁、成诚

OneFlow静态图的训练效率远高于动态图（eager模式）。本文试图通过一个简单例子，结合v0.8.0版本的代码，解读一下静态图和运行时的实现机制。

在开始之前，建议先读一下参考资料中《OneFlow框架的系统设计（https://zhuanlan.zhihu.com/p/337851255）》等系列文章。对静态图、运行时的基本概念和设计理念有基本的了解，会更容易理解代码。

1 

代码示例

下面的示例代码来自官方文档（https://docs.oneflow.org/master/basics/08_nn_graph.html），是一个线性模型的前向计算。后续主要基于这段代码进行分析。

import oneflow as flow
import oneflow.nn as nn
 
 
class ModuleMyLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(flow.randn(in_features, out_features))
        self.bias = nn.Parameter(flow.randn(out_features))
 
 
    def forward(self, input):
        return flow.matmul(input, self.weight) + self.bias
 
 
linear_model = ModuleMyLinear(4, 3)
 
 
class GraphMyLinear(nn.Graph):
    def __init__(self):
        super().__init__()
        # ModuleBlock
        self.model = linear_model
 
 
    def build(self, input):
        # ModuleBlock.__call__
        return self.model(input)
 
 
graph_mylinear = GraphMyLinear()
input = flow.randn(1, 4)
out = graph_mylinear(input)
print(out)

2 

oneflow包的初始化

import oneflow在初始化包（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/__init__.py）时，与静态图相关的主要操作如下：

• GetEnv（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/__init__.py#L228）

  • EnvGlobalObjectsScope::Init（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L126）

    • 启动各个节点的控制面（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L160-L162）网络连接

    • 初始化VM（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L180）

    • 启动各个节点的数据面网络连接（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L184-L188）

    • 初始化KernelObserver（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L192-L203）

• NewDefaultSession（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/__init__.py#L229）

  • RegsiterSession（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/multi_client_session.py#L39） 创建 Session，并注册为 default session（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/session_util.cpp#L89）

  • 创建 Python MultiClientSession 并保存到dict（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/session_context.py#L40），但并不 TryInit

    • 创建 C++ MultiClientSessionContext（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/multi_client_session.py#L41） 但并不 TryInit

EnvGlobalObjectsScope::Init中先创建一个全局的ProcessCtx（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L132）对象。然后根据环境变量等配置，在各个进程间创建gRPC和CommNet的连接，分别负责控制面和数据面的数据传输。其中在Bootstrap过程中会初始化全局的ProcessCtx（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/rpc/lib/grpc.cpp#L42），给每个进程分配一个全局唯一的rank编号（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/rpc/lib/global_process_ctx.cpp#L28）（machine_id（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/rpc/lib/global_process_ctx.cpp#L24））。

本文不涉及网络层面的操作，只讨论同一进程内各线程间的交互。

3 

Module类

虽然可以直接用op和tensor构造模型，但是op的粒度太细了，直接用op构造模型会比较繁琐。

Module（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/module.py#L54）是由op和tensor构成的、可复用的子模块。利用Module可以更高效、更快捷的构建复杂模型。oneflow.nn（https://github.com/Oneflow-Inc/oneflow/blob/d825243aa7aff5cba8bd3a901b4cc56c2b1a36af/python/oneflow/nn/__init__.py）模块导出了很多预定义的Module。

Module定义了自己的属性设置逻辑（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/module.py#L262），核心逻辑是

• 如果value是Parameter类型，就保存到Module._parameters中

• 如果value是Module类型，就保存到Module._modules中

• 如果value是Tensor类型，就保存到Module._buffers中

• 否则按常规属性处理

Module可以包含子Module，形成树结构。因为Module通过setattr将子Module和Parameter都保存到字典结构中，可以方便的遍历所有Module及其参数tensor。

4 

Graph类

4.1 构造函数

Graph的构造函数中GetDefaultSession（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L145）得到的session，就是导入oneflow包时NewDefaultSession（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/__init__.py#L229）构建的session。当时没有初始化，而是在Graph构造时进行初始化（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L147）。对应的C++函数是MultiClientSessionContext::TryInit（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/multi_client_session_context.cpp#L67），执行时会创建各种全局的资源管理器，比如：
 

• LazyJobBuildAndInferCtxMgr

• BufferMgr

• RegstMgr

• ActorMsgBus

• ThreadMgr

4.2 __setattr__: 将Module和Tensor封装为Block

Graph.__setattr__ 支持通过设置属性的方式把一个 Module 添加到 Graph 中，之后改 Module 就可以被 Graph 调用了。添加到 Graph 中的 Module，会被包装到 Block 里面，Block 起到了代理执行的作用，它会给原 Eager 下的 Module 扩展出静态执行需要的一些特殊功能。

添加到 Graph 中的 Module 和原 Module 共享了状态（Parameter、Buffer）和 forward 执行逻辑。共享 forward 执行逻辑使得静态和动态执行计算逻辑相同。共享状态则可以使动态图下的模型状态被静态图复用。基于此，两个 Graph，一个用于训练，一个用于预测，他们都复用统一模型 Module，这样训练和预测 Graph 也就实现了模型共享。

setattr最重要的动作就是对_add_block的调用（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L1332），_add_block中主要是调用get_block_cls并保存结果（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L1326）。get_block_cls（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L39）的作用是将Module及其所有Tensor属性都转为对应的Block对象。为什么要做这个动作呢？主要是静态图编译需要借助Block类型来实现代理执行的功能，这些功能不适合直接写到 eager 下的 Module 和 Tensor 上。

这个转换是在ModuleBlock构造时调用set_origin（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L131）完成的。对于子Module，会递归调用get_block_cls函数（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L145），这样所有子Module及其Tensor属性都会被转换为对应的Block对象。

所以，上述示例代码中，GraphMyLinear实际存储的是ModuleBlock，Graph.build执行时获取的model属性也是ModuleBlock对象，ModuleBlock.origin才是ModuleMyLinear。

Graph.__setattr__不允许将Tensor对象设置为属性（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L1340）。Tensor只能存到Module中，因为 Module 是做状态共享的基本单位，而 Graph 是不允许复用的。

4.3 针对不同任务，定义不同的计算图

根据Oneflow Model Zoo的模型示例（https://github.com/Oneflow-Inc/models/blob/1b291f78d8f60e5f04ee0c5962e4611cc4bab40a/Vision/classification/image/alexnet/graph/train.py），train/eval等阶段可以创建不同的Graph子类。动态图下提供了 Module、Optimizer、Dataloader等模块，这些模型都可以被添加到 Graph 中。不同的组合可以构建不同类型的任务。

在这些不同阶段，Graph构造函数的行为、build函数的输入输出都有各自特点。了解这些，看后续代码时会更容易理解各个参数的具体含义。

• 构造函数

  • train阶段，需要添加Module、损失函数、优化器和dataloader

  • eval阶段，只需要添加Module和dataloader

• build函数

  • train

    • 导入样本和label

    • 调用Module得到前向计算结果

    • 计算损失

    • 计算梯度

    • 返回loss

  • eval

    • 导入样本和label

    • 调用Module得到预估结果

    • 返回预估结果和label

4.4 小结

上述几个类型的关系如下：

下面描述了GraphMyLinear的构造流程

* `__init__`
  * `Graph.__init__`
  * self.model = linear_model
    * `Graph.__setattr__`
      * _add_block
        * get_block_cls: 递归地把Module转为ModuleBlock
        * `ModuleBlock.__init__`
          * ModuleBlock.set_origin
            * `ModuleBlock._origin = origin` (Module)
            * 对origin的sub modules, parameters, buffers递归调用get_block_cls
            * `ModuleBlock.__setattr__`

5 

逻辑图的编译

计算机语言的编译，是将高级语言的语句编译为汇编或机器指令。深度学习框架对计算任务的编译，是将用户的特定语句操作转换为DAG图。oneflow中用Job（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job.proto#L30）描述逻辑的计算图。

不同于eager模式的动态图，静态图在开始执行前可以得到整个计算任务的所有信息，可以对DAG进行多轮优化。每轮优化都是输入一个Job、得到一个新Job。

最后，根据分布式环境配置，将逻辑图Job转换为物理执行的计算图Plan（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/plan.proto#L34）。在物理图中，一个op可能分布在多个节点/进程。

启动DAG计算需要调用Graph.__call__，这个函数的执行主要分以下几个步骤：

• __call__

  • _compile（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L221） if not _is_compiled

    • build_graph（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L741）

      • __build_graph（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L759）

    • finish_complie_and_init_runtime（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L742）

  • __run（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L226）

逻辑图编译主要在__build_graph中进行。finish_complie_and_init_runtime会继续做一些优化pass，然后构建物理图、初始化运行时Actor系统。__run会启动一次DAG的运算。

5.1 graph_build_context: 为逻辑图编译设置基本环境

在 Graph 中，build 函数里面的代码执行都在 graph_build_context 的作用域下，这样实现了动态转静态的功能。

__build_graph中的graph_build_context（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L851）虽然只有一行代码，但却做了几件非常重要的事情。

首先在context作用域内设置全局的lazy_mode为True（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L46）。在这个context作用域内，所有op都由LazyInterpreter解释执行。

其次，在JobBuildAndInferCtx（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L47）作用域内，JobBuildAndInferCtx_Open（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_buil