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

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

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

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

1 代码示例

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


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-In...）时，与静态图相关的主要操作如下：

GetEnv（https://github.com/Oneflow-In...）
- EnvGlobalObjectsScope::Init（https://github.com/Oneflow-In...）
  - 启动各个节点的控制面（https://github.com/Oneflow-In...）网络连接
  - 初始化VM（https://github.com/Oneflow-In...）
  - 启动各个节点的数据面网络连接（https://github.com/Oneflow-In...）
  - 初始化KernelObserver（https://github.com/Oneflow-In...）
NewDefaultSession（https://github.com/Oneflow-In...）
- RegsiterSession（https://github.com/Oneflow-In...）创建 Session，并注册为 default session（https://github.com/Oneflow-In...）
- 创建 Python MultiClientSession 并保存到dict（https://github.com/Oneflow-In...），但并不 TryInit
  - 创建 C++ MultiClientSessionContext（https://github.com/Oneflow-In...）但并不 TryInit

EnvGlobalObjectsScope::Init中先创建一个全局的ProcessCtx（https://github.com/Oneflow-In...）对象。然后根据环境变量等配置，在各个进程间创建gRPC和CommNet的连接，分别负责控制面和数据面的数据传输。其中在Bootstrap过程中会初始化全局的ProcessCtx（https://github.com/Oneflow-In...），给每个进程分配一个全局唯一的rank编号（https://github.com/Oneflow-In...）（machine_id（https://github.com/Oneflow-In...））。

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

3 Module类

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

Module（https://github.com/Oneflow-In...）是由op和tensor构成的、可复用的子模块。利用Module可以更高效、更快捷的构建复杂模型。oneflow.nn（https://github.com/Oneflow-In...）模块导出了很多预定义的Module。

Module定义了自己的属性设置逻辑（https://github.com/Oneflow-In...），核心逻辑是

如果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-In...）得到的session，就是导入oneflow包时NewDefaultSession（https://github.com/Oneflow-In...）构建的session。当时没有初始化，而是在Graph构造时进行初始化（https://github.com/Oneflow-In...）。对应的C++函数是MultiClientSessionContext::TryInit（https://github.com/Oneflow-In...），执行时会创建各种全局的资源管理器，比如：

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-In...），_add_block中主要是调用get_block_cls并保存结果（https://github.com/Oneflow-In...）。get_block_cls（https://github.com/Oneflow-In...）的作用是将Module及其所有Tensor属性都转为对应的Block对象。为什么要做这个动作呢？主要是静态图编译需要借助Block类型来实现代理执行的功能，这些功能不适合直接写到 eager 下的 Module 和 Tensor 上。
这个转换是在ModuleBlock构造时调用set_origin（https://github.com/Oneflow-In...）完成的。对于子Module，会递归调用get_block_cls函数（https://github.com/Oneflow-In...），这样所有子Module及其Tensor属性都会被转换为对应的Block对象。

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

Graph.__setattr__不允许将Tensor对象设置为属性（https://github.com/Oneflow-In...）。Tensor只能存到Module中，因为 Module 是做状态共享的基本单位，而 Graph 是不允许复用的。

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

根据Oneflow Model Zoo的模型示例（https://github.com/Oneflow-In...），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-In...）描述逻辑的计算图。

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

最后，根据分布式环境配置，将逻辑图Job转换为物理执行的计算图Plan（https://github.com/Oneflow-In...）。在物理图中，一个op可能分布在多个节点/进程。

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

call
- _compile（https://github.com/Oneflow-In...） if not _is_compiled
  - build_graph（https://github.com/Oneflow-In...）
    - __build_graph（https://github.com/Oneflow-In...）
  - finish_complie_and_init_runtime（https://github.com/Oneflow-In...）
- __run（https://github.com/Oneflow-In...）

逻辑图编译主要在__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-In...）虽然只有一行代码，但却做了几件非常重要的事情。

首先在context作用域内设置全局的lazy_mode为True（https://github.com/Oneflow-In...）。在这个context作用域内，所有op都由LazyInterpreter解释执行。

其次，在JobBuildAndInferCtx（https://github.com/Oneflow-In...）作用域内，JobBuildAndInferCtx_Open（https://github.com/Oneflow-In...）调用类似如下C++代码

// oneflow/api/python/job_build/job_build_and_infer.h
// oneflow/core/job/job_build_and_infer_ctx_mgr.cpp

// 如前所述，LazyJobBuildAndInferCtxMgr 在 MultiClientSessionContext::TryInit 执行时初始化。
// LazyJobBuildAndInferCtxMgr mgr;
mgr.OpenJobBuildAndInferCtx(job_name);

OpenJobBuildAndInferCtx会新建一个Job对象（https://github.com/Oneflow-In...）、一个LazyJobBuildAndInferCtx对象（https://github.com/Oneflow-In...）。LazyJobBuildAndInferCtx负责根据用户定制的op等操作，修改Job，其中最主要的功能是添加新 Op。

