胖虎的工具箱-编程导航 胖虎的工具箱-编程导航

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

1年前 (2022) 程序员胖胖胖虎阿
229 0 0

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

  • GetEnv

    • EnvGlobalObjectsScope::Init

      • 启动各个节点的控制面网络连接
      • 初始化VM
      • 启动各个节点的数据面网络连接
      • 初始化KernelObserver

  • NewDefaultSession

    • RegsiterSession 创建 Session，并注册为 default session

    • 创建 Python MultiClientSession 并保存到dict，但并不 TryInit

      • 创建 C++ MultiClientSessionContext 但并不 TryInit

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：静态图与运行时

下面描述了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__，这个函数的执行主要分以下几个步骤：

  • __call__

    • _compile if not _is_compiled

      • build_graph

        • __build_graph
      • finish_complie_and_init_runtime
    • __run

逻辑图编译主要在__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：静态图与运行时

__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的主要调用流程如下：

  • infer_ctx->AddAndInferConsistentOp
  • AddAndInferOp
  • ConstructOp
  • CheckAndConstructOp
  • NewObj

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

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

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

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：静态图与运行时

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.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处理）。

相关实体的协作关系如下

  • Actor是调度的基本单元。

    • actor_id就是task_id，是在编译Plan时就确定的。task是编译器概念，actor是对等的运行时概念。
    • task_id有特定的编码格式，从中可以解析出machine_id和thread_id。
    • 在跨网络的整个物理图Plan中，actor id相当于地址，通过它可以定位唯一的actor实体。

  • Actor之间通过ActorMsgBus::SendMsg进行消息通讯。

    • ActorMsg包含源和目的actor id。
    • 如果是进程内通讯，ActorMsgBus将消息发给Thread，Thread转给Actor处理消息。
    • 如果是跨进程消息，ActorMsgBus通过CommNet发送消息，接收方的CommNet应该会根据actor id获得线程id，从ThreadMgr查到Thread，将消息交给Thread处理。

  • Thread通过EnqueueActorMsg接收发给自己管理的Actor的消息。

    • 如果是本线程内的actor之间的消息，直接将消息放到local队列。否则放到Channel的异步消息队列。
    • Thread::PollMsgChannel会在构造Thread时启动的物理线程中接收Channel的消息，普通消息转给Actor处理。
  • Actor收到消息后，负责LaunchKernel执行计算，计算结束后通过ActorMsgBus通知上下游Actor。

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

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.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中的调用流程如下：

  • AsyncSendNaiveProducedRegstMsgToConsumer

    • VirtualAsyncSendNaiveProducedRegstMsgToConsumer

      • HandleProducedNaiveDataRegstToConsumer

        • HandleRegstToConsumer

          • EnqueueAsyncMsg

            • 如果目标线程是当前线程，ActorMsgBus::SendMsg
            • 否则，将消息加入async_msg_queue_
          • 增加 total_reading_cnt_（这个变量表示已经发消息给下游、但未收到的ack数量）
        • naive_produced_rs_.PopFrontRegsts
    • AsyncSendProducedCtrlRegstMsgToConsumer

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

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

  • AsyncSendNaiveConsumedRegstMsgToProducer
  • AsyncRetInplaceConsumedRegstIfNoConsumer
    在Actor::HandlerNormal中收到kRegstMsg消息后，将消息添加到consumed_rs_，以保证当前Actor在收到所有依赖数据后是ReadReady的。

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

7.6 Act执行的动作

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

  • AsyncLaunchKernel
  • ek.kernel->Launch，启动Kernel计算
  • Forward
  • ForwardDataContent
  • buffer->Pull
  • 给regst的存储地址mut_dptr赋值

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

8 启动静态图的计算

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

  • RunLazyNNGraph
  • builder->LaunchLazyJob
  • LaunchLazyJobInstructionType
  • Buffer::Push

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

9 运行时的退出机制

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

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

9.1 Actor的退出

  • NNGraph在析构时，会关闭所有的Buffer对象。

  • Buffer在关闭时，会设置is_closed_ = true并通知所有监听者。但是Pull会继续处理完已经提交的计算。

    • 所以，Buffer应该是主要用于进程内的通信和异步协调的一个类。
  • WaitAndSendIdsKernel这时候正在等待新一轮计算开始，结果收到Pull返回的kBufferStatusErrorClosed。

  • WaitAndSendIdsActor::IsCustomizedReadReady以后就一直返回false，IsReadReady也返回false。

    • 这之后，ActUntilFail只会执行异步消息发送（不再进入while循环）
  • WaitAndSendIdsActor::HandlerNormal仍然会处理其它Actor发来的消息。但因为IsCustomizedReadReady返回false，会进入AsyncSendEORDMsgForAllProducedRegstDesc执行。它会给每个下游发送kEordMsg消息。

  • Actor在收到上游发来的kEordMsg消息后，递减remaining_eord_cnt_。

    • remaining_eord_cnt_被初始化为Actor的输入regst的数量。

  • total_reading_cnt_是当前Actor生产的、已经发给consumer、但尚未收到ack的消息数量。

    • Actor目前仍可以正常接收consumer发来的ack消息。

  • 当上述2个变量都为0时，意味着所有上游都发出了kEordMsg消息，也收到了所有下游的ack消息。Actor就给Thread返回1。

    • 如果上述两个变量有不为0的，就修改handler，由HandlerZombie处理后续收到的消息。
  • Thread收到Actor返回的1后，将它从自己的存储中删除，并递减运行Actor的数量。

9.2 Thread的退出

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

10 分布式场景的静态图

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

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表示

  • forward
  • full
  • compiled
  • plan

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
  • OneFlow框架的系统设计（上篇）
  • OneFlow框架的系统设计（中篇）
  • OneFlow框架的系统设计（下篇）
  • 一个Job在OneFlow中的执行过程—上篇
  • 一个Job在OneFlow中的执行过程—中篇
  • 一个Job在OneFlow中的执行过程—下篇
  • 静态图模块 nn.Graph
  • OneFlow系统设计
  • torch.nn.Module
版权声明：程序员胖胖胖虎阿 发表于 2022年11月25日 上午6:00。
转载请注明：OneFlow源码阅读7：静态图与运行时 | 胖虎的工具箱-编程导航

相关文章

使用最新版flink tidb cdc 同步数据到StarRocks
程序员胖胖胖虎阿
270
一文读懂 Spring Boot 配置文件 application.properties ！
程序员胖胖胖虎阿
145
IntelliJ IDEA正版 破解教程(Intellij IDEA2022永久激活码每日更新)
程序员胖胖胖虎阿
228
Spring Boot 深度调优，JVM 调优（深度好文）
程序员胖胖胖虎阿
200
五年北京，这个改变我命运的城市，终于要离开了
程序员胖胖胖虎阿
186
c++ value category
程序员胖胖胖虎阿
106

暂无评论

暂无评论...
致力于称为最好的程序员导航网站

友链申请 免责声明 关于我们

Copyright © 2022  胖虎的工具箱-编程导航   陇ICP备2022001249号