OneFlow源码阅读7：静态图与运行时

发布时间：2023-03-17 17:00

oneflow静态图的训练效率远高于动态图（eager模式）。本文试图通过一个简单例子，结合0.8.0的代码，看一下静态图和运行时的实现机制。
在开始之前，建议先读一下参考资料中《OneFlow框架的系统设计》等系列文章。对静态图、运行时的基本概念和设计理念有基本的了解，会更容易理解代码。

1 代码示例

下面的示例代码来自官方文档，是一个线性模型的前向计算。后续主要基于这段代码进行分析。

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在初始化包时，与静态图相关的主要操作如下：

EnvGlobalObjectsScope::Init中先创建一个全局的ProcessCtx对象。然后根据环境变量等配置，在各个进程间创建gRPC和CommNet的连接，分别负责控制面和数据面的数据传输。其中在Bootstrap过程中会初始化全局的ProcessCtx，给每个进程分配一个全局唯一的rank编号（machine_id）。
本文不涉及网络层面的操作，只讨论同一进程内各线程间的交互。

3 Module类

虽然可以直接用op和tensor构造模型，但是op的粒度太细了，直接用op构造模型会比较繁琐。
Module是由op和tensor构成的、可复用的子模块。利用Module可以更高效、更快捷的构建复杂模型。oneflow.nn模块导出了很多预定义的Module。

Module定义了自己的属性设置逻辑，核心逻辑是

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

  • LazyJobBuildAndInferCtxMgr
  • BufferMgr
  • RegstMgr
  • ActorMsgBus
  • ThreadMgr

4.2 __setattr__: 将Module和Tensor封装为Block

Graph.__setattr__不允许将Tensor对象设置为属性。Tensor只能存到Module中。
setattr最重要的动作就是对_add_block的调用，_add_block中主要是调用get_block_cls并保存结果。
get_block_cls的作用是将Module及其所有Tensor属性都转为对应的Block对象。为什么要做这个动作呢？主要是静态图编译需要借助Block类型的一些功能，这些功能不适合直接加诸于Module和Tensor。
这个转换是在ModuleBlock构造时调用set_origin完成的。对于子Module，会递归调用get_block_cls函数，这样所有子Module及其Tensor属性都会被转换为对应的Block对象。
所以，上述示例代码中，GraphMyLinear实际存储的是ModuleBlock，Graph.build执行时获取的model属性也是ModuleBlock对象，ModuleBlock.origin才是ModuleMyLinear。

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

根据Oneflow Model Zoo的模型示例，train/eval/infer等阶段可以创建不同的Graph子类。在这些不同阶段，Graph构造函数的行为、build函数的输入输出都有各自特点。了解这些，看后续代码时会更容易理解各个参数的具体含义。

  • 构造函数

    • train阶段，构造函数会引入Module、损失函数、优化器和data_loader
    • eval阶段，只需要引入Module和data_loader
    • infer阶段，只需要引入Module

  • build函数

    • train

      • 导入样本和label
      • 调用Module得到前向计算结果
      • 计算损失
      • 计算梯度
      • 返回loss

    • eval

      • 导入样本和label
      • 调用Module得到预估结果
      • 返回预估结果和label

    • infer

      • 传入的参数包括样本
      • 调用Module得到预估结果
      • 返回预估结果

4.4 小结

上述几个类型的关系如下：
OneFlow源码阅读7：静态图与运行时_第1张图片

下面描述了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描述逻辑的计算图。
不同于eager模式的动态图，静态图在开始执行前可以得到整个计算任务的所有信息，可以对DAG进行多轮优化。每轮优化都是输入一个Job、得到一个新Job。
最后，根据分布式环境配置，将逻辑图Job转换为物理执行的计算图Plan。在物理图中，一个op可能分布在多个节点/进程。

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

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

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

__build_graph中的graph_build_context虽然只有一行代码，但却做了几件非常重要的事情。

首先在context作用域内设置全局的lazy_mode为True。在这个context作用域内，所有op都由LazyInterpreter解释执行。

其次，在JobBuildAndInferCtx作用域内，JobBuildAndInferCtx_Open调用类似如下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对象、一个LazyJobBuildAndInferCtx对象。LazyJobBuildAndInferCtx负责根据用户定制的op等操作，修改Job。

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

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

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

__build_io内会用args（即input）和kwargs构造一个ArgsTree。示例代码中kwargs是空的。
然后遍历ArgsTree，对args和kwargs的每个tensor都调用传入的build_func，对于input来说，就是build_graph_input_arg。后面会看到，model的output也会调用__build_io，所以这个函数名的意思应该就是对model的输入、输出进行静态图的构图工作。

build_graph_input_arg内部会构造一个FeedInputOpExpr，提交给解释器执行。因为是在lazy作用域内，由LazyInterpreter解释执行，LazyInterpreter会将对应的op插入静态图。

附： build input时ArgsTree的内部结构
OneFlow源码阅读7：静态图与运行时_第2张图片

__build_io(input) 中 ArgsTree 的内部数据组织示意

  • _named_io_args: NamedArg

    • _value: tuple

      • [0]: NamedArg

        • _value: tuple of NamedArg

          • [0]: NamedArg

            • _value: args tensor from Graph.__call__

      • [1]: NamedArg

        • _value: empty kwargs from Graph.__call__

通过pdb命令可以查看变量: p args_tree._named_io_args._value[0]._value[0]._value.to_numpy()

5.2.1 将op添加到逻辑图

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

OperatorConf中，多种op配置共享op_type字段，protobuf oneof的op_type_case常量作为注册NewObj的key。
系统预定义的op在oneflow/core/operator下，例如UserOp。

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

JobBuildAndInferCtx相关的类关系如下：
OneFlow源码阅读7：静态图与运行时_第3张图片

5.2.2 lazy tensor 和 eager tensor 的区别

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

那么lazy tensor和eager tensor的区别是什么呢？我的理解是，eager tensor是要即时计算的，所以需要携带数据；而lazy tensor仅在静态图编译阶段用于推导，只需要描述性质的元信息，不需要类似样本那样的数据。同时，静态图编译是在lazy模式下运行，使用lazy tensor在各种检查校验时应该会更顺畅（？）。
后面会看到，静态图的运行期已经没有tensor的概念。运行期看到的只是更广义的Regst存储，可能代表tensor/blob，也可能是其它控制信息。静态图运行时的输入，应该是通过op直接读取tensor的blob（或者复用地址？）到regst；输出应该是op写到regst，通过blob构造eager tensor。

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

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

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

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

state_block的类型是TensorBlock。所有的state_block的lazy_origin_builder().method都被设置为调用build_graph_state。

给build_graph_state设置个断点能让整个调用过程显形，主要的调用栈如下：

-> 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，其中的_origin是ModuleMyLinear，进入到它的forward方法，执行到flow.matmul(input, self.weight) + self.bias时，会触发调用ModuleBlock.__getattr__，进而调用_get_from_states，调用TensorBlock.try_build()。这里执行的就是进入lazy模式之前设置的build_graph_state。从而增加一个FeedVariableOp到计算图。

为什么设置和调用会距离这么远呢？可能是因为，如果在引用weights等Parameter时再try_build，会不方便处理多个block共享一个tensor的情况。
再后面的步骤就是调用__build_io插入FetchOutputOp。也就是说，获取model的output也是一个op。

到目前为止，前向计算图就构建完成了。它的json表示可以参考附录。net.op是计算图的节点，通过input等属性可以看出节点之间的连接关系。
示例代码的前向计算图如下。从这个图可以看到，input、output、weights等都是op。
OneFlow源码阅读7：静态图与运行时_第4张图片

5.4 逻辑图优化

在__build_graph中会调用CurJobBuildAndInferCtx_Complete对静态图进行多轮优化，对应的C++函数是LazyJobBuildAndInferCtx::Complete()。

这之后生成的Job是full_job。本文的示例代码比较简单，并不是典型的计算场景，其forwar和ful计算图l的拓扑是一样的。
到这里，逻辑图构建的主体部分就结束了。

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

5.5 物理图的编译