5.2 __build_io：为计算图添加input和output Op

self.__build_io("input", graph_build_util.build_graph_input_arg, *args, **kwargs)

上面这行代码（https://github.com/Oneflow-In...）的作用是，对于用户传递给graph_mylinear(input)的input参数，针对其中的每个tensor都在逻辑计算图中插入一个FeedInputOp（https://github.com/Oneflow-In...）节点。也就是说，model的输入（比如样本tensor，具体参考4.3节），在静态图中也视为一个op操作。

__build_io内会用args（即input）和kwargs构造一个ArgsTree。ArgsTree 把 Python 下的输入、输出抽象成了一个树，输入、输出可以是嵌套的 Tuple、List、Dict，元素是 Tensor，嵌套的结构刚好可以表示为树，而 Tensor 是树中的叶子节点。示例代码中kwargs是空的。

遍历ArgsTree，对args和kwargs的每个tensor都调用传入的build_func，对于input来说，就是build_graph_input_arg（https://github.com/Oneflow-In...）。后面会看到，model的output也会调用__build_io，所以这个函数名的意思应该就是对model的输入、输出进行静态图的构图工作。

build_graph_input_arg内部会构造一个FeedInputOpExpr（https://github.com/Oneflow-In...），提交给解释器执行。因为是在lazy作用域内，由LazyInterpreter解释执行（https://github.com/Oneflow-In...），LazyInterpreter会将对应的op插入静态图。

附：build input时ArgsTree的内部结构

5.2.1 将op添加到逻辑图

LazyInterpreter::ApplyImpl（https://github.com/Oneflow-In...）在执行时，GetCurInferCtx()（https://github.com/Oneflow-In...）返回的就是graph_build_context中OpenJobBuildAndInferCtx（https://github.com/Oneflow-In...）创建的那个LazyJobBuildAndInferCtx对象，这个对象负责逻辑图的构建。添加op的主要调用流程如下：

infer_ctx->AddAndInferConsistentOp（https://github.com/Oneflow-In...）
AddAndInferOp（https://github.com/Oneflow-In...）
ConstructOp（https://github.com/Oneflow-In...）
CheckAndConstructOp（https://github.com/Oneflow-In...）
NewObj（https://github.com/Oneflow-In...）

OperatorConf中，多种op配置共享op_type字段（https://github.com/Oneflow-In...），protobuf oneof的op_type_case常量作为注册NewObj的key。

系统预定义的op在oneflow/core/operator（https://github.com/Oneflow-In...）下，例如UserOp（https://github.com/Oneflow-In...）。

AddAndInferOp将返回的Operator保存到LazyJobBuildAndInferCtx的字典中。后续的函数调用，主要是进行推导并修改静态图Job，使得各个节点构成一个DAG。

JobBuildAndInferCtx相关的类关系如下：

5.2.2 lazy tensor 和 eager tensor 的区别

LazyInterpreter::ApplyImpl的最后，会调用BuildTensor（https://github.com/Oneflow-In...）构造一个lazy tensor，作为build_graph_input_arg的返回值（https://github.com/Oneflow-In...）。所以__build_io返回的lazy_args（https://github.com/Oneflow-In...）是lazy tensor，它将替代eager的args（https://github.com/Oneflow-In...）（也就是用户输入的input）参与后续的计算图构建。

那么lazy tensor和eager tensor的区别是什么呢？eager tensor是要即时计算的，所以需要真实数据；而lazy tensor仅在静态图编译阶段用于推导，只需要描述性质的元信息。静态图编译是在lazy模式下运行，只是使用lazy tensor 做计算机构图和校验。

后面会看到，静态图的运行期已经没有tensor的概念。运行期看到的只是更广义的Regst存储，可能代表tensor/blob，也可能是其它控制信息。静态图运行时的输入，是直接读取外部 eager tensor的内存数据到到regst；输出应该是op写到regst，通过blob构造eager tensor。

5.3 build: 将UserOp和FeedVariableOp添加到逻辑图

__build_graph中的self.build()（https://github.com/Oneflow-In...）会调用GraphMyLinear.build()，以及ModuleMyLinear.forward()。因为是在lazy模式下运行，matmul和add都会调用UserOpExpr重载版本的LazyInterpreter::ApplyImpl（https://github.com/Oneflow-In...），进而调用AddAndInferConsistentOp（https://github.com/Oneflow-In...）进行构图操作。

需要说明的是，在引用Module的Parameter属性时（如weight/bias），会触发FeedVariableOp的构图操作（https://github.com/Oneflow-In...）、调用对应版本的LazyInterpreter::ApplyImpl（https://github.com/Oneflow-In...）。这个是怎么执行的呢？

__build_graph中，在进入lazy模式之前，先调用了_create_states_builder（https://github.com/Oneflow-In...）。其中self._state()（https://github.com/Oneflow-In...）返回所有Module的所有Parameter（包括子Module）。

state_block的类型是TensorBlock（https://github.com/Oneflow-In...）。所有的state_block的lazy_origin_builder().method（https://github.com/Oneflow-In...）都被设置为调用build_graph_state（https://github.com/Oneflow-In...）。

给build_graph_state（https://github.com/Oneflow-In...）设置个断点能让整个调用过程显形，主要的调用栈如下：


-> out = graph_mylinear(input)
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/graph.py(221)__call__()
-> self._compile(*args, **kwargs)
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/graph.py(741)_compile()
-> _, eager_outputs = self.build_graph(*args, **kwargs)
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/graph.py(759)build_graph()
-> outputs = self.__build_graph(*args, **kwargs)
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/graph.py(864)__build_graph()
-> outputs = self.build(*lazy_args, **lazy_kwargs)
  /mnt/project/machine-learning/oneflow/oneflow/test.py(21)build()
-> return self.model(input)
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/block.py(234)__call__()
-> result = self.__block_forward(*args, **kwargs)
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/block.py(266)__block_forward()
-> result = self._origin.__class__.forward(self, *args, **kwargs)
  /mnt/project/machine-learning/oneflow/oneflow/test.py(11)forward()
-> return flow.matmul(input, self.weight) + self.bias
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/block.py(483)__getattr__()
-> p_state = self._get_from_states(name, "_parameters")
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/block.py(521)_get_from_states()
-> _s_block.try_build()
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/block.py(679)try_build()
-> self._lazy_origin_builder.try_build(self)
  /usr/local/lib64/python3.6/site-packages/oneflow/nn/graph/block.py(627)try_build()
-> self.result = self.method()
> /usr/local/lib64/python3.6/site-packages/oneflow/framework/graph_build_util.py(227)build_graph_state()
-> op_name, var_conf_str, ["in_0"], ["out_0"]

这个调用过程比较容易困扰的是，执行对象会在Grpah、GraphMyLinear、ModuleMyLinear、ModuleBlock之间切换。

前面在讨论Graph的构造时已经提过，执行self.model(input)时，Graph.__getattr__返回的属性model是ModuleBlock对象，所以实际调用的是ModuleBlock.__call__。

在这个函数内调用__block_forward（https://github.com/Oneflow-In...），其中的_origin（https://github.com/Oneflow-In...）是ModuleMyLinear，进入到它的forward方法，执行到flow.matmul(input, self.weight) + self.bias时，matmul 会被LazyOpInterpreter 所执行，在 LazyOpInterpreter 中调用 AddAndInferConsistentOp（https://github.com/Oneflow-In...）
，在 Job 中添加一个 matmul operator。同理后面的加法会在 job 中添加一个 add operator。

self.weight 和 self.bias 会触发调用ModuleBlock.__getattr__，进而调用_get_from_states（https://github.com/Oneflow-In...），调用TensorBlock.try_build()（https://github.com/Oneflow-In...）。这里执行的就是进入lazy模式之前设置的build_graph_state（https://github.com/Oneflow-In...）。从而增加一个FeedVariableOp到计算图（https://github.com/Oneflow-In...）。为什么设置和调用会距离这么远呢？主要是为了让参数尽量和消费参数的 Operator 在一个作用域下，所以实现成了惰性求值来达到延迟计算的目的。

再后面的步骤就是调用__build_io（https://github.com/Oneflow-In...）插入FetchOutputOp（https://github.com/Oneflow-In...）。也就是说，获取model的output也是一个op。

到目前为止，前向计算图就构建完成了。它的json表示可以参考附录。net.op是计算图的节点，通过input等属性可以看出节点之间的连接关系。

示例代码的前向计算图如下。从这个图可以看到，input、output、weights等都是op。

5.4 逻辑图优化

在__build_graph中会调用CurJobBuildAndInferCtx_Complete对静态图进行多轮优化（https://github.com/Oneflow-In...），对应的C++函数是LazyJobBuildAndInferCtx::Complete()（https://github.com/Oneflow-In...）。

这之后生成的Job是full_job。本文的示例代码比较简单，并不是典型的计算场景，其forwar和ful计算图的拓扑是一样的。实际大部的图优化都实现在这个阶段，如 Op fusion、AMP、ZeRO、常量折叠等等。

到这里，逻辑图构建的主体部分就结束了。

随后会构建一个CNNGraph对象（https://github.com/Oneflow-In...），对应的C++类型是NNGraph（https://github.com/Oneflow-In...）。这个对象将负责构建物理计算图Plan。它也是整个运行时的拥有者和维护者。这个对象析构时，整个运行时也会有序终止并释放资源。

5.5 物理图的编译

接下来就是执行finish_complie_and_init_runtime（https://github.com/Oneflow-In...），其中的核心调用是self._c_nn_graph.complie_and_init_runtime()（https://github.com/Oneflow-In...），对应的C++函数是NNGraph::CompileAndInitRuntime（https://github.com/Oneflow-In...）。

