OneFlow源碼解析：靜態(tài)圖與運行時

來源：CSDN博客 | 2023-01-06 19:07:24 |

作者｜鄭建華更新｜許嘯宇、張文驍、成誠OneFlow靜態(tài)圖的訓(xùn)練效率遠(yuǎn)高于動態(tài)圖（eager模式）。本文試圖通過一個簡單例子，結(jié)合v0.8.0版本的代碼，解讀一下靜態(tài)圖和運行時的實現(xiàn)機制。

在開始之前，建議先讀一下參考資料中《OneFlow框架的系統(tǒng)設(shè)計（https://zhuanlan.zhihu.com/p/337851255）》等系列文章。對靜態(tài)圖、運行時的基本概念和設(shè)計理念有基本的了解，會更容易理解代碼。

(相關(guān)資料圖)

代碼示例

下面的示例代碼來自官方文檔（https://docs.oneflow.org/master/basics/08_nn_graph.html），是一個線性模型的前向計算。后續(xù)主要基于這段代碼進(jìn)行分析。

import oneflow as flowimport oneflow.nn as nnclass 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.biaslinear_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)

oneflow包的初始化

import oneflow在初始化包（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/__init__.py）時，與靜態(tài)圖相關(guān)的主要操作如下：

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）

啟動各個節(jié)點的控制面（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L160-L162）網(wǎng)絡(luò)連接

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

啟動各個節(jié)點的數(shù)據(jù)面網(wǎng)絡(luò)連接（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）?創(chuàng)建 Session，并注冊為 default session（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/session_util.cpp#L89）

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

創(chuàng)建 C++ MultiClientSessionContext（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/multi_client_session.py#L41）?但并不 TryInit

EnvGlobalObjectsScope::Init中先創(chuàng)建一個全局的ProcessCtx（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L132）對象。然后根據(jù)環(huán)境變量等配置，在各個進(jìn)程間創(chuàng)建gRPC和CommNet的連接，分別負(fù)責(zé)控制面和數(shù)據(jù)面的數(shù)據(jù)傳輸。其中在Bootstrap過程中會初始化全局的ProcessCtx（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/rpc/lib/grpc.cpp#L42），給每個進(jìn)程分配一個全局唯一的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））。

本文不涉及網(wǎng)絡(luò)層面的操作，只討論同一進(jìn)程內(nèi)各線程間的交互。

Module類

雖然可以直接用op和tensor構(gòu)造模型，但是op的粒度太細(xì)了，直接用op構(gòu)造模型會比較繁瑣。

Module（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/module.py#L54）是由op和tensor構(gòu)成的、可復(fù)用的子模塊。利用Module可以更高效、更快捷的構(gòu)建復(fù)雜模型。oneflow.nn（https://github.com/Oneflow-Inc/oneflow/blob/d825243aa7aff5cba8bd3a901b4cc56c2b1a36af/python/oneflow/nn/__init__.py）模塊導(dǎo)出了很多預(yù)定義的Module。

Module定義了自己的屬性設(shè)置邏輯（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中

否則按常規(guī)屬性處理

Module可以包含子Module，形成樹結(jié)構(gòu)。因為Module通過setattr將子Module和Parameter都保存到字典結(jié)構(gòu)中，可以方便的遍歷所有Module及其參數(shù)tensor。

Graph類

4.1 構(gòu)造函數(shù)

Graph的構(gòu)造函數(shù)中GetDefaultSession（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L145）得到的session，就是導(dǎo)入oneflow包時NewDefaultSession（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/__init__.py#L229）構(gòu)建的session。當(dāng)時沒有初始化，而是在Graph構(gòu)造時進(jìn)行初始化（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L147）。對應(yīng)的C++函數(shù)是MultiClientSessionContext::TryInit（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/multi_client_session_context.cpp#L67），執(zhí)行時會創(chuàng)建各種全局的資源管理器，比如： ?

LazyJobBuildAndInferCtxMgr

BufferMgr

RegstMgr

ActorMsgBus

ThreadMgr

4.2?__setattr__: 將Module和Tensor封裝為Block

Graph.__setattr__ 支持通過設(shè)置屬性的方式把一個 Module 添加到 Graph 中，之后改 Module 就可以被 Graph 調(diào)用了。添加到 Graph 中的 Module，會被包裝到 Block 里面，Block 起到了代理執(zhí)行的作用，它會給原 Eager 下的 Module 擴展出靜態(tài)執(zhí)行需要的一些特殊功能。

添加到 Graph 中的 Module 和原 Module 共享了狀態(tài)（Parameter、Buffer）和 forward 執(zhí)行邏輯。共享 forward 執(zhí)行邏輯使得靜態(tài)和動態(tài)執(zhí)行計算邏輯相同。共享狀態(tài)則可以使動態(tài)圖下的模型狀態(tài)被靜態(tài)圖復(fù)用?；诖?，兩個 Graph，一個用于訓(xùn)練，一個用于預(yù)測，他們都復(fù)用統(tǒng)一模型 Module，這樣訓(xùn)練和預(yù)測 Graph 也就實現(xiàn)了模型共享。

setattr最重要的動作就是對_add_block的調(diào)用（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L1332），_add_block中主要是調(diào)用get_block_cls并保存結(jié)果（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屬性都轉(zhuǎn)為對應(yīng)的Block對象。為什么要做這個動作呢？主要是靜態(tài)圖編譯需要借助Block類型來實現(xiàn)代理執(zhí)行的功能，這些功能不適合直接寫到 eager 下的 Module 和 Tensor 上。

這個轉(zhuǎn)換是在ModuleBlock構(gòu)造時調(diào)用set_origin（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L131）完成的。對于子Module，會遞歸調(diào)用get_block_cls函數(shù)（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L145），這樣所有子Module及其Tensor屬性都會被轉(zhuǎn)換為對應(yīng)的Block對象。

所以，上述示例代碼中，GraphMyLinear實際存儲的是ModuleBlock，Graph.build執(zhí)行時獲取的model屬性也是ModuleBlock對象，ModuleBlock.origin才是ModuleMyLinear。

Graph.__setattr__不允許將Tensor對象設(shè)置為屬性（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L1340）。Tensor只能存到Module中，因為 Module 是做狀態(tài)共享的基本單位，而 Graph 是不允許復(fù)用的。

4.3 針對不同任務(wù)，定義不同的計算圖

根據(jù)Oneflow Model Zoo的模型示例（https://github.com/Oneflow-Inc/models/blob/1b291f78d8f60e5f04ee0c5962e4611cc4bab40a/Vision/classification/image/alexnet/graph/train.py），train/eval等階段可以創(chuàng)建不同的Graph子類。動態(tài)圖下提供了 Module、Optimizer、Dataloader等模塊，這些模型都可以被添加到 Graph 中。不同的組合可以構(gòu)建不同類型的任務(wù)。