接下来就是执行finish_complie_and_init_runtime，其中的核心调用是self._c_nn_graph.complie_and_init_runtime()，对应的C++函数是NNGraph::CompileAndInitRuntime。
在这个函数中，JobCompleter().Complete()会继续对逻辑图做几轮修改优化，Compiler().Compile()将逻辑图转为物理图，并继续对Plan进行修改优化。

Plan的编译是在master节点进行的。master节点会将Plan通过gRPC推送给各个worker节点，worker节点从master拉取物理计算图。

之后调用NewRuntimeBuffers创建Buffer对象，Buffer应该是主要用于进程内的信息同步。

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

示例代码生成的compiled_job和物理图Plan的json参见附录。
最终生成的compiled逻辑图如下。框架自动插入了很多系统控制节点。
OneFlow源码阅读7：静态图与运行时_第5张图片

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这行代码会初始化运行时。主要做的事情包括：

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

6.1 Runtime创建Thread

在Runtime的构造函数中，DumpThreadIdsFromPlan会将Plan中属于当前进程的task的thread id存入thread_ids_变量。AddThreads创建这些Thread对象。
Thread在构造时会创建一个物理线程，线程执行的是PollMsgChannel方法，Thread就是在这里持续等待需要处理的新消息。
Thread只处理两类命令消息：线程终止消息，创建Actor的消息。其它消息交给Actor::ProcessMsg处理。

6.2 Runtime通知Thread创建Actor

在Runtime的构造函数中，tasks被分为两类：source_tasks和other_tasks。在示例代码中，source_tasks是没有输入边的task。
从代码逻辑看，在Plan proto中，task的consumed_regst_desc_id字段是一个map。如果这个map的所有key都是in_ctrl，这个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函数会给ActorMsgBus发送构建Actor的kConstructActor消息。

6.3 ActorMsgBus和Thread的消息处理

从接口看，ActorMsgBus负责消息的发送（Actor通过ActorMsgBus发送消息），Thread负责消息的接收（普通消息转给Actor处理）。

相关实体的协作关系如下

这些对象之间的消息传递关系如下图所示
OneFlow源码阅读7：静态图与运行时_第6张图片

6.4 Runtime激活Actor系统

到这里，Actor之间的协作关系基本清楚了。但是整个Actor系统还处于静止待命的状态。Actor是消息驱动的，总要有消息触发才能让它们转起来。

Runtime在构造函数中，给没有输入依赖的source_tasks发送kStart消息，让Actor系统处于可运行状态。

从DAG的角度看，只要激活source_tasks就行，这些节点给下游发送消息，自然会触发所有Actor的执行。

但这个kStart消息并没有启动计算图的一轮计算。因为这是在Runtime的构造函数中，还处于运行时初始化阶段。发送kStart更像是让Actor系统处于激活状态。我们在讨论完Actor之后，再看看计算图的每一轮计算是怎么触发的。

7 Actor

7.1 Actor的创建

Thread在创建Actor时，会先尝试创建为LightActor，如果不成功，再尝试用预先注册的工厂创建Actor。

有几种TaskType可以用于LightActor：

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

目前大约有20多种Actor的子类型。其它Actor类型根据TaskType预先注册。例如WaitAndSendIdsActor。

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

Actor相关的类关系如下（包含关系只是表示可以访问到相关信息，并不意味着创建或着拥有该类型对象）
OneFlow源码阅读7：静态图与运行时_第7张图片

7.2 Actor的初始化

Actor的构造函数一般都是空的，构建之后需要执行Init函数进行初始化。

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

在Actor::Init中，首先调用ConstructKernel创建kernel实例。和Operator类似，kernel也是以OpTypeCase作为注册的key，例如WaitAndSendIdsKernel。一个Actor通常只有一个kernel。

之后调用NewRegsts创建Regst。Tensor是用户侧的概念。对应的运行时概念是Regst，它持有Kernel需要读写的内存。Regst的概念比Tensor更宽泛，比如框架自动添加的控制Op也会用到Regst。
Actor将自己创建的Regst保存到produced_regsts_。

TakeOverNaiveConsumed只记录需要消费的regst id，但并不push到consumed_regsts_。

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