在这个函数中，JobCompleter().Complete()（https://github.com/Oneflow-In...）会继续对逻辑图做几轮修改优化，补全 Runtime 执行所需要的附加信息，Compiler().Compile()（https://github.com/Oneflow-In...）将逻辑图转为分设备的物理图，并继续对Plan进行修改优化。

Plan的编译是在master节点进行的（https://github.com/Oneflow-In...）。master节点会将Plan通过gRPC推送给各个worker节点（https://github.com/Oneflow-In...），worker节点从master拉取物理计算图（https://github.com/Oneflow-In...）。

之后调用NewRuntimeBuffers创建Buffer对象（https://github.com/Oneflow-In...），Buffer应该是主要用于进程内的信息同步。

然后就准备初始化运行时了。

示例代码生成的compiled_job和物理图Plan的json参见附录。

最终生成的compiled逻辑图如下。框架自动插入了很多系统控制节点。

5.6 Plan的结构

示例代码输出的Plan json数据见附录。

Plan在逻辑上和compiled_job是等价的。这里主要关注task/op之间的关系。
Plan.task中的每个元素是一个task，其中的exec_sequence.exec_node对应job中的op，通常只有一个op（数组可以支持sub graph）。

exec_node.kernel_conf.op_attribute描述了op信息。其中op_conf包含op name信息。

kernel_conf.op_attribute.op_conf就是Job中的OperatorConf。

kernel_conf.op_attribute.arg_signature.bn_in_op2lbi体现了task/op之间的连接关系。

bn_in_op就是blob name in op，即op输入的blob name。

以System-AutoTick-DstSubsetTick_21为例

{
  "out": {
    "op_name": "System-AutoTick-DstSubsetTick_21",
    "blob_name": "out"
  },
  "in_0": {
    "op_name": "System-EagerCriticalSection-Interface-End-Tick-19",
    "blob_name": "out"
  },
  "in_1": {
    "op_name": "System-AutoTick-SrcSubsetTick_20",
    "blob_name": "out"
  }
}

exec_node.bn_in_op2regst_desc_id在task层面体现了连接关系。这个map中的key表示输入输出，value是register id。

{
"out": "29",
"in_0": "27",
"in_1": "28"
}

task.produced_regst_desc描述了对应task生产的register，consumer_task_id是消费者，
produced_regst_desc.out.regst_desc_type.data_regst_desc.lbi2blob_desc.lbi就是这个register的logic blob id。

task.consumed_regst_desc_id描述了对应task消费的register信息

6 运行时的初始化

NNGraph::CompileAndInitRuntime中，new Runtime这行代码会初始化运行时（https://github.com/Oneflow-In...）。主要做的事情包括：

创建Thread
通知Thread创建Actor，Actor会创建Regst和Kernel
给没有输入的source_tasks发送启动信号kStart

6.1 Runtime创建Thread

在Runtime的构造函数中，DumpThreadIdsFromPlan（https://github.com/Oneflow-In...）会将Plan中属于当前进程的task的thread id存入thread_ids_变量。AddThreads创建这些Thread对象（https://github.com/Oneflow-In...）。

Thread在构造时会创建一个物理线程（ https://github.com/Oneflow-In...），线程执行的是PollMsgChannel方法（https://github.com/Oneflow-In...），Thread就是在这里持续等待需要处理的新消息。

Thread只处理两类命令消息：线程终止消息，创建Actor的消息。其它消息交给Actor::ProcessMsg处理（https://github.com/Oneflow-In...）。

6.2 Runtime通知Thread创建Actor

在Runtime的构造函数中，tasks被分为两类：source_tasks和other_tasks。在示例代码中，source_tasks（https://github.com/Oneflow-In...）是没有输入边的task。

从代码逻辑看，在Plan proto中，task的consumed_regst_desc_id字段是一个map。如果这个map的所有key都是in_ctrl（https://github.com/Oneflow-In...），这个task就是source_tasks。

一些source_tasks的示例如下：

System-Src-WaitAndSendIds_16
System-AutoTick-AppendDeviceTick_9
System-EagerCriticalSection-Interface-End-Tick-19
System-EagerCriticalSection-Interface-End-Tick-25

Runtime调用HandoutTasks函数（https://github.com/Oneflow-In...）会给ActorMsgBus发送构建Actor的kConstructActor消息（https://github.com/Oneflow-In...）。

6.3 ActorMsgBus和Thread的消息处理

从接口看，ActorMsgBus （https://github.com/Oneflow-In...）负责消息的发送（Actor通过ActorMsgBus发送消息），Thread::PollMsgChannel（https://github.com/Oneflow-In...）负责消息的接收和处理。

7 Actor

7.1 Actor的创建

Thread在创建Actor时，会先尝试创建为LightActor（https://github.com/Oneflow-In...），如果不成功，再尝试用预先注册的工厂创建Actor。

有几种TaskType可以用于LightActor（https://github.com/Oneflow-In...）：

kNormalForward，比如matmul、add等user op。
kCopyHd
kTick
kCollectiveBoxingGeneric

目前大约有20多种Actor的子类型。其它Actor类型根据TaskType（https://github.com/Oneflow-In...）预先注册。例如WaitAndSendIdsActor。

示例代码的各个节点对应的actor类型参见附录。

Actor相关的类关系如下（包含关系只是表示可以访问到相关信息，并不意味着创建或着拥有该类型对象）

7.2 Actor的初始化

Actor的构造函数一般都是空的，构建之后需要执行Init（https://github.com/Oneflow-In...）函数进行初始化。

LightActor继承自ActorBase，不是Actor的子类，有自己的Init函数实现。这里只讨论Actor的初始化。

在Actor::Init（https://github.com/Oneflow-In...）中，首先调用ConstructKernel（https://github.com/Oneflow-In...）创建kernel实例。和Operator类似，kernel也是以OpTypeCase作为注册的key，例如WaitAndSendIdsKernel（https://github.com/Oneflow-In...）。一个Actor通常只有一个kernel。

之后调用NewRegsts创建Regst（https://github.com/Oneflow-In...）。Tensor是用户侧的概念。对应的运行时概念是Regst（https://github.com/Oneflow-In...），它持有Kernel需要读写的内存。Regst的概念比Tensor更宽泛，比如框架自动添加的控制Op也会用到Regst。

Actor将自己创建的Regst保存到produced_regsts_（https://github.com/Oneflow-In...）。

TakeOverNaiveConsumed（https://github.com/Oneflow-In...）只记录需要消费的regst id，但并不push到consumed_regsts_。

TakeOverNaiveProduced（https://github.com/Oneflow-In...）既记录生产的regst id，也push到naive_produced_rs_（https://github.com/Oneflow-In...）。这种区别是为了首次执行计算时，actor能顺利执行。后面分析Actor的消息处理时会再回过头来讨论一下。

调用InitBnInOp2BlobInfo会初始化BlobInfo（https://github.com/Oneflow-In...）。

之后就是调用VirtualActorInit（https://github.com/Oneflow-In...），这里允许各个Actor子类定制自己的初始化逻辑。通常会调用OF_SET_MSG_HANDLER宏（https://github.com/Oneflow-In...）设置Actor的消息处理函数。

7.3 Actor的消息处理

LightActor 首先会根据消息类型分别处理 kRegstMsg 和 kEordMsg 消息。HandleRegstMsg（https://github.com/Oneflow-In...）中根据 RegstMsg 的 type (kProduced 或 kComsumed) 来分别处理各种读写状态计数。

然后判断读写计数是否达到了判断条件，如果达到了意味着满足了读写 regst 的条件，然后就执行 ActOnce（https://github.com/Oneflow-In...）。

LightActor::ActOnce 会在第一次执行时去 InitBnInOp2Blob 和 InitActMsg。InitBnInOp2Blob 初始化 resgt 中的 bn 与 Blob 的映射关系，为 kernel 提供通过 bn 访问 Blob 的功能。InitActMsg 会初始化好所有需要发送的消息避免后继发消息时重复的构建消息。

然后就是 LaunchKernel，接着会 ResetState 重置 regst 状态。

LaunchKernel 后就会把之前构建好的消息发送出去，同步消息会直接入队 thread 消息队列，异步消息通过 callback 发送到 ActorMsgBus。

普通 Actor::ProcessMsg 会调用 msg handler 来处理消息，最常见的 msg handler 就是 Actor::HandlerNormal（https://github.com/Oneflow-In...）。

Actor::HandlerNormal 中流程跟 LightActor 中类似，会根据不同的 regst 类型来分别处理，Actor 中对 regst 的状态管理方式与 LightActor 不同，LightActor 中的方式更加高效，Actor 中能处理一些特殊情况。

消息处理完毕后，就会调用 ActUntilFail，ActUntilFail 会判断 IsReadReady 和 IsWriteReady 来决定是否可以进行 Act。

最常见的 NaiveActor::Act() 就是执行 AsyncLaunchKernel。

Act 完成后，就开始朝上下游发送 regst 消息。

还有一些特殊的 Actor，我们以WaitAndSendIdsActor为例，观察一下这类Actor的消息处理机制。

之所以选择这个例子，一是这个Actor比较简单；二是这是一个典型的source task，想看一下计算图是怎么被触发启动计算的。

Thread收到的消息如果不是kStopThread或kConstructActor，就调用Actor::ProcessMsg（https://github.com/Oneflow-In...），将消息转给Actor处理。

ProcessMsg函数只是简单的将消息转给handler处理（https://github.com/Oneflow-In...）。

WaitAndSendIdsActor::VirtualActorInit中，handler被设置为HandlerWaitToStart（https://github.com/Oneflow-In...）。