在這些不同階段，Graph構(gòu)造函數(shù)的行為、build函數(shù)的輸入輸出都有各自特點。了解這些，看后續(xù)代碼時會更容易理解各個參數(shù)的具體含義。

構(gòu)造函數(shù)

train階段，需要添加Module、損失函數(shù)、優(yōu)化器和dataloader

eval階段，只需要添加Module和dataloader

build函數(shù)

train

導(dǎo)入樣本和label

調(diào)用Module得到前向計算結(jié)果

計算損失

計算梯度

返回loss

eval

導(dǎo)入樣本和label

調(diào)用Module得到預(yù)估結(jié)果

返回預(yù)估結(jié)果和label

4.4 小結(jié)

上述幾個類型的關(guān)系如下：

下面描述了GraphMyLinear的構(gòu)造流程

* `__init__` * `Graph.__init__` * self.model = linear_model * `Graph.__setattr__` * _add_block * get_block_cls: 遞歸地把Module轉(zhuǎn)為ModuleBlock * `ModuleBlock.__init__` * ModuleBlock.set_origin * `ModuleBlock._origin = origin` (Module) * 對origin的sub modules, parameters, buffers遞歸調(diào)用get_block_cls * `ModuleBlock.__setattr__`

邏輯圖的編譯

計算機語言的編譯，是將高級語言的語句編譯為匯編或機器指令。深度學(xué)習(xí)框架對計算任務(wù)的編譯，是將用戶的特定語句操作轉(zhuǎn)換為DAG圖。oneflow中用Job（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job.proto#L30）描述邏輯的計算圖。

不同于eager模式的動態(tài)圖，靜態(tài)圖在開始執(zhí)行前可以得到整個計算任務(wù)的所有信息，可以對DAG進(jìn)行多輪優(yōu)化。每輪優(yōu)化都是輸入一個Job、得到一個新Job。

最后，根據(jù)分布式環(huán)境配置，將邏輯圖Job轉(zhuǎn)換為物理執(zhí)行的計算圖Plan（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/plan.proto#L34）。在物理圖中，一個op可能分布在多個節(jié)點/進(jìn)程。

啟動DAG計算需要調(diào)用Graph.__call__，這個函數(shù)的執(zhí)行主要分以下幾個步驟：

__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中進(jìn)行。finish_complie_and_init_runtime會繼續(xù)做一些優(yōu)化pass，然后構(gòu)建物理圖、初始化運行時Actor系統(tǒng)。__run會啟動一次DAG的運算。

5.1 graph_build_context: 為邏輯圖編譯設(shè)置基本環(huán)境

在 Graph 中，build 函數(shù)里面的代碼執(zhí)行都在 graph_build_context 的作用域下，這樣實現(xiàn)了動態(tài)轉(zhuǎn)靜態(tài)的功能。

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

首先在context作用域內(nèi)設(shè)置全局的lazy_mode為True（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L46）。在這個context作用域內(nèi)，所有op都由LazyInterpreter解釋執(zhí)行。

其次，在JobBuildAndInferCtx（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L47）作用域內(nèi)，JobBuildAndInferCtx_Open（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L57）調(diào)用類似如下C++代碼

// oneflow/api/python/job_build/job_build_and_infer.h// oneflow/core/job/job_build_and_infer_ctx_mgr.cpp// 如前所述，LazyJobBuildAndInferCtxMgr 在 MultiClientSessionContext::TryInit 執(zhí)行時初始化。// LazyJobBuildAndInferCtxMgr mgr;mgr.OpenJobBuildAndInferCtx(job_name);

OpenJobBuildAndInferCtx會新建一個Job對象（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job_build_and_infer_ctx_mgr.cpp#L32）、一個LazyJobBuildAndInferCtx對象（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job_build_and_infer_ctx_mgr.cpp#L34）。LazyJobBuildAndInferCtx負(fù)責(zé)根據(jù)用戶定制的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-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L854-L856）的作用是，對于用戶傳遞給graph_mylinear(input)的input參數(shù)，針對其中的每個tensor都在邏輯計算圖中插入一個FeedInputOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/system_ops.h#L48）節(jié)點。也就是說，model的輸入（比如樣本tensor，具體參考4.3節(jié)），在靜態(tài)圖中也視為一個op操作。

__build_io內(nèi)會用args（即input）和kwargs構(gòu)造一個ArgsTree。ArgsTree 把 Python 下的輸入、輸出抽象成了一個樹，輸入、輸出可以是嵌套的 Tuple、List、Dict，元素是 Tensor，嵌套的結(jié)構(gòu)剛好可以表示為樹，而 Tensor 是樹中的葉子節(jié)點。示例代碼中kwargs是空的。

遍歷ArgsTree，對args和kwargs的每個tensor都調(diào)用傳入的build_func，對于input來說，就是build_graph_input_arg（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L206）。后面會看到，model的output也會調(diào)用__build_io，所以這個函數(shù)名的意思應(yīng)該就是對model的輸入、輸出進(jìn)行靜態(tài)圖的構(gòu)圖工作。

build_graph_input_arg內(nèi)部會構(gòu)造一個FeedInputOpExpr（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L213），提交給解釋器執(zhí)行。因為是在lazy作用域內(nèi)，由LazyInterpreter解釋執(zhí)行（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L471），LazyInterpreter會將對應(yīng)的op插入靜態(tài)圖。

附：build input時ArgsTree的內(nèi)部結(jié)構(gòu)

__build_io(input)?中 ArgsTree 的內(nèi)部數(shù)據(jù)組織示意

_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（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L471）在執(zhí)行時，GetCurInferCtx()（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L500）返回的就是graph_build_context中OpenJobBuildAndInferCtx（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L57）創(chuàng)建的那個LazyJobBuildAndInferCtx對象，這個對象負(fù)責(zé)邏輯圖的構(gòu)建。添加op的主要調(diào)用流程如下：

infer_ctx->AddAndInferConsistentOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L503）

AddAndInferOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job_build_and_infer_ctx.cpp#L563）

ConstructOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job_build_and_infer_ctx.cpp#L580）

CheckAndConstructOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/operator/operator.cpp#L1216）

NewObj（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/operator/operator.cpp#L51）

OperatorConf中，多種op配置共享op_type字段（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/operator/op_conf.proto#L412），protobuf oneof的op_type_case常量作為注冊NewObj的key。

系統(tǒng)預(yù)定義的op在oneflow/core/operator（https://github.com/Oneflow-Inc/oneflow/tree/release/v0.8.0/oneflow/core/operator）下，例如UserOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/operator/user_op.h#L24）。

AddAndInferOp將返回的Operator保存到LazyJobBuildAndInferCtx的字典中。后續(xù)的函數(shù)調(diào)用，主要是進(jìn)行推導(dǎo)并修改靜態(tài)圖Job，使得各個節(jié)點構(gòu)成一個DAG。

JobBuildAndInferCtx相關(guān)的類關(guān)系如下：

5.2.2 lazy tensor 和 eager tensor 的區(qū)別

LazyInterpreter::ApplyImpl的最后，會調(diào)用BuildTensor（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L518）構(gòu)造一個lazy tensor，作為build_graph_input_arg的返回值（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L216）。所以__build_io返回的lazy_args（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L854）是lazy tensor，它將替代eager的args（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L828）（也就是用戶輸入的input）參與后續(xù)的計算圖構(gòu)建。

那么lazy tensor和eager tensor的區(qū)別是什么呢？eager tensor是要即時計算的，所以需要真實數(shù)據(jù)；而lazy tensor僅在靜態(tài)圖編譯階段用于推導(dǎo)，只需要描述性質(zhì)的元信息。靜態(tài)圖編譯是在lazy模式下運行，只是使用lazy tensor 做計算機構(gòu)圖和校驗。

后面會看到，靜態(tài)圖的運行期已經(jīng)沒有tensor的概念。運行期看到的只是更廣義的Regst存儲，可能代表tensor/blob，也可能是其它控制信息。靜態(tài)圖運行時的輸入，是直接讀取外部 eager tensor的內(nèi)存數(shù)據(jù)到到regst；輸出應(yīng)該是op寫到regst，通過blob構(gòu)造eager tensor。

5.3 build: 將UserOp和FeedVariableOp添加到邏輯圖

__build_graph中的self.build()（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L861）會調(diào)用GraphMyLinear.build()，以及ModuleMyLinear.forward()。因為是在lazy模式下運行，matmul和add都會調(diào)用UserOpExpr重載版本的LazyInterpreter::ApplyImpl（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L832），進(jìn)而調(diào)用AddAndInferConsistentOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L940）進(jìn)行構(gòu)圖操作。

需要說明的是，在引用Module的Parameter屬性時（如weight/bias），會觸發(fā)FeedVariableOp的構(gòu)圖操作（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L226）、調(diào)用對應(yīng)版本的LazyInterpreter::ApplyImpl（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L527）。這個是怎么執(zhí)行的呢？

__build_graph中，在進(jìn)入lazy模式之前，先調(diào)用了_create_states_builder（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L843）。其中self._state()（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L667）返回所有Module的所有Parameter（包括子Module）。

state_block的類型是TensorBlock（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L631）。所有的state_block的lazy_origin_builder().method（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L647）都被設(shè)置為調(diào)用build_graph_state（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L683-L688）。

給build_graph_state（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L220）設(shè)置個斷點能讓整個調(diào)用過程顯形，主要的調(diào)用棧如下：

-> 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"]

這個調(diào)用過程比較容易困擾的是，執(zhí)行對象會在Grpah、GraphMyLinear、ModuleMyLinear、ModuleBlock之間切換。

前面在討論Graph的構(gòu)造時已經(jīng)提過，執(zhí)行self.model(input)時，Graph.__getattr__返回的屬性model是ModuleBlock對象，所以實際調(diào)用的是ModuleBlock.__call__。

在這個函數(shù)內(nèi)調(diào)用__block_forward（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L234），其中的_origin（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L266）是ModuleMyLinear，進(jìn)入到它的forward方法，執(zhí)行到flow.matmul(input, self.weight) + self.bias時，matmul 會被LazyOpInterpreter 所執(zhí)行，在 LazyOpInterpreter 中調(diào)用 AddAndInferConsistentOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L503）

，在 Job 中添加一個 matmul operator。同理后面的加法會在 job 中添加一個 add operator。

self.weight 和 self.bias 會觸發(fā)調(diào)用ModuleBlock.__getattr__，進(jìn)而調(diào)用_get_from_states（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L483），調(diào)用TensorBlock.try_build()（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L521）。這里執(zhí)行的就是進(jìn)入lazy模式之前設(shè)置的build_graph_state（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L220）。從而增加一個FeedVariableOp到計算圖（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L527）。為什么設(shè)置和調(diào)用會距離這么遠(yuǎn)呢？主要是為了讓參數(shù)盡量和消費參數(shù)的 Operator 在一個作用域下，所以實現(xiàn)成了惰性求值來達(dá)到延遲計算的目的。

再后面的步驟就是調(diào)用__build_io（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L869-L875）插入FetchOutputOp（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L589）。也就是說，獲取model的output也是一個op。

到目前為止，前向計算圖就構(gòu)建完成了。它的json表示可以參考附錄。net.op是計算圖的節(jié)點，通過input等屬性可以看出節(jié)點之間的連接關(guān)系。

示例代碼的前向計算圖如下。從這個圖可以看到，input、output、weights等都是op。

5.4 邏輯圖優(yōu)化

在__build_graph中會調(diào)用CurJobBuildAndInferCtx_Complete對靜態(tài)圖進(jìn)行多輪優(yōu)化（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L923），對應(yīng)的C++函數(shù)是LazyJobBuildAndInferCtx::Complete()（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job_build_and_infer_ctx.cpp#L975）。

這之后生成的Job是full_job。本文的示例代碼比較簡單，并不是典型的計算場景，其forwar和ful計算圖的拓?fù)涫且粯拥?。實際大部的圖優(yōu)化都實現(xiàn)在這個階段，如 Op fusion、AMP、ZeRO、常量折疊等等。

到這里，邏輯圖構(gòu)建的主體部分就結(jié)束了。

隨后會構(gòu)建一個CNNGraph對象（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L947），對應(yīng)的C++類型是NNGraph（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.h#L33）。這個對象將負(fù)責(zé)構(gòu)建物理計算圖Plan。它也是整個運行時的擁有者和維護(hù)者。這個對象析構(gòu)時，整個運行時也會有序終止并釋放資源。

5.5 物理圖的編譯

接下來就是執(zhí)行finish_complie_and_init_runtime（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L742），其中的核心調(diào)用是self._c_nn_graph.complie_and_init_runtime()（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/graph.py#L802），對應(yīng)的C++函數(shù)是NNGraph::CompileAndInitRuntime（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L265）。

在這個函數(shù)中，JobCompleter().Complete()（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L280）會繼續(xù)對邏輯圖做幾輪修改優(yōu)化，補全 Runtime 執(zhí)行所需要的附加信息，Compiler().Compile()（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L285）將邏輯圖轉(zhuǎn)為分設(shè)備的物理圖，并繼續(xù)對Plan進(jìn)行修改優(yōu)化。

Plan的編譯是在master節(jié)點進(jìn)行的（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L282）。master節(jié)點會將Plan通過gRPC推送給各個worker節(jié)點（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L308），worker節(jié)點從master拉取物理計算圖（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L310）。