调用InitBnInOp2BlobInfo会初始化BlobInfo。
之后就是调用VirtualActorInit，这里允许各个Actor子类定制自己的初始化逻辑。通常会调用OF_SET_MSG_HANDLER宏设置Actor的消息处理函数。

7.3 Actor的消息处理

我们以WaitAndSendIds为例，观察一下Actor的消息处理机制。其Actor是WaitAndSendIdsActor，kernel是WaitAndSendIdsKernel。

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

Thread收到的消息如果不是kStopThread或kConstructActor，就调用Actor::ProcessMsg，将消息转给Actor处理。
ProcessMsg函数只是简单的将消息转给handler处理。

WaitAndSendIdsActor::VirtualActorInit中，handler被设置为HandlerWaitToStart。
Runtime的构造函数中，发送的第一批消息是给source_tasks的kStart消息，这个消息就由HandlerWaitToStart函数处理。

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

HandlerNormal中，如果是kCmdMsg，只允许是kStart。通过消息类型校验后，会直接调用ActUntilFail。

7.4 Act执行的条件

ActUntilFail中，Act方法是各个子类自己实现的，一般主要是启动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中的调用流程如下：

注意naive_produced_rs_.PopFrontRegsts会将Regst指针从队列中删掉，相应的计数减1。
而在Actor::HandlerNormal中处理收到的kRegstMsg消息时，如果是consumer发来的ack消息，会调用TryUpdtStateAsProducedRegst，将Regst再添加到 naive_produced_rs_ 中，以保证当前Actor在收到所有ack后是WriteReady的；同时递减total_reading_cnt_。

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

LightActor有自己的消息处理机制，大致原理应该是差不多的。

7.6 Act执行的动作

根据上述讨论，Actor收到kRegstMsg后也会进入ActUntilFail执行。如果读写都是Ready，就执行Act。以WaitAndSendIdsActor为例，主要调用链路如下：

buffer->Pull会等待条件变量的通知。现在，看上去所有Actor都已准备就绪，只等发令枪一响就开跑了。

8 启动静态图的计算

Graph.__run会扣动发令枪的板机，启动计算图的一轮计算。
主要调用流程如下：

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

9 运行时的退出机制

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

运行时的退出始于NNGraph对象的析构。

9.1 Actor的退出

9.2 Thread的退出

10 分布式场景的静态图

分布式的compile_job、物理图Plan和单机场景有明显变化。
比如，每个进程都有一套WaitAndSendIds等控制节点。这也容易理解，因为每个节点都要执行__run和Buffer::Push/Pull，都要启动本进程的Actors执行计算。
matmul和broadcast_add等user op也会在两个节点进行计算。
OneFlow源码阅读7：静态图与运行时_第8张图片

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 参考资料

相关推荐

C#中AutoResetEvent控制线程用法小结

linux中用conda安装大于3.6版本的R

React.js前端导出Excel的方式

atcoder ABC 232 B~E题解

C++实现ETW进行进程变动监控详解

【Unity2D】射线·碰撞投射·方法总结

【面试：并发篇16：多线程：wait/notify详解】原理及错误用法(虚假唤醒等)

android redux 架构,在React Native中使用Redux架构

深入理解数据结构原理(1)—栈(Stack)

实践数据湖iceberg 第七课 实时写入到iceberg

【车牌识别】基于GRNN广义回归神经网络的车牌识别matlab仿真

行业认可 | 悬镜安全荣膺DevSecOps创新赛道领航者，获选年度创新力十强

玩转Linux与运维岗（20）

免费SSL证书申请及部署实践

3.蜂鸣器音乐代码

微信小程序开发 - 视图与逻辑

我在Linux开发板上跑的第一个Qt程序

redis哨兵集群搭建详细过程

C语言进阶（字符串函数及模拟实现strncpy&strncat&strncmp&strstr&strtok&strerror）2

Spark3.x入门到精通-阶段三(深度剖析spark处理数据全流程)

ItVuer - 免责声明 - 关于我们 - 联系我们

本网站信息来源于互联网，如有侵权请联系：561261067@qq.com

桂ICP备16001015号