Runtime的构造函数中，发送的第一批消息是给source_tasks的kStart消息，这个消息就由HandlerWaitToStart函数处理。

HandlerWaitToStart校验消息类型后，将handler设置为HandlerNormal（https://github.com/Oneflow-In...）（这也是大部分Actor的默认handler），然后调用ProcessMsg（https://github.com/Oneflow-In...），实际就是调用新设置的handler HandlerNormal。

HandlerNormal中，如果是kCmdMsg，只允许是kStart（https://github.com/Oneflow-In...）。通过消息类型校验后，会直接调用ActUntilFail（https://github.com/Oneflow-In...）。

7.4 Act执行的条件

LightActor 和 Actor 判断能否进行 Act 采用了不同的策略，LightActor 的效率更高，Actor 能处理一些特殊情况。

对于 LightActor，当在读的register计数 total_reading_cnt_ 归 0，可消费的register计数 ready_consumed_ 增加到 max_ready_consumed_，前者表示所有的消费者已经读取当前 LightActor 的 Regst，后者表示当前 LightActor 消费的所有 Regst 已经到达（由上游发送的 Regst 消息）。

对于 Actor，Actor::ActUntilFail中，Act方法（https://github.com/Oneflow-In...）是各个子类自己实现的，一般主要是启动kernel计算。

但是在执行Act之前，需要先确认：

Act执行依赖的数据是否都已经就绪？（IsReadReady）
Act生产出来的数据，消费方是否已经用完、并收到ack消息确认？（IsWriteReady）

Actor有4个与此相关的成员变量

RegstSlot naive_produced_rs_;
RegstSlot inplace_produced_rs_;
RegstSlot naive_consumed_rs_;
RegstSlot inplace_consumed_rs_;

xx_produced_rs_存储的是当前Actor的下游consumer返回的、已经使用完毕的ack regst信息。（当前Actor生产的Regst存储在produced_regsts_中。）
运行时在初始化的过程中，所有Actor都没有运行过，任何Actor都不可能收到ack消息，所以在Actor初始化时，要预先填充xx_produced_rs_，这样才能保证Actor在首次运行前是WriteReady的，才能顺利启动执行。

xx_consumed_rs_存储的是上游依赖发来的数据。它不需要预先填充。因为source_tasks没有输入依赖，自然就是ReadReady的；而xx_produced_rs_在初始化时的预先填充又保证它是WriteReady的，所以source_tasks可以直接运行。source_tasks的输出消息发给下游，下游也会变为ReadReady，而下游在初始化后也保证是WriteReady的。整个Actor系统就可以这样运转起来了。

7.5 Actor上下游之间的通知机制

Act执行完毕后，需要将结果数据发给下游consumer。以 WaitAndSendIds 的 Naive Produced 为例，ActUntilFail中的调用流程如下：

AsyncSendNaiveProducedRegstMsgToConsumer（https://github.com/Oneflow-In...）
- VirtualAsyncSendNaiveProducedRegstMsgToConsumer（https://github.com/Oneflow-In...）
  - HandleProducedNaiveDataRegstToConsumer（https://github.com/Oneflow-In...）
    - HandleRegstToConsumer（https://github.com/Oneflow-In...）
      - EnqueueAsyncMsg（https://github.com/Oneflow-In...）
        
        如果目标线程是当前线程，ActorMsgBus::SendMsg（https://github.com/Oneflow-In...）
        
        否则，将消息加入async_msg_queue_（https://github.com/Oneflow-In...）
      - 增加 total_reading_cnt_（https://github.com/Oneflow-In...）（这个变量表示已经发消息给下游、但未收到的ack数量）
    - naive_produced_rs_.PopFrontRegsts（https://github.com/Oneflow-In...）
  - AsyncSendProducedCtrlRegstMsgToConsumer

注意naive_produced_rs_.PopFrontRegsts（https://github.com/Oneflow-In...）会将Regst指针从队列中删掉，相应的可用（https://github.com/Oneflow-In...）register计数减1（https://github.com/Oneflow-In...）。

而在Actor::HandlerNormal中处理收到的kRegstMsg消息（https://github.com/Oneflow-In...）时，如果是consumer发来的ack消息，会调用TryUpdtStateAsProducedRegst（https://github.com/Oneflow-In...），将Regst再添加到 naive_produced_rs_ 中（https://github.com/Oneflow-In...），以保证当前Actor在收到所有ack后是WriteReady的；同时递减在读的 register 计数total_reading_cnt_。

Actor对依赖的上游消息的处理是类似的。通过以下函数调用给上游发送ack消息、通知 register 已经用完，可以继续更新了：

AsyncSendNaiveConsumedRegstMsgToProducer（https://github.com/Oneflow-In...）
AsyncRetInplaceConsumedRegstIfNoConsumer（https://github.com/Oneflow-In...）在Actor::HandlerNormal中收到kRegstMsg消息后，将消息添加到consumed_rs_（https://github.com/Oneflow-In...），以保证当前Actor在收到所有依赖数据后是ReadReady的。

LightActor有自己的消息处理机制（https://github.com/Oneflow-In...），大致原理应该是差不多的。

7.6 Act执行的动作

根据上述讨论，Actor收到kRegstMsg后也会进入ActUntilFail执行。如果读写都是Ready，就执行Act（https://github.com/Oneflow-In...）。以WaitAndSendIdsActor为例，主要调用链路如下：

AsyncLaunchKernel（https://github.com/Oneflow-In...）
ek.kernel->Launch（https://github.com/Oneflow-In...），启动Kernel计算
Forward（https://github.com/Oneflow-In...）
ForwardDataContent（https://github.com/Oneflow-In...）
buffer->Pull（https://github.com/Oneflow-In...）
给regst的存储地址mut_dptr赋值（https://github.com/Oneflow-In...）

buffer->Pull会等待条件变量的通知（https://github.com/Oneflow-In...）。现在，看上去所有Actor都已准备就绪，只等发令枪一响就开跑了。

8 启动静态图的计算

Graph.__run（https://github.com/Oneflow-In...）会扣动发令枪的板机，启动计算图的一轮计算。

主要调用流程如下：

RunLazyNNGraph（https://github.com/Oneflow-In...）
builder->LaunchLazyJob（https://github.com/Oneflow-In...）
LaunchLazyJobInstructionType（https://github.com/Oneflow-In...）
Buffer::Push（https://github.com/Oneflow-In...）

这里的Buffer::Push就是WaitAndSendIdsKernel在等待的起跑信号。

9 运行时的退出机制

整个运行时包含很多对象和资源，安全有序的退出是庞杂而又细致的工作。这里仅以WaitAndSendIds为例，从一个侧面观察一下运行时的退出机制。

运行时的退出始于NNGraph对象的析构（https://github.com/Oneflow-In...）。

9.1 Actor的退出

NNGraph在析构时，会关闭所有的Buffer对象（https://github.com/Oneflow-In...）。
Buffer在关闭时，会设置is_closed_ = true并通知所有监听者（https://github.com/Oneflow-In...）。但是Pull会继续处理完已经提交的计算。
- 所以，Buffer应该是主要用于进程内的通信和异步协调的一个类。
WaitAndSendIdsKernel这时候正在等待新一轮计算开始（https://github.com/Oneflow-In...），结果收到Pull返回的kBufferStatusErrorClosed（https://github.com/Oneflow-In...）。
WaitAndSendIdsActor::IsCustomizedReadReady以后就一直返回false（https://github.com/Oneflow-In...），IsReadReady也返回false（https://github.com/Oneflow-In...）。
- 这之后，ActUntilFail只会执行异步消息发送（https://github.com/Oneflow-In...）（不再进入while循环）
WaitAndSendIdsActor::HandlerNormal仍然会处理其它Actor发来的消息（https://github.com/Oneflow-In...）。但因为IsCustomizedReadReady返回false，会进入AsyncSendEORDMsgForAllProducedRegstDesc（https://github.com/Oneflow-In...）执行。它会给每个下游发送kEordMsg消息（https://github.com/Oneflow-In...）。
Actor在收到上游发来的kEordMsg消息后，递减remaining_eord_cnt_（https://github.com/Oneflow-In...）。
- remaining_eord_cnt_被初始化为Actor的输入regst的数量（https://github.com/Oneflow-In...）。
total_reading_cnt_是当前Actor生产的、已经发给consumer、但尚未收到ack的消息数量。
- Actor目前仍可以正常接收consumer发来的ack消息。
当上述2个变量都为0时（https://github.com/Oneflow-In...），意味着所有上游都发出了kEordMsg消息，也收到了所有下游的ack消息。Actor就给Thread返回1（https://github.com/Oneflow-In...）。
- 如果上述两个变量有不为0的，就修改handler，由HandlerZombie（https://github.com/Oneflow-In...）处理后续收到的消息。
Thread收到Actor返回的1后（https://github.com/Oneflow-In...），将它从自己的存储中删除（https://github.com/Oneflow-In...），并递减运行Actor的数量。

9.2 Thread的退出

NNGraph重置runtime_导致运行时对象被析构（https://github.com/Oneflow-In...）。
Runtime删除所有Thread（https://github.com/Oneflow-In...）。
ThreadMgr给所有Thread发送kStopThread消息（https://github.com/Oneflow-In...）。同时，重置指针导致Thread析构（https://github.com/Oneflow-In...）。
Thread的物理线程退出PollMsgChannel循环（https://github.com/Oneflow-In...）。
Thread等待物理线程结束，关闭channel（https://github.com/Oneflow-In...）。

10 分布式场景的静态图

分布式的compile_job、物理图Plan和单机场景有明显变化。

比如，每个进程都有一套WaitAndSendIds等控制节点。这也容易理解，因为每个节点都要执行__run和Buffer::Push/Pull，都要启动本进程的Actors执行计算。

matmul和broadcast_add等user op也会在两个节点进行计算。

10.1 示例代码

启动方式参考Global Tensor的官方文档。

import oneflow as flow
import oneflow.nn as nn

P0 = flow.placement("cpu", ranks=[0, 1])
a0_sbp = flow.sbp.split(0)

class ModuleMyLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(flow.randn(in_features, out_features,
            placement=P0, sbp=flow.sbp.broadcast))
        self.bias = nn.Parameter(flow.randn(1, out_features,
            placement=P0, sbp=flow.sbp.broadcast))

    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(5, 4, placement=P0, sbp=flow.sbp.split(1))
out = graph_mylinear(input)
print(out)

11 附录

11.1 断点

11.1.1 Python断点示例

# python3 -m pdb test.py
break test.py:25
break oneflow/nn/graph/graph.py:221
break oneflow/nn/graph/graph.py:741
break oneflow/nn/graph/graph.py:745
break oneflow/nn/graph/graph.py:759
break oneflow/nn/graph/graph.py:828
break oneflow/nn/graph/graph.py:777
break oneflow/nn/graph/graph.py:1066
break oneflow/nn/graph/graph.py:1133
break oneflow/framework/graph_build_util.py:227

11.1.2 C++断点示例

启动命令

source /mnt/oneflow/build/source.sh
gdb --args python3 /mnt/oneflow/test.py
# set breakpoints
# run

断点示例


set breakpoint pending on
break oneflow::ActorMsg::BuildEordMsg
break oneflow/core/common/buffer.h:80
break oneflow::(anonymous namespace)::CheckAndConstructOp
break oneflow::WaitAndSendIdsActor::Act
break oneflow::WaitAndSendIdsActor::HandlerWaitToStart
break oneflow/core/lazy/actor/light_actor.cpp:452
break oneflow/core/lazy/actor/light_actor.cpp:485
break oneflow::ForeignInputKernel::ForwardDataContent
break oneflow::vm::LaunchLazyJobInstructionType::Compute

11.2 静态图的json表示

11.3 actor type

naive_actor


System-AutoTick-AppendDeviceTick_9
System-AutoTick-DstSubsetTick_12
System-AutoTick-DstSubsetTick_21
System-AutoTick-DstSubsetTick_27
System-AutoTick-Prepend-DeviceTick_7
System-AutoTick-SrcSubsetTick_20
System-AutoTick-SrcSubsetTick_26
System-AutoTick-SrcSubsetTick_8
System-AutoTick-Tick_11
System-AutoTick-Tick_13
System-EagerCriticalSection-Callback-23
System-EagerCriticalSection-Callback-29
System-EagerCriticalSection-Interface-Begin-Tick-18
System-EagerCriticalSection-Interface-Begin-Tick-24
System-EagerCriticalSection-Interface-End-Tick-19
System-EagerCriticalSection-Interface-End-Tick-25
System-EagerCriticalSection-Wait-22
System-EagerCriticalSection-Wait-28

light_actor

_GraphMyLinear_0_input.0.0_2
_GraphMyLinear_0_output.0.0_2
model.bias
model-broadcast_add-1
model-matmul-0
model.weight
System-AutoTick-SinkTick_15
System-SyncAllRanksSinkTick_14

wait_and_send_ids_actor
System-Src-WaitAndSendIds_16

call_back_notify_actor
System-Sink-CallbackNotify_17

12 参考资料

oneflow v0.8.0（https://github.com/Oneflow-In...）
OneFlow框架的系统设计（上篇）（https://zhuanlan.zhihu.com/p/...）
OneFlow框架的系统设计（中篇）（https://zhuanlan.zhihu.com/p/...）
OneFlow框架的系统设计（下篇）（https://zhuanlan.zhihu.com/p/...）
一个Job在OneFlow中的执行过程—上篇（https://zhuanlan.zhihu.com/p/...）
一个Job在OneFlow中的执行过程—中篇（https://zhuanlan.zhihu.com/p/...）
一个Job在OneFlow中的执行过程—下篇（https://zhuanlan.zhihu.com/p/...）
静态图模块 nn.Graph（https://docs.oneflow.org/mast...）
OneFlow系统设计（https://docs.oneflow.org/v0.4...）
torch.nn.Module（https://pytorch.org/docs/1.10...）

欢迎 Star、试用 OneFlow 最新版本：https://github.com/Oneflow-In...

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

1

代码示例

2

oneflow包的初始化

3

Module类

4

Graph类

5

逻辑图的编译

6

运行时的初始化

7

Actor

8

启动静态图的计算

9

运行时的退出机制

10

分布式场景的静态图

11

附录

12

参考资料

你可能感兴趣的:(人工智能深度学习)