之后調(diào)用NewRuntimeBuffers創(chuàng)建Buffer對象（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L322），Buffer應(yīng)該是主要用于進(jìn)程內(nèi)的信息同步。

然后就準(zhǔn)備初始化運行時了。

示例代碼生成的compiled_job和物理圖Plan的json參見附錄。

最終生成的compiled邏輯圖如下?？蚣茏詣硬迦肓撕芏嘞到y(tǒng)控制節(jié)點。

5.6 Plan的結(jié)構(gòu)

示例代碼輸出的Plan json數(shù)據(jù)見附錄。

Plan在邏輯上和compiled_job是等價的。這里主要關(guān)注task/op之間的關(guān)系。

Plan.task中的每個元素是一個task，其中的exec_sequence.exec_node對應(yīng)job中的op，通常只有一個op（數(shù)組可以支持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體現(xiàn)了task/op之間的連接關(guān)系。

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層面體現(xiàn)了連接關(guān)系。這個map中的key表示輸入輸出，value是register id。

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

task.produced_regst_desc描述了對應(yīng)task生產(chǎn)的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描述了對應(yīng)task消費的register信息

運行時的初始化

NNGraph::CompileAndInitRuntime中，new Runtime這行代碼會初始化運行時（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L331）。主要做的事情包括：

創(chuàng)建Thread

通知Thread創(chuàng)建Actor，Actor會創(chuàng)建Regst和Kernel

給沒有輸入的source_tasks發(fā)送啟動信號kStart

6.1 Runtime創(chuàng)建Thread

在Runtime的構(gòu)造函數(shù)中，DumpThreadIdsFromPlan（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/runtime.cpp#L65）會將Plan中屬于當(dāng)前進(jìn)程的task的thread id存入thread_ids_變量。AddThreads創(chuàng)建這些Thread對象（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/runtime.cpp#L69）。

Thread在構(gòu)造時會創(chuàng)建一個物理線程（?https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/thread/thread.cpp#L39），線程執(zhí)行的是PollMsgChannel方法（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/thread/thread.cpp#L44），Thread就是在這里持續(xù)等待需要處理的新消息。

Thread只處理兩類命令消息：線程終止消息，創(chuàng)建Actor的消息。其它消息交給Actor::ProcessMsg處理（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/thread/thread.cpp#L83）。

6.2 Runtime通知Thread創(chuàng)建Actor

在Runtime的構(gòu)造函數(shù)中，tasks被分為兩類：source_tasks和other_tasks。在示例代碼中，source_tasks（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/runtime.cpp#L84-L85）是沒有輸入邊的task。

從代碼邏輯看，在Plan proto中，task的consumed_regst_desc_id字段是一個map。如果這個map的所有key都是in_ctrl（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/runtime.cpp#L54），這個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調(diào)用HandoutTasks函數(shù)（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/runtime.cpp#L100-L101）會給ActorMsgBus發(fā)送構(gòu)建Actor的kConstructActor消息（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/runtime.cpp#L49）。

6.3 ActorMsgBus和Thread的消息處理

從接口看，ActorMsgBus?（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor_message_bus.cpp#L24）負(fù)責(zé)消息的發(fā)送（Actor通過ActorMsgBus發(fā)送消息），Thread::PollMsgChannel（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/thread/thread.cpp#L60）負(fù)責(zé)消息的接收和處理。

相關(guān)實體的協(xié)作關(guān)系如下

Actor是自調(diào)度的基本單元，接受消息然后工作，工作完后再繼續(xù)發(fā)送消息。

actor_id就是task_id，是在編譯Plan時就確定的。task是編譯時概念，actor是對等的運行時概念。

task_id有特定的編碼格式（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/graph/task_id.cpp#L21-L29），從中可以解析出machine_id（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/graph/task_id.cpp#L73）和thread_id（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/graph/task_id.cpp#L77）。

在跨網(wǎng)絡(luò)的整個物理圖Plan中，actor id相當(dāng)于地址，通過它可以定位唯一的actor實體。

Actor 通過 ActorMsgBus::SendMsg（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor_message_bus.cpp#L24）發(fā)送 ActorMsg（https://github.com/Oneflow-Inc/oneflow/blob/4856d691051accd72f13f4139d281e411977b297/oneflow/core/lazy/actor/actor_message.h#L34）消息。

ActorMsg包含源和目的actor id（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor_message.h#L84-L85）。

如果是進(jìn)程內(nèi)通訊（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor_message_bus.cpp#L26），將通過 ActorMsgBus::SendMsgWithoutCommNet?（https://github.com/Oneflow-Inc/oneflow/blob/4856d691051accd72f13f4139d281e411977b297/oneflow/core/lazy/actor/actor_message_bus.cpp#L49）把 ActorMsg 朝目的 actor 所在的 thread 入隊消息（https://github.com/Oneflow-Inc/oneflow/blob/4856d691051accd72f13f4139d281e411977b297/oneflow/core/thread/thread.h#L40）。

Thread::EnqueueActorMsg 會判斷當(dāng)前 thread 是否是 actor thread，如果是則入本地隊列，否則則入 actor thead 的 channel 隊列。

如果ActorMsg是跨進(jìn)程消息，ActorMsgBus通過CommNet發(fā)送消息（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor_message_bus.cpp#L42-L44），接收方的CommNet應(yīng)該會根據(jù)actor id獲得線程id，從ThreadMgr查到Thread，將消息交給Thread處理。

Thread::PollMsgChannel（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/thread/thread.cpp#L60）負(fù)責(zé)消息的接收和處理。

如果線程本地隊列l(wèi)ocal_msg_queue_為空，則從thread的channel隊列中取出全部ActorMsg放入本地隊列（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/thread/thread.cpp#L63）。

從本地隊列中取出一個ActorMsg，然后開始處理。

處理一些特殊的kCmdMsg消息（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/thread/thread.cpp#L67-L79），然后普通消息交給Actor自行處理（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/thread/thread.cpp#L83）。

Actor收到消息后，會判斷是否滿足了Act的條件，如果滿足，則會執(zhí)行Act，從而調(diào)用LaunchKernel執(zhí)行計算，Act執(zhí)行結(jié)束后通過ActorMsgBus發(fā)消息通知上下游Actor。

這些對象之間的消息傳遞關(guān)系如下圖所示

6.4 激活source Actor

目前的實現(xiàn)中，Actor全部是自調(diào)度的，只能接受來自其他Actor的消息。Actor中有一類比較特殊的source actors，它們與source tasks對應(yīng)。

source actors 沒有上游 actor，它們會朝下游actor發(fā)送消息從而激活所有的Actor運行。

source actors 本身是如何執(zhí)行的呢？它們在接受到 kStart 消息后就會一直 Act 直到進(jìn)入退出流程。但是其 kernel 會阻塞在 Buffer（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/common/buffer.h#L26）處，一直等待其他線程往 buffer 中添加數(shù)據(jù)后，阻塞會被激活，然后 kernel 執(zhí)行讀取，kernel 完成后，actor 的 Act 結(jié)束，往下游發(fā)送消息。

source actors 由于會發(fā)生阻塞，所以其必須有單獨的 actor thread。

Runtime 初始化的的最后一步就是朝各 source actors 發(fā)送 kStart 消息用以激活它們，但 source actors 只有接受到 buffer 的數(shù)據(jù)后才會往下執(zhí)行，然后朝下游 actors 發(fā)送消息，使所有的 actors 都執(zhí)行起來。

Actor

7.1 Actor的創(chuàng)建

Thread在創(chuàng)建Actor時，會先嘗試創(chuàng)建為LightActor（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/thread/thread.cpp#L104），如果不成功，再嘗試用預(yù)先注冊的工廠創(chuàng)建Actor。

有幾種TaskType可以用于LightActor（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/light_actor.cpp#L677-L689）：

kNormalForward，比如matmul、add等user op。

kCopyHd

kTick

kCollectiveBoxingGeneric

目前大約有20多種Actor的子類型。其它Actor類型根據(jù)TaskType（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/task.proto#L8）預(yù)先注冊。例如WaitAndSendIdsActor。

示例代碼的各個節(jié)點對應(yīng)的actor類型參見附錄。

Actor相關(guān)的類關(guān)系如下（包含關(guān)系只是表示可以訪問到相關(guān)信息，并不意味著創(chuàng)建或著擁有該類型對象）

7.2 Actor的初始化

Actor的構(gòu)造函數(shù)一般都是空的，構(gòu)建之后需要執(zhí)行Init（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L129）函數(shù)進(jìn)行初始化。

LightActor繼承自ActorBase，不是Actor的子類，有自己的Init函數(shù)實現(xiàn)。這里只討論Actor的初始化。

在Actor::Init（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L129）中，首先調(diào)用ConstructKernel（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L138）創(chuàng)建kernel實例。和Operator類似，kernel也是以O(shè)pTypeCase作為注冊的key，例如WaitAndSendIdsKernel（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/kernel/wait_and_send_ids_kernel.cpp#L51）。一個Actor通常只有一個kernel。

之后調(diào)用NewRegsts創(chuàng)建Regst（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L152）。Tensor是用戶側(cè)的概念。對應(yīng)的運行時概念是Regst（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/register/register.h#L24），它持有Kernel需要讀寫的內(nèi)存。Regst的概念比Tensor更寬泛，比如框架自動添加的控制Op也會用到Regst。

Actor將自己創(chuàng)建的Regst保存到produced_regsts_（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L153）。

TakeOverNaiveConsumed（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L182）只記錄需要消費的regst id，但并不push到consumed_regsts_。

TakeOverNaiveProduced（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L183）既記錄生產(chǎn)的regst id，也push到naive_produced_rs_（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L249）。這種區(qū)別是為了首次執(zhí)行計算時，actor能順利執(zhí)行。后面分析Actor的消息處理時會再回過頭來討論一下。

調(diào)用InitBnInOp2BlobInfo會初始化BlobInfo（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L184）。

之后就是調(diào)用VirtualActorInit（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.cpp#L185），這里允許各個Actor子類定制自己的初始化邏輯。通常會調(diào)用OF_SET_MSG_HANDLER宏（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/lazy/actor/actor.h#L76-L80）設(shè)置Actor的消息處理函數(shù)。

7.3 Actor的消息處理

LightActor 首先會根據(jù)消息類型分別處理 kRegstMsg 和 kEordMsg 消息。HandleRegstMsg（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/lazy/actor/light_actor.cpp#L424）中根據(jù) RegstMsg 的 type (kProduced 或 kComsumed) 來分別處理各種讀寫狀態(tài)計數(shù)。

然后判斷讀寫計數(shù)是否達(dá)到了判斷條件，如果達(dá)到了意味著滿足了讀寫 regst 的條件，然后就執(zhí)行 ActOnce（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/lazy/actor/light_actor.cpp#L451）。

LightActor::ActOnce 會在第一次執(zhí)行時去 InitBnInOp2Blob 和 InitActMsg。InitBnInOp2Blob 初始化 resgt 中的 bn 與 Blob 的映射關(guān)系，為 kernel 提供通過 bn 訪問 Blob 的功能。InitActMsg 會初始化好所有需要發(fā)送的消息避免后繼發(fā)消息時重復(fù)的構(gòu)建消息。

然后就是 LaunchKernel，接著會 ResetState 重置 regst 狀態(tài)。

LaunchKernel 后就會把之前構(gòu)建好的消息發(fā)送出去，同步消息會直接入隊 thread 消息隊列，異步消息通過 callback 發(fā)送到 ActorMsgBus。

普通 Actor::ProcessMsg 會調(diào)用 msg handler 來處理消息，最常見的 msg handler 就是 Actor::HandlerNormal（https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/lazy/actor/actor.cpp#L329）。

Actor::HandlerNormal 中流程跟 LightActor 中類似，會根據(jù)不同的 regst 類型來分別處理，Actor 中對 regst 的狀態(tài)管理方式與 LightActor 不同，LightActor 中的方式更加高效，Actor 中能處理一些特殊情況。

消息處理完畢后，就會調(diào)用 ActUntilFail，ActUntilFail 會判斷 IsReadReady 和 IsWriteReady 來決定是否可以進(jìn)行 Act。

最常見的 NaiveActor::Act() 就是執(zhí)行 AsyncLaunchKernel。

Act 完成后，就開始朝上下游發(fā)送 regst 消息。

還有一些特殊的 Actor，我們以WaitAndSendIdsActor為例，觀察一下這類Actor的消息處理機制。

之所以選擇這個例子，一是這個Actor比較簡單；二是這是一個典型的source task，想看一下計算圖是怎么被觸發(fā)啟動計算的。

Thread收到的消息如果不是kStopThread或kConstructActor，就調(diào)用Actor::ProcessMsg（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/thread/thread.cpp#L83），將消息轉(zhuǎn)給Actor處理。

ProcessMsg函數(shù)只是簡單的將消息轉(zhuǎn)給handler處理（https://github.com/Oneflow-Inc/oneflow/blob/b6bf3f8843679111eb1edf79deefce814d250f4e/oneflow/core/lazy/actor/actor.h#L38）。

WaitAndSendIdsActor::VirtualActorInit中，handler被設(shè)置為HandlerWaitToStart（https://github.com/Oneflow-Inc/oneflow/blob/22f70a1719f371a54512633bb92086580d9c3c89/oneflow/core/lazy/actor/wait_and_send_ids_actor.cpp#L53）。

Runtime的構(gòu)造函數(shù)中，發(fā)送的第一批消息是給source_tasks的kStart消息，這個消息就由HandlerWaitToStart函數(shù)處理。

HandlerWaitToStart校驗消息類型后，將handler設(shè)置為HandlerNormal（https://github.com/Oneflow-Inc/oneflow/blob/b17a9cd6b930b5817c63623fb682bd708377a93b/oneflow/core/job/runtime.cpp#L109）（這也是大部分Actor的默認(rèn)handler），然后調(diào)用ProcessMsg（https://github.com/Oneflow-Inc/oneflow/blob/22f70a1719f371a54512633bb92086580d9c3c89/oneflow/core/lazy/actor/wait_and_send_ids_actor.cpp#L74），實際就是調(diào)用新設(shè)置的handler HandlerNormal。

HandlerNormal中，如果是kCmdMsg，只允許是kStart（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L377）。通過消息類型校驗后，會直接調(diào)用ActUntilFail（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L378）。

7.4 Act執(zhí)行的條件

LightActor 和 Actor 判斷能否進(jìn)行 Act 采用了不同的策略，LightActor 的效率更高，Actor 能處理一些特殊情況。

對于 LightActor，當(dāng)在讀的register計數(shù) total_reading_cnt_ 歸 0，可消費的register計數(shù) ready_consumed_ 增加到 max_ready_consumed_，前者表示所有的消費者已經(jīng)讀取當(dāng)前 LightActor 的 Regst，后者表示當(dāng)前 LightActor 消費的所有 Regst 已經(jīng)到達(dá)（由上游發(fā)送的 Regst 消息）。

對于 Actor，Actor::ActUntilFail中，Act方法（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L424）是各個子類自己實現(xiàn)的，一般主要是啟動kernel計算。

但是在執(zhí)行Act之前，需要先確認(rèn)：

Act執(zhí)行依賴的數(shù)據(jù)是否都已經(jīng)就緒？（IsReadReady）

Act生產(chǎn)出來的數(shù)據(jù)，消費方是否已經(jīng)用完、并收到ack消息確認(rèn)？（IsWriteReady）

Actor有4個與此相關(guān)的成員變量

RegstSlot naive_produced_rs_;

RegstSlot inplace_produced_rs_;

RegstSlot naive_consumed_rs_;

RegstSlot inplace_consumed_rs_;

xx_produced_rs_存儲的是當(dāng)前Actor的下游consumer返回的、已經(jīng)使用完畢的ack regst信息。（當(dāng)前Actor生產(chǎn)的Regst存儲在produced_regsts_中。）

運行時在初始化的過程中，所有Actor都沒有運行過，任何Actor都不可能收到ack消息，所以在Actor初始化時，要預(yù)先填充xx_produced_rs_，這樣才能保證Actor在首次運行前是WriteReady的，才能順利啟動執(zhí)行。

xx_consumed_rs_存儲的是上游依賴發(fā)來的數(shù)據(jù)。它不需要預(yù)先填充。因為source_tasks沒有輸入依賴，自然就是ReadReady的；而xx_produced_rs_在初始化時的預(yù)先填充又保證它是WriteReady的，所以source_tasks可以直接運行。source_tasks的輸出消息發(fā)給下游，下游也會變?yōu)镽eadReady，而下游在初始化后也保證是WriteReady的。整個Actor系統(tǒng)就可以這樣運轉(zhuǎn)起來了。

7.5 Actor上下游之間的通知機制

Act執(zhí)行完畢后，需要將結(jié)果數(shù)據(jù)發(fā)給下游consumer。以 WaitAndSendIds 的 Naive Produced 為例，ActUntilFail中的調(diào)用流程如下：

AsyncSendNaiveProducedRegstMsgToConsumer（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L427）

VirtualAsyncSendNaiveProducedRegstMsgToConsumer（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L441）

HandleProducedNaiveDataRegstToConsumer（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L446）

HandleRegstToConsumer（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L577）

EnqueueAsyncMsg（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L523）

如果目標(biāo)線程是當(dāng)前線程，ActorMsgBus::SendMsg（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L662）

否則，將消息加入async_msg_queue_（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L664）

增加 total_reading_cnt_（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L526）（這個變量表示已經(jīng)發(fā)消息給下游、但未收到的ack數(shù)量）

naive_produced_rs_.PopFrontRegsts（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L581）

AsyncSendProducedCtrlRegstMsgToConsumer

注意naive_produced_rs_.PopFrontRegsts（https://github.com/Oneflow-Inc/oneflow/blob/06a6af1c7f760ba4b12d2dfb8f73d7fda5c7dbab/oneflow/core/lazy/actor/register_slot.cpp#L53）會將Regst指針從隊列中刪掉，相應(yīng)的可用（https://github.com/Oneflow-Inc/oneflow/blob/06a6af1c7f760ba4b12d2dfb8f73d7fda5c7dbab/oneflow/core/lazy/actor/register_slot.cpp#L49）register計數(shù)減1（https://github.com/Oneflow-Inc/oneflow/blob/06a6af1c7f760ba4b12d2dfb8f73d7fda5c7dbab/oneflow/core/lazy/actor/register_slot.cpp#L49）。

而在Actor::HandlerNormal中處理收到的kRegstMsg消息（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L340）時，如果是consumer發(fā)來的ack消息，會調(diào)用TryUpdtStateAsProducedRegst（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L355），將Regst再添加到 naive_produced_rs_ 中（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L654），以保證當(dāng)前Actor在收到所有ack后是WriteReady的；同時遞減在讀的 register 計數(shù)total_reading_cnt_。

Actor對依賴的上游消息的處理是類似的。通過以下函數(shù)調(diào)用給上游發(fā)送ack消息、通知 register 已經(jīng)用完，可以繼續(xù)更新了：

AsyncSendNaiveConsumedRegstMsgToProducer（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L431）

AsyncRetInplaceConsumedRegstIfNoConsumer（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L432）在Actor::HandlerNormal中收到kRegstMsg消息后，將消息添加到consumed_rs_（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L344），以保證當(dāng)前Actor在收到所有依賴數(shù)據(jù)后是ReadReady的。

LightActor有自己的消息處理機制（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/light_actor.cpp#L299），大致原理應(yīng)該是差不多的。

7.6 Act執(zhí)行的動作

根據(jù)上述討論，Actor收到kRegstMsg后也會進(jìn)入ActUntilFail執(zhí)行。如果讀寫都是Ready，就執(zhí)行Act（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L424）。以WaitAndSendIdsActor為例，主要調(diào)用鏈路如下：

AsyncLaunchKernel（https://github.com/Oneflow-Inc/oneflow/blob/22f70a1719f371a54512633bb92086580d9c3c89/oneflow/core/lazy/actor/wait_and_send_ids_actor.cpp#L58）

ek.kernel->Launch（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L562），啟動Kernel計算

Forward（https://github.com/Oneflow-Inc/oneflow/blob/eae9ff38f074479d79ce24b0f6e0594f82126171/oneflow/core/kernel/kernel.cpp#L52）

ForwardDataContent（https://github.com/Oneflow-Inc/oneflow/blob/eae9ff38f074479d79ce24b0f6e0594f82126171/oneflow/core/kernel/kernel.cpp#L65）

buffer->Pull（https://github.com/Oneflow-Inc/oneflow/blob/b17a9cd6b930b5817c63623fb682bd708377a93b/oneflow/core/kernel/wait_and_send_ids_kernel.cpp#L40）

給regst的存儲地址mut_dptr賦值（https://github.com/Oneflow-Inc/oneflow/blob/b17a9cd6b930b5817c63623fb682bd708377a93b/oneflow/core/kernel/wait_and_send_ids_kernel.cpp#L47）

buffer->Pull會等待條件變量的通知（https://github.com/Oneflow-Inc/oneflow/blob/49f60e682518436dfeb37344a15902a959e0e4f2/oneflow/core/common/buffer.h#L60）。現(xiàn)在，看上去所有Actor都已準(zhǔn)備就緒，只等發(fā)令槍一響就開跑了。

啟動靜態(tài)圖的計算

Graph.__run（https://github.com/Oneflow-Inc/oneflow/blob/81edd938826a7ea903174d682348847658b64653/python/oneflow/nn/graph/graph.py#L226）會扣動發(fā)令槍的板機，啟動計算圖的一輪計算。

主要調(diào)用流程如下：

RunLazyNNGraph（https://github.com/Oneflow-Inc/oneflow/blob/81edd938826a7ea903174d682348847658b64653/python/oneflow/nn/graph/graph.py#L1076）

builder->LaunchLazyJob（https://github.com/Oneflow-Inc/oneflow/blob/8f672eea116cae4a73bb7309e7496b08d7ec9a32/oneflow/core/framework/nn_graph.cpp#L568）

LaunchLazyJobInstructionType（https://github.com/Oneflow-Inc/oneflow/blob/8f672eea116cae4a73bb7309e7496b08d7ec9a32/oneflow/core/framework/instructions_builder.cpp#L179）

Buffer::Push（https://github.com/Oneflow-Inc/oneflow/blob/8f672eea116cae4a73bb7309e7496b08d7ec9a32/oneflow/core/framework/instructions_builder.cpp#L179）

這里的Buffer::Push就是WaitAndSendIdsKernel在等待的起跑信號。

運行時的退出機制

整個運行時包含很多對象和資源，安全有序的退出是龐雜而又細(xì)致的工作。這里僅以WaitAndSendIds為例，從一個側(cè)面觀察一下運行時的退出機制。

運行時的退出始于NNGraph對象的析構(gòu)（https://github.com/Oneflow-Inc/oneflow/blob/8f672eea116cae4a73bb7309e7496b08d7ec9a32/oneflow/core/framework/nn_graph.cpp#L76）。

9.1 Actor的退出

NNGraph在析構(gòu)時，會關(guān)閉所有的Buffer對象（https://github.com/Oneflow-Inc/oneflow/blob/8f672eea116cae4a73bb7309e7496b08d7ec9a32/oneflow/core/framework/nn_graph.cpp#L82）。

Buffer在關(guān)閉時，會設(shè)置is_closed_ = true并通知所有監(jiān)聽者（https://github.com/Oneflow-Inc/oneflow/blob/49f60e682518436dfeb37344a15902a959e0e4f2/oneflow/core/common/buffer.h#L81）。但是Pull會繼續(xù)處理完已經(jīng)提交的計算。

所以，Buffer應(yīng)該是主要用于進(jìn)程內(nèi)的通信和異步協(xié)調(diào)的一個類。

WaitAndSendIdsKernel這時候正在等待新一輪計算開始（https://github.com/Oneflow-Inc/oneflow/blob/b17a9cd6b930b5817c63623fb682bd708377a93b/oneflow/core/kernel/wait_and_send_ids_kernel.cpp#L40），結(jié)果收到Pull返回的kBufferStatusErrorClosed（https://github.com/Oneflow-Inc/oneflow/blob/49f60e682518436dfeb37344a15902a959e0e4f2/oneflow/core/common/buffer.h#L61）。

WaitAndSendIdsActor::IsCustomizedReadReady以后就一直返回false（https://github.com/Oneflow-Inc/oneflow/blob/22f70a1719f371a54512633bb92086580d9c3c89/oneflow/core/lazy/actor/wait_and_send_ids_actor.cpp#L68），IsReadReady也返回false（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L533）。

這之后，ActUntilFail只會執(zhí)行異步消息發(fā)送（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L437）（不再進(jìn)入while循環(huán)）

WaitAndSendIdsActor::HandlerNormal仍然會處理其它Actor發(fā)來的消息（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L340）。但因為IsCustomizedReadReady返回false，會進(jìn)入AsyncSendEORDMsgForAllProducedRegstDesc（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L394）執(zhí)行。它會給每個下游發(fā)送kEordMsg消息（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L614）。

Actor在收到上游發(fā)來的kEordMsg消息后，遞減remaining_eord_cnt_（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L331）。

remaining_eord_cnt_被初始化為Actor的輸入regst的數(shù)量（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L171）。

total_reading_cnt_是當(dāng)前Actor生產(chǎn)的、已經(jīng)發(fā)給consumer、但尚未收到ack的消息數(shù)量。

Actor目前仍可以正常接收consumer發(fā)來的ack消息。

當(dāng)上述2個變量都為0時（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L395），意味著所有上游都發(fā)出了kEordMsg消息，也收到了所有下游的ack消息。Actor就給Thread返回1（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L397）。

如果上述兩個變量有不為0的，就修改handler，由HandlerZombie（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L399）處理后續(xù)收到的消息。

Thread收到Actor返回的1后（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/thread/thread.cpp#L84），將它從自己的存儲中刪除（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/thread/thread.cpp#L89），并遞減運行Actor的數(shù)量。

9.2 Thread的退出

NNGraph重置runtime_導(dǎo)致運行時對象被析構(gòu)（https://github.com/Oneflow-Inc/oneflow/blob/8f672eea116cae4a73bb7309e7496b08d7ec9a32/oneflow/core/framework/nn_graph.cpp#L83）。

Runtime刪除所有Thread（https://github.com/Oneflow-Inc/oneflow/blob/b17a9cd6b930b5817c63623fb682bd708377a93b/oneflow/core/job/runtime.cpp#L117）。

ThreadMgr給所有Thread發(fā)送kStopThread消息（https://github.com/Oneflow-Inc/oneflow/blob/c8c6d351fa28c5ebce948d69c06670a783f83f74/oneflow/core/thread/thread_manager.cpp#L64）。同時，重置指針導(dǎo)致Thread析構(gòu)（https://github.com/Oneflow-Inc/oneflow/blob/c8c6d351fa28c5ebce948d69c06670a783f83f74/oneflow/core/thread/thread_manager.cpp#L66）。

Thread的物理線程退出PollMsgChannel循環(huán)（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/thread/thread.cpp#L68）。

Thread等待物理線程結(jié)束，關(guān)閉channel（https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/thread/thread.cpp#L52）。

10?

分布式場景的靜態(tài)圖

分布式的compile_job、物理圖Plan和單機場景有明顯變化。

比如，每個進(jìn)程都有一套WaitAndSendIds等控制節(jié)點。這也容易理解，因為每個節(jié)點都要執(zhí)行__run和Buffer::Push/Pull，都要啟動本進(jìn)程的Actors執(zhí)行計算。

matmul和broadcast_add等user op也會在兩個節(jié)點進(jìn)行計算。 ?

10.1 示例代碼

啟動方式參考Global Tensor的官方文檔。

import oneflow as flowimport oneflow.nn as nnP0 = 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.biaslinear_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.pybreak test.py:25break oneflow/nn/graph/graph.py:221break oneflow/nn/graph/graph.py:741break oneflow/nn/graph/graph.py:745break oneflow/nn/graph/graph.py:759break oneflow/nn/graph/graph.py:828break oneflow/nn/graph/graph.py:777break oneflow/nn/graph/graph.py:1066break oneflow/nn/graph/graph.py:1133break oneflow/framework/graph_build_util.py:227

11.1.2 C++斷點示例

啟動命令

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

斷點示例

set breakpoint pending onbreak oneflow::ActorMsg::BuildEordMsgbreak oneflow/core/common/buffer.h:80break oneflow::(anonymous namespace)::CheckAndConstructOpbreak oneflow::WaitAndSendIdsActor::Actbreak oneflow::WaitAndSendIdsActor::HandlerWaitToStartbreak oneflow/core/lazy/actor/light_actor.cpp:452break oneflow/core/lazy/actor/light_actor.cpp:485break oneflow::ForeignInputKernel::ForwardDataContentbreak oneflow::vm::LaunchLazyJobInstructionType::Compute

11.2 靜態(tài)圖的json表示

forward（https://quip.com/OMc4A0HOOr0C）

full（https://quip.com/JLaMAHGBLXmK）

compiled（https://quip.com/tXjuAiS3J0Ab）

plan（https://quip.com/a0DMAAIte6PQ）

11.3 actor type

naive_actor

System-AutoTick-AppendDeviceTick_9System-AutoTick-DstSubsetTick_12System-AutoTick-DstSubsetTick_21System-AutoTick-DstSubsetTick_27System-AutoTick-Prepend-DeviceTick_7System-AutoTick-SrcSubsetTick_20System-AutoTick-SrcSubsetTick_26System-AutoTick-SrcSubsetTick_8System-AutoTick-Tick_11System-AutoTick-Tick_13System-EagerCriticalSection-Callback-23System-EagerCriticalSection-Callback-29System-EagerCriticalSection-Interface-Begin-Tick-18System-EagerCriticalSection-Interface-Begin-Tick-24System-EagerCriticalSection-Interface-End-Tick-19System-EagerCriticalSection-Interface-End-Tick-25System-EagerCriticalSection-Wait-22System-EagerCriticalSection-Wait-28

light_actor

_GraphMyLinear_0_input.0.0_2_GraphMyLinear_0_output.0.0_2model.biasmodel-broadcast_add-1model-matmul-0model.weightSystem-AutoTick-SinkTick_15System-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-Inc/oneflow/tree/release/v0.8.0）

OneFlow框架的系統(tǒng)設(shè)計（上篇）（https://zhuanlan.zhihu.com/p/337851255）

OneFlow框架的系統(tǒng)設(shè)計（中篇）（https://zhuanlan.zhihu.com/p/338699487）

OneFlow框架的系統(tǒng)設(shè)計（下篇）（https://zhuanlan.zhihu.com/p/339208452）

一個Job在OneFlow中的執(zhí)行過程—上篇（https://zhuanlan.zhihu.com/p/344531540）

一個Job在OneFlow中的執(zhí)行過程—中篇（https://zhuanlan.zhihu.com/p/355654002）

一個Job在OneFlow中的執(zhí)行過程—下篇（https://zhuanlan.zhihu.com/p/363689736）

靜態(tài)圖模塊 nn.Graph（https://docs.oneflow.org/master/basics/08_nn_graph.html）

OneFlow系統(tǒng)設(shè)計（https://docs.oneflow.org/v0.4.0/basics_topics/essentials_of_oneflow.html）

torch.nn.Module（https://pytorch.org/docs/1.10/generated/torch.nn.Module.html）

其他人都在看

OneFlow源碼解析：自動微分機制

ChatGPT的一小步，NLP范式轉(zhuǎn)變的一大步

李白：你的模型權(quán)重很不錯，可惜被我沒收了

OpenAI掌門Sam Altman：AI下一個發(fā)展階段

32篇年度最佳AI論文;Python編譯器Codon開源

比快更快，開源Stable Diffusion刷新作圖速度

OneEmbedding:單卡訓(xùn)練TB級推薦模型不是夢

歡迎Star、試用OneFlow最新版本：GitHub - Oneflow-Inc/oneflow: OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient.OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient. - GitHub - Oneflow-Inc/oneflow: OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient.https://github.com/Oneflow-Inc/oneflow/

關(guān)鍵詞：

OneFlow源碼解析：靜態(tài)圖與運行時

本周資訊推薦

2020上半年全球企業(yè)區(qū)塊鏈發(fā)明專利排行榜：阿里巴巴以1457件位列第一

數(shù)據(jù)：區(qū)塊鏈相關(guān)企業(yè)目前共有4.56萬家在業(yè)存續(xù)的超過4萬家

2020年區(qū)塊鏈領(lǐng)域全球授權(quán)專利報告：支付寶以212件專利數(shù)位列全球第一

米粉節(jié)神秘大獎：小米3萬元全屋智能產(chǎn)品僅需1元

考拉海購宣布升級商品全鏈路溯源系統(tǒng) 引入?yún)^(qū)塊鏈溯源技術(shù)

熱點資訊

ios7如何越獄 ios7越獄步驟越獄前準(zhǔn)備工作

剪卡教程怎么剪手機卡?

自拍神器是什么？自拍神器怎么選擇？選購方法教給你

劍靈小神龍證書有什么用劍靈小神龍證書怎么用?-環(huán)球百事通

網(wǎng)絡(luò)受限制或無連接怎么辦應(yīng)該怎么結(jié)局呢？

OneFlow源碼解析：靜態(tài)圖與運行時

圖片推薦

要聞

本周資訊推薦

熱點資訊