作者|鄭建華 更新|許嘯宇、張文驍、成誠OneFlow靜態(tài)圖的訓(xùn)練效率遠(yuǎn)高于動(dòng)態(tài)圖(eager模式)。本文試圖通過一個(gè)簡單例子,結(jié)合v0.8.0版本的代碼,解讀一下靜態(tài)圖和運(yùn)行時(shí)的實(shí)現(xiàn)機(jī)制。
在開始之前,建議先讀一下參考資料中《OneFlow框架的系統(tǒng)設(shè)計(jì)(https://zhuanlan.zhihu.com/p/337851255)》等系列文章。對靜態(tài)圖、運(yùn)行時(shí)的基本概念和設(shè)計(jì)理念有基本的了解,會(huì)更容易理解代碼。
(相關(guān)資料圖)
1?
代碼示例
下面的示例代碼來自官方文檔(https://docs.oneflow.org/master/basics/08_nn_graph.html),是一個(gè)線性模型的前向計(jì)算。后續(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)
2?
oneflow包的初始化
import oneflow在初始化包(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/__init__.py)時(shí),與靜態(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)
啟動(dòng)各個(gè)節(jié)點(diǎn)的控制面(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)
啟動(dòng)各個(gè)節(jié)點(diǎn)的數(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)建一個(gè)全局的ProcessCtx(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/env_global_objects_scope.cpp#L132)對象。然后根據(jù)環(huán)境變量等配置,在各個(gè)進(jìn)程間創(chuàng)建gRPC和CommNet的連接,分別負(fù)責(zé)控制面和數(shù)據(jù)面的數(shù)據(jù)傳輸。其中在Bootstrap過程中會(huì)初始化全局的ProcessCtx(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/rpc/lib/grpc.cpp#L42),給每個(gè)進(jìn)程分配一個(gè)全局唯一的rank編號(hào)(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)各線程間的交互。
3?
Module類
雖然可以直接用op和tensor構(gòu)造模型,但是op的粒度太細(xì)了,直接用op構(gòu)造模型會(huì)比較繁瑣。
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)。因?yàn)镸odule通過setattr將子Module和Parameter都保存到字典結(jié)構(gòu)中,可以方便的遍歷所有Module及其參數(shù)tensor。
4?
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包時(shí)NewDefaultSession(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/__init__.py#L229)構(gòu)建的session。當(dāng)時(shí)沒有初始化,而是在Graph構(gòu)造時(shí)進(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í)行時(shí)會(huì)創(chuàng)建各種全局的資源管理器,比如: ?
LazyJobBuildAndInferCtxMgr
BufferMgr
RegstMgr
ActorMsgBus
ThreadMgr
4.2?__setattr__: 將Module和Tensor封裝為Block
Graph.__setattr__ 支持通過設(shè)置屬性的方式把一個(gè) Module 添加到 Graph 中,之后改 Module 就可以被 Graph 調(diào)用了。添加到 Graph 中的 Module,會(huì)被包裝到 Block 里面,Block 起到了代理執(zhí)行的作用,它會(huì)給原 Eager 下的 Module 擴(kuò)展出靜態(tài)執(zhí)行需要的一些特殊功能。
添加到 Graph 中的 Module 和原 Module 共享了狀態(tài)(Parameter、Buffer)和 forward 執(zhí)行邏輯。共享 forward 執(zhí)行邏輯使得靜態(tài)和動(dòng)態(tài)執(zhí)行計(jì)算邏輯相同。共享狀態(tài)則可以使動(dòng)態(tài)圖下的模型狀態(tài)被靜態(tài)圖復(fù)用。基于此,兩個(gè) Graph,一個(gè)用于訓(xùn)練,一個(gè)用于預(yù)測,他們都復(fù)用統(tǒng)一模型 Module,這樣訓(xùn)練和預(yù)測 Graph 也就實(shí)現(xiàn)了模型共享。
setattr最重要的動(dòng)作就是對_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對象。為什么要做這個(gè)動(dòng)作呢?主要是靜態(tài)圖編譯需要借助Block類型來實(shí)現(xiàn)代理執(zhí)行的功能,這些功能不適合直接寫到 eager 下的 Module 和 Tensor 上。
這個(gè)轉(zhuǎn)換是在ModuleBlock構(gòu)造時(shí)調(diào)用set_origin(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/nn/graph/block.py#L131)完成的。對于子Module,會(huì)遞歸調(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屬性都會(huì)被轉(zhuǎn)換為對應(yīng)的Block對象。
所以,上述示例代碼中,GraphMyLinear實(shí)際存儲(chǔ)的是ModuleBlock,Graph.build執(zhí)行時(shí)獲取的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中,因?yàn)?Module 是做狀態(tài)共享的基本單位,而 Graph 是不允許復(fù)用的。
4.3 針對不同任務(wù),定義不同的計(jì)算圖
根據(jù)Oneflow Model Zoo的模型示例(https://github.com/Oneflow-Inc/models/blob/1b291f78d8f60e5f04ee0c5962e4611cc4bab40a/Vision/classification/image/alexnet/graph/train.py),train/eval等階段可以創(chuàng)建不同的Graph子類。動(dòng)態(tài)圖下提供了 Module、Optimizer、Dataloader等模塊,這些模型都可以被添加到 Graph 中。不同的組合可以構(gòu)建不同類型的任務(wù)。
在這些不同階段,Graph構(gòu)造函數(shù)的行為、build函數(shù)的輸入輸出都有各自特點(diǎn)。了解這些,看后續(xù)代碼時(shí)會(huì)更容易理解各個(gè)參數(shù)的具體含義。
構(gòu)造函數(shù)
train階段,需要添加Module、損失函數(shù)、優(yōu)化器和dataloader
eval階段,只需要添加Module和dataloader
build函數(shù)
train
導(dǎo)入樣本和label
調(diào)用Module得到前向計(jì)算結(jié)果
計(jì)算損失
計(jì)算梯度
返回loss
eval
導(dǎo)入樣本和label
調(diào)用Module得到預(yù)估結(jié)果
返回預(yù)估結(jié)果和label
4.4 小結(jié)
上述幾個(gè)類型的關(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__`
5?
邏輯圖的編譯
計(jì)算機(jī)語言的編譯,是將高級(jí)語言的語句編譯為匯編或機(jī)器指令。深度學(xué)習(xí)框架對計(jì)算任務(wù)的編譯,是將用戶的特定語句操作轉(zhuǎn)換為DAG圖。oneflow中用Job(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job.proto#L30)描述邏輯的計(jì)算圖。
不同于eager模式的動(dòng)態(tài)圖,靜態(tài)圖在開始執(zhí)行前可以得到整個(gè)計(jì)算任務(wù)的所有信息,可以對DAG進(jìn)行多輪優(yōu)化。每輪優(yōu)化都是輸入一個(gè)Job、得到一個(gè)新Job。
最后,根據(jù)分布式環(huán)境配置,將邏輯圖Job轉(zhuǎn)換為物理執(zhí)行的計(jì)算圖Plan(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/plan.proto#L34)。在物理圖中,一個(gè)op可能分布在多個(gè)節(jié)點(diǎn)/進(jìn)程。
啟動(dòng)DAG計(jì)算需要調(diào)用Graph.__call__,這個(gè)函數(shù)的執(zhí)行主要分以下幾個(gè)步驟:
__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會(huì)繼續(xù)做一些優(yōu)化pass,然后構(gòu)建物理圖、初始化運(yùn)行時(shí)Actor系統(tǒng)。__run會(huì)啟動(dòng)一次DAG的運(yùn)算。
5.1 graph_build_context: 為邏輯圖編譯設(shè)置基本環(huán)境
在 Graph 中,build 函數(shù)里面的代碼執(zhí)行都在 graph_build_context 的作用域下,這樣實(shí)現(xiàn)了動(dòng)態(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)。在這個(gè)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í)行時(shí)初始化。// LazyJobBuildAndInferCtxMgr mgr;mgr.OpenJobBuildAndInferCtx(job_name);
OpenJobBuildAndInferCtx會(huì)新建一個(gè)Job對象(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/job_build_and_infer_ctx_mgr.cpp#L32)、一個(gè)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:為計(jì)算圖添加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ù),針對其中的每個(gè)tensor都在邏輯計(jì)算圖中插入一個(gè)FeedInputOp(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/system_ops.h#L48)節(jié)點(diǎn)。也就是說,model的輸入(比如樣本tensor,具體參考4.3節(jié)),在靜態(tài)圖中也視為一個(gè)op操作。
__build_io內(nèi)會(huì)用args(即input)和kwargs構(gòu)造一個(gè)ArgsTree。ArgsTree 把 Python 下的輸入、輸出抽象成了一個(gè)樹,輸入、輸出可以是嵌套的 Tuple、List、Dict,元素是 Tensor,嵌套的結(jié)構(gòu)剛好可以表示為樹,而 Tensor 是樹中的葉子節(jié)點(diǎn)。示例代碼中kwargs是空的。
遍歷ArgsTree,對args和kwargs的每個(gè)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)。后面會(huì)看到,model的output也會(huì)調(diào)用__build_io,所以這個(gè)函數(shù)名的意思應(yīng)該就是對model的輸入、輸出進(jìn)行靜態(tài)圖的構(gòu)圖工作。
build_graph_input_arg內(nèi)部會(huì)構(gòu)造一個(gè)FeedInputOpExpr(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/python/oneflow/framework/graph_build_util.py#L213),提交給解釋器執(zhí)行。因?yàn)槭窃趌azy作用域內(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會(huì)將對應(yīng)的op插入靜態(tài)圖。
附:build input時(shí)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í)行時(shí),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)建的那個(gè)LazyJobBuildAndInferCtx對象,這個(gè)對象負(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,使得各個(gè)節(jié)點(diǎn)構(gòu)成一個(gè)DAG。
JobBuildAndInferCtx相關(guān)的類關(guān)系如下:
5.2.2 lazy tensor 和 eager tensor 的區(qū)別
LazyInterpreter::ApplyImpl的最后,會(huì)調(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)造一個(gè)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ù)的計(jì)算圖構(gòu)建。
那么lazy tensor和eager tensor的區(qū)別是什么呢?eager tensor是要即時(shí)計(jì)算的,所以需要真實(shí)數(shù)據(jù);而lazy tensor僅在靜態(tài)圖編譯階段用于推導(dǎo),只需要描述性質(zhì)的元信息。靜態(tài)圖編譯是在lazy模式下運(yùn)行,只是使用lazy tensor 做計(jì)算機(jī)構(gòu)圖和校驗(yàn)。
后面會(huì)看到,靜態(tài)圖的運(yùn)行期已經(jīng)沒有tensor的概念。運(yùn)行期看到的只是更廣義的Regst存儲(chǔ),可能代表tensor/blob,也可能是其它控制信息。靜態(tài)圖運(yùn)行時(shí)的輸入,是直接讀取外部 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)會(huì)調(diào)用GraphMyLinear.build(),以及ModuleMyLinear.forward()。因?yàn)槭窃趌azy模式下運(yùn)行,matmul和add都會(huì)調(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屬性時(shí)(如weight/bias),會(huì)觸發(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)。這個(gè)是怎么執(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è)置個(gè)斷點(diǎn)能讓整個(gè)調(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"]
這個(gè)調(diào)用過程比較容易困擾的是,執(zhí)行對象會(huì)在Grpah、GraphMyLinear、ModuleMyLinear、ModuleBlock之間切換。
前面在討論Graph的構(gòu)造時(shí)已經(jīng)提過,執(zhí)行self.model(input)時(shí),Graph.__getattr__返回的屬性model是ModuleBlock對象,所以實(shí)際調(diào)用的是ModuleBlock.__call__。
在這個(gè)函數(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時(shí),matmul 會(huì)被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 中添加一個(gè) matmul operator。同理后面的加法會(huì)在 job 中添加一個(gè) add operator。
self.weight 和 self.bias 會(huì)觸發(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)。從而增加一個(gè)FeedVariableOp到計(jì)算圖(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/op_interpreter/lazy_op_interpreter.cpp#L527)。為什么設(shè)置和調(diào)用會(huì)距離這么遠(yuǎn)呢?主要是為了讓參數(shù)盡量和消費(fèi)參數(shù)的 Operator 在一個(gè)作用域下,所以實(shí)現(xiàn)成了惰性求值來達(dá)到延遲計(jì)算的目的。
再后面的步驟就是調(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也是一個(gè)op。
到目前為止,前向計(jì)算圖就構(gòu)建完成了。它的json表示可以參考附錄。net.op是計(jì)算圖的節(jié)點(diǎn),通過input等屬性可以看出節(jié)點(diǎn)之間的連接關(guān)系。
示例代碼的前向計(jì)算圖如下。從這個(gè)圖可以看到,input、output、weights等都是op。
5.4 邏輯圖優(yōu)化
在__build_graph中會(huì)調(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。本文的示例代碼比較簡單,并不是典型的計(jì)算場景,其forwar和ful計(jì)算圖的拓?fù)涫且粯拥?。?shí)際大部的圖優(yōu)化都實(shí)現(xiàn)在這個(gè)階段,如 Op fusion、AMP、ZeRO、常量折疊等等。
到這里,邏輯圖構(gòu)建的主體部分就結(jié)束了。
隨后會(huì)構(gòu)建一個(gè)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)。這個(gè)對象將負(fù)責(zé)構(gòu)建物理計(jì)算圖Plan。它也是整個(gè)運(yùn)行時(shí)的擁有者和維護(hù)者。這個(gè)對象析構(gòu)時(shí),整個(gè)運(yùn)行時(shí)也會(huì)有序終止并釋放資源。
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)。
在這個(gè)函數(shù)中,JobCompleter().Complete()(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L280)會(huì)繼續(xù)對邏輯圖做幾輪修改優(yōu)化,補(bǔ)全 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é)點(diǎn)進(jìn)行的(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L282)。master節(jié)點(diǎn)會(huì)將Plan通過gRPC推送給各個(gè)worker節(jié)點(diǎn)(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/framework/nn_graph.cpp#L308),worker節(jié)點(diǎn)從master拉取物理計(jì)算圖(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)備初始化運(yùn)行時(shí)了。
示例代碼生成的compiled_job和物理圖Plan的json參見附錄。
最終生成的compiled邏輯圖如下??蚣茏詣?dòng)插入了很多系統(tǒng)控制節(jié)點(diǎn)。
5.6 Plan的結(jié)構(gòu)
示例代碼輸出的Plan json數(shù)據(jù)見附錄。
Plan在邏輯上和compiled_job是等價(jià)的。這里主要關(guān)注task/op之間的關(guān)系。
Plan.task中的每個(gè)元素是一個(gè)task,其中的exec_sequence.exec_node對應(yīng)job中的op,通常只有一個(gè)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)系。這個(gè)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是消費(fèi)者,
produced_regst_desc.out.regst_desc_type.data_regst_desc.lbi2blob_desc.lbi就是這個(gè)register的logic blob id。
task.consumed_regst_desc_id描述了對應(yīng)task消費(fèi)的register信息
6?
運(yùn)行時(shí)的初始化
NNGraph::CompileAndInitRuntime中,new Runtime這行代碼會(huì)初始化運(yùn)行時(shí)(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會(huì)創(chuàng)建Regst和Kernel
給沒有輸入的source_tasks發(fā)送啟動(dòng)信號(hào)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)會(huì)將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)造時(shí)會(huì)創(chuàng)建一個(gè)物理線程(?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字段是一個(gè)map。如果這個(gè)map的所有key都是in_ctrl(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/job/runtime.cpp#L54),這個(gè)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)會(huì)給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)實(shí)體的協(xié)作關(guān)系如下
Actor是自調(diào)度的基本單元,接受消息然后工作,工作完后再繼續(xù)發(fā)送消息。
actor_id就是task_id,是在編譯Plan時(shí)就確定的。task是編譯時(shí)概念,actor是對等的運(yùn)行時(shí)概念。
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ò)的整個(gè)物理圖Plan中,actor id相當(dāng)于地址,通過它可以定位唯一的actor實(shí)體。
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 入隊(duì)消息(https://github.com/Oneflow-Inc/oneflow/blob/4856d691051accd72f13f4139d281e411977b297/oneflow/core/thread/thread.h#L40)。
Thread::EnqueueActorMsg 會(huì)判斷當(dāng)前 thread 是否是 actor thread,如果是則入本地隊(duì)列,否則則入 actor thead 的 channel 隊(duì)列。
如果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)該會(huì)根據(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é)消息的接收和處理。
如果線程本地隊(duì)列l(wèi)ocal_msg_queue_為空,則從thread的channel隊(duì)列中取出全部ActorMsg放入本地隊(duì)列(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/thread/thread.cpp#L63)。
從本地隊(duì)列中取出一個(gè)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收到消息后,會(huì)判斷是否滿足了Act的條件,如果滿足,則會(huì)執(zhí)行Act,從而調(diào)用LaunchKernel執(zhí)行計(jì)算,Act執(zhí)行結(jié)束后通過ActorMsgBus發(fā)消息通知上下游Actor。
這些對象之間的消息傳遞關(guān)系如下圖所示
6.4 激活source Actor
目前的實(shí)現(xiàn)中,Actor全部是自調(diào)度的,只能接受來自其他Actor的消息。Actor中有一類比較特殊的source actors,它們與source tasks對應(yīng)。
source actors 沒有上游 actor,它們會(huì)朝下游actor發(fā)送消息從而激活所有的Actor運(yùn)行。
source actors 本身是如何執(zhí)行的呢?它們在接受到 kStart 消息后就會(huì)一直 Act 直到進(jìn)入退出流程。但是其 kernel 會(huì)阻塞在 Buffer(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.1/oneflow/core/common/buffer.h#L26) 處,一直等待其他線程往 buffer 中添加數(shù)據(jù)后,阻塞會(huì)被激活,然后 kernel 執(zhí)行讀取,kernel 完成后,actor 的 Act 結(jié)束,往下游發(fā)送消息。
source actors 由于會(huì)發(fā)生阻塞,所以其必須有單獨(dú)的 actor thread。
Runtime 初始化的的最后一步就是朝各 source actors 發(fā)送 kStart 消息用以激活它們,但 source actors 只有接受到 buffer 的數(shù)據(jù)后才會(huì)往下執(zhí)行,然后朝下游 actors 發(fā)送消息,使所有的 actors 都執(zhí)行起來。
7?
Actor
7.1 Actor的創(chuàng)建
Thread在創(chuàng)建Actor時(shí),會(huì)先嘗試創(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。
示例代碼的各個(gè)節(jié)點(diǎn)對應(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ù)實(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實(shí)例。和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)。一個(gè)Actor通常只有一個(gè)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)的運(yùn)行時(shí)概念是Regst(https://github.com/Oneflow-Inc/oneflow/blob/release/v0.8.0/oneflow/core/register/register.h#L24),它持有Kernel需要讀寫的內(nèi)存。Regst的概念比Tensor更寬泛,比如框架自動(dòng)添加的控制Op也會(huì)用到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)只記錄需要消費(fèi)的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í)行計(jì)算時(shí),actor能順利執(zhí)行。后面分析Actor的消息處理時(shí)會(huì)再回過頭來討論一下。
調(diào)用InitBnInOp2BlobInfo會(huì)初始化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),這里允許各個(gè)Actor子類定制自己的初始化邏輯。通常會(huì)調(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 首先會(huì)根據(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)計(jì)數(shù)。
然后判斷讀寫計(jì)數(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 會(huì)在第一次執(zhí)行時(shí)去 InitBnInOp2Blob 和 InitActMsg。InitBnInOp2Blob 初始化 resgt 中的 bn 與 Blob 的映射關(guān)系,為 kernel 提供通過 bn 訪問 Blob 的功能。InitActMsg 會(huì)初始化好所有需要發(fā)送的消息避免后繼發(fā)消息時(shí)重復(fù)的構(gòu)建消息。
然后就是 LaunchKernel,接著會(huì) ResetState 重置 regst 狀態(tài)。
LaunchKernel 后就會(huì)把之前構(gòu)建好的消息發(fā)送出去,同步消息會(huì)直接入隊(duì) thread 消息隊(duì)列,異步消息通過 callback 發(fā)送到 ActorMsgBus。
普通 Actor::ProcessMsg 會(huì)調(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 中類似,會(huì)根據(jù)不同的 regst 類型來分別處理,Actor 中對 regst 的狀態(tài)管理方式與 LightActor 不同,LightActor 中的方式更加高效,Actor 中能處理一些特殊情況。
消息處理完畢后,就會(huì)調(diào)用 ActUntilFail,ActUntilFail 會(huì)判斷 IsReadReady 和 IsWriteReady 來決定是否可以進(jìn)行 Act。
最常見的 NaiveActor::Act() 就是執(zhí)行 AsyncLaunchKernel。
Act 完成后,就開始朝上下游發(fā)送 regst 消息。
還有一些特殊的 Actor,我們以WaitAndSendIdsActor為例,觀察一下這類Actor的消息處理機(jī)制。
之所以選擇這個(gè)例子,一是這個(gè)Actor比較簡單;二是這是一個(gè)典型的source task,想看一下計(jì)算圖是怎么被觸發(fā)啟動(dòng)計(jì)算的。
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消息,這個(gè)消息就由HandlerWaitToStart函數(shù)處理。
HandlerWaitToStart校驗(yàn)消息類型后,將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),實(shí)際就是調(diào)用新設(shè)置的handler HandlerNormal。
HandlerNormal中,如果是kCmdMsg,只允許是kStart(https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L377)。通過消息類型校驗(yàn)后,會(huì)直接調(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計(jì)數(shù) total_reading_cnt_ 歸 0,可消費(fèi)的register計(jì)數(shù) ready_consumed_ 增加到 max_ready_consumed_,前者表示所有的消費(fèi)者已經(jīng)讀取當(dāng)前 LightActor 的 Regst,后者表示當(dāng)前 LightActor 消費(fèi)的所有 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)是各個(gè)子類自己實(shí)現(xiàn)的,一般主要是啟動(dòng)kernel計(jì)算。
但是在執(zhí)行Act之前,需要先確認(rèn):
Act執(zhí)行依賴的數(shù)據(jù)是否都已經(jīng)就緒?(IsReadReady)
Act生產(chǎn)出來的數(shù)據(jù),消費(fèi)方是否已經(jīng)用完、并收到ack消息確認(rèn)?(IsWriteReady)
Actor有4個(gè)與此相關(guān)的成員變量
RegstSlot naive_produced_rs_;
RegstSlot inplace_produced_rs_;
RegstSlot naive_consumed_rs_;
RegstSlot inplace_consumed_rs_;
xx_produced_rs_存儲(chǔ)的是當(dāng)前Actor的下游consumer返回的、已經(jīng)使用完畢的ack regst信息。(當(dāng)前Actor生產(chǎn)的Regst存儲(chǔ)在produced_regsts_中。)
運(yùn)行時(shí)在初始化的過程中,所有Actor都沒有運(yùn)行過,任何Actor都不可能收到ack消息,所以在Actor初始化時(shí),要預(yù)先填充xx_produced_rs_,這樣才能保證Actor在首次運(yùn)行前是WriteReady的,才能順利啟動(dòng)執(zhí)行。
xx_consumed_rs_存儲(chǔ)的是上游依賴發(fā)來的數(shù)據(jù)。它不需要預(yù)先填充。因?yàn)閟ource_tasks沒有輸入依賴,自然就是ReadReady的;而xx_produced_rs_在初始化時(shí)的預(yù)先填充又保證它是WriteReady的,所以source_tasks可以直接運(yùn)行。source_tasks的輸出消息發(fā)給下游,下游也會(huì)變?yōu)镽eadReady,而下游在初始化后也保證是WriteReady的。整個(gè)Actor系統(tǒng)就可以這樣運(yùn)轉(zhuǎn)起來了。
7.5 Actor上下游之間的通知機(jī)制
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)(這個(gè)變量表示已經(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)會(huì)將Regst指針從隊(duì)列中刪掉,相應(yīng)的可用(https://github.com/Oneflow-Inc/oneflow/blob/06a6af1c7f760ba4b12d2dfb8f73d7fda5c7dbab/oneflow/core/lazy/actor/register_slot.cpp#L49)register計(jì)數(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)時(shí),如果是consumer發(fā)來的ack消息,會(huì)調(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的;同時(shí)遞減在讀的 register 計(jì)數(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有自己的消息處理機(jī)制(https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/light_actor.cpp#L299),大致原理應(yīng)該是差不多的。
7.6 Act執(zhí)行的動(dòng)作
根據(jù)上述討論,Actor收到kRegstMsg后也會(huì)進(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),啟動(dòng)Kernel計(jì)算
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的存儲(chǔ)地址mut_dptr賦值(https://github.com/Oneflow-Inc/oneflow/blob/b17a9cd6b930b5817c63623fb682bd708377a93b/oneflow/core/kernel/wait_and_send_ids_kernel.cpp#L47)
buffer->Pull會(huì)等待條件變量的通知(https://github.com/Oneflow-Inc/oneflow/blob/49f60e682518436dfeb37344a15902a959e0e4f2/oneflow/core/common/buffer.h#L60)?,F(xiàn)在,看上去所有Actor都已準(zhǔn)備就緒,只等發(fā)令槍一響就開跑了。
8?
啟動(dòng)靜態(tài)圖的計(jì)算
Graph.__run(https://github.com/Oneflow-Inc/oneflow/blob/81edd938826a7ea903174d682348847658b64653/python/oneflow/nn/graph/graph.py#L226)會(huì)扣動(dòng)發(fā)令槍的板機(jī),啟動(dòng)計(jì)算圖的一輪計(jì)算。
主要調(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在等待的起跑信號(hào)。
9?
運(yùn)行時(shí)的退出機(jī)制
整個(gè)運(yùn)行時(shí)包含很多對象和資源,安全有序的退出是龐雜而又細(xì)致的工作。這里僅以WaitAndSendIds為例,從一個(gè)側(cè)面觀察一下運(yùn)行時(shí)的退出機(jī)制。
運(yùn)行時(shí)的退出始于NNGraph對象的析構(gòu)(https://github.com/Oneflow-Inc/oneflow/blob/8f672eea116cae4a73bb7309e7496b08d7ec9a32/oneflow/core/framework/nn_graph.cpp#L76)。
9.1 Actor的退出
NNGraph在析構(gòu)時(shí),會(huì)關(guān)閉所有的Buffer對象(https://github.com/Oneflow-Inc/oneflow/blob/8f672eea116cae4a73bb7309e7496b08d7ec9a32/oneflow/core/framework/nn_graph.cpp#L82)。
Buffer在關(guān)閉時(shí),會(huì)設(shè)置is_closed_ = true并通知所有監(jiān)聽者(https://github.com/Oneflow-Inc/oneflow/blob/49f60e682518436dfeb37344a15902a959e0e4f2/oneflow/core/common/buffer.h#L81)。但是Pull會(huì)繼續(xù)處理完已經(jīng)提交的計(jì)算。
所以,Buffer應(yīng)該是主要用于進(jìn)程內(nèi)的通信和異步協(xié)調(diào)的一個(gè)類。
WaitAndSendIdsKernel這時(shí)候正在等待新一輪計(jì)算開始(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只會(huì)執(zhí)行異步消息發(fā)送(https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L437)(不再進(jìn)入while循環(huán))
WaitAndSendIdsActor::HandlerNormal仍然會(huì)處理其它Actor發(fā)來的消息(https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L340)。但因?yàn)镮sCustomizedReadReady返回false,會(huì)進(jìn)入AsyncSendEORDMsgForAllProducedRegstDesc(https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/lazy/actor/actor.cpp#L394)執(zhí)行。它會(huì)給每個(gè)下游發(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個(gè)變量都為0時(shí)(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)。
如果上述兩個(gè)變量有不為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),將它從自己的存儲(chǔ)中刪除(https://github.com/Oneflow-Inc/oneflow/blob/55b822e4d3c88757d11077d7546981309125c73f/oneflow/core/thread/thread.cpp#L89),并遞減運(yùn)行Actor的數(shù)量。
9.2 Thread的退出
NNGraph重置runtime_導(dǎo)致運(yùn)行時(shí)對象被析構(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)。同時(shí),重置指針導(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ī)場景有明顯變化。
比如,每個(gè)進(jìn)程都有一套WaitAndSendIds等控制節(jié)點(diǎn)。這也容易理解,因?yàn)槊總€(gè)節(jié)點(diǎn)都要執(zhí)行__run和Buffer::Push/Pull,都要啟動(dòng)本進(jìn)程的Actors執(zhí)行計(jì)算。
matmul和broadcast_add等user op也會(huì)在兩個(gè)節(jié)點(diǎn)進(jìn)行計(jì)算。 ?
10.1 示例代碼
啟動(dòng)方式參考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 斷點(diǎn)
11.1.1 Python斷點(diǎn)示例
# 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++斷點(diǎn)示例
啟動(dòng)命令
source /mnt/oneflow/build/source.shgdb --args python3 /mnt/oneflow/test.py# set breakpoints# run
斷點(diǎn)示例
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è)計(jì)(上篇)(https://zhuanlan.zhihu.com/p/337851255)
OneFlow框架的系統(tǒng)設(shè)計(jì)(中篇)(https://zhuanlan.zhihu.com/p/338699487)
OneFlow框架的系統(tǒng)設(shè)計(jì)(下篇)(https://zhuanlan.zhihu.com/p/339208452)
一個(gè)Job在OneFlow中的執(zhí)行過程—上篇(https://zhuanlan.zhihu.com/p/344531540)
一個(gè)Job在OneFlow中的執(zhí)行過程—中篇(https://zhuanlan.zhihu.com/p/355654002)
一個(gè)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è)計(jì)(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源碼解析:自動(dòng)微分機(jī)制
ChatGPT的一小步,NLP范式轉(zhuǎn)變的一大步
李白:你的模型權(quán)重很不錯(cuò),可惜被我沒收了
OpenAI掌門Sam Altman:AI下一個(gè)發(fā)展階段
32篇年度最佳AI論文;Python編譯器Codon開源
比快更快,開源Stable Diffusion刷新作圖速度
OneEmbedding:單卡訓(xùn)練TB級(jí)推薦模型不是夢
關(guān)鍵詞: