深度解讀：讓你掌握OneFlow框架的系統(tǒng)設計(下篇)

本文主要介紹OneFlow系統(tǒng)的運行時（Runtime）的運行流程，以及參與運行時的各個模塊是如何協(xié)同工作的，還探討了OneFlow的Actor機制如何解決流水線和流控問題（Control Flow）。

如果你對OneFlow這套致簡致快的框架設計感興趣，或者對深度學習框架、分布式系統(tǒng)感興趣的話，本文就會讓你全面掌握OneFlow的系統(tǒng)設計。相信讀完這篇文章，你就會理解我們是如何看待分布式深度學習訓練的，我們?yōu)槭裁匆@樣設計，這樣設計的好處是什么，以及我們?yōu)槭裁聪嘈臤neFlow這套設計是分布式深度學習訓練框架的最優(yōu)設計。

目 錄

1. 深度學習框架原理

2. OneFlow系統(tǒng)架構設計（簡略版）

3. OneFlow完整運行流程與各模塊的交互方式

3.1 分布式集群環(huán)境初始化

3.2 Python端搭建計算圖

3.3 編譯期：OneFlow(JobSet) -> MergedPlan

3.4 編譯期：Compiler(Job)->Plan

3.5 運行時：Runtime(Plan)

本篇內容為下篇，上篇請見本次推送第一條，中篇請見本次推送第二條。

3.5. 運行時：Runtime（Plan）

OneFlow的運行時（Runtime）極其簡單，總共分3步：

1. 創(chuàng)建所有的Global對象

2. 根據(jù)Plan創(chuàng)建本機上的所有Actor

3. 給源節(jié)點的Actor發(fā)送ActorCmd::kStart啟動信號。隨后整個計算圖中的Actor依次啟動，runtime啟動完畢。

由于OneFlow的運行時僅是一張全部由Actor組成的計算圖（對于分布式訓練，是一張跨機器的Actor計算圖），當每個機器都把本機上的Actor建立起來以后，僅需要給那些源節(jié)點Actor發(fā)動啟動信號，啟動信號就會在整個計算圖中傳導開來，每個Actor就開始根據(jù)自身的狀態(tài)機工作了。當整個訓練（Runtime）要結束時，也是這些源節(jié)點發(fā)送關閉信號（kEordMsg），關閉信號也會隨著Actor之間的通信逐漸傳導到整個計算圖，所有的Actor就會根據(jù)eord信號依次關閉。當所有的Actor都關閉后，Runtime就可以下線了。

對運行時的Actor機制介紹可以參考知乎文章：《都2020年了，為什么我們相信OneFlow會成功》 中的章節(jié)三：OneFlow的特色一：Actor機制——用一套簡潔的機制解決所有分布式深度學習框架中的技術難題。

3.5.1 創(chuàng)建所有的Global對象

此處會依次創(chuàng)建運行時所需的所有全局對象。

• CommNet：CommNet是OneFlow分布式訓練中負責多機數(shù)據(jù)傳輸和消息通信的模塊。底層有基于Epoll的實現(xiàn)和基于RDMA的實現(xiàn)。

• boxing::collective::CollectiveBoxingExecutor & boxing::collective::CollectiveBoxingDeviceCtxPoller：負責執(zhí)行集合通信操作（NCCL）

• MemoryAllocator：負責內存（Host內存 和 GPU顯存）的申請與釋放

• RegstMgr：負責創(chuàng)建所有的Regst （Mgr是Manager的縮寫）

• ActorMsgBus：負責運行時Actor之間的消息通信 （Msg是Message的縮寫）

• ThreadMgr：負責創(chuàng)建和管理所有的Thread

3.5.2 ThreadMgr與Thread

在創(chuàng)建Global對象ThreadMgr時，ThreadMgr會根據(jù)Plan中本機上的所有TaskProto中的ThreadID創(chuàng)建對應的Thread。

Thread

Thread負責創(chuàng)建、運行、銷毀Actor。一個Thread會管理多個Actor，Actor收到的消息（ActorMsg）都需要通過Thread中的消息隊列獲取。

Thread的消息隊列分為兩級：

• msg_channel_（繼承自Channel對象）接收跨線程的ActorMsg

• local_msg_queue_ （就是一個隊列std::queue）接收本線程內的消息通信

通過local_msg_queue_可以加速消息傳遞的過程。

每個Thread內部都有一個輪詢線程actor_thread_負責輪詢消息隊列PollMsgChannel，將輪詢到的消息解析，調用該消息的接收者Actor，并讓該Actor處理該消息ProcessMsg。

GpuThread

Thread分為CPU和GPU Thread。CpuThread除了啟動輪詢線程以外沒有其他多余的工作了。GpuThread有兩個額外的部分：1）創(chuàng)建ThreadCtx，里面包含了GPU的CUDA stream handle和CUDA callback event channel；2）啟動一個額外的輪詢線程callback event poller，負責從ThreadCtx中的callback event channel中輪詢獲取callback，并執(zhí)行該callback（原因是GPU上的任務是異步執(zhí)行）。對GPU的架構和使用，我們放在下面的Device章節(jié)介紹。

3.5.3 ActorMsgBus與ActorMsg

ActorMsgBus

每臺機器都會有一個Global對象ActorMsgBus負責消息通信。只有一個主要的接口：SendMsg

ActorMsgBus相當于一個消息的路由，會判斷該消息的目的地是否是本機，如果是本機，則通過ThreadMgr找到對應的Thread，然后EnqueueActorMsg。如果消息的目的地是其他機器，則通過Global對象CommNet將該消息發(fā)送給其他機器。其他機器的Global對象再收到這個消息以后會通知本機的ActorMsgBus做消息處理。這樣就完成了一個消息從消息的生產(chǎn)者Actor到消費者Actor的傳遞。

運行時Actor消息通信機制

示意圖見下圖：

當一個Actor需要給另一個Actor發(fā)消息時，會判斷接收者Actor：

• 是否是本線程內：

• 如果是，則ActorMsgBus會找到本機內的對應線程Thread，傳入到該Thread的Msg channel中

• 否則：調用本機器的CommNet對象傳輸該消息。接收者所在機器的CommNet對象收到消息后會轉給該機器的ActorMsgBus處理。該機器的ActorMsgBus會找到對應的線程Thread將該消息傳入線程的MsgChannel中

• 如果是，則直接壓入Thread的LocalMsgQueue中 （最快）

• 否則：調用本機器的ActorMsgBus傳輸數(shù)據(jù)。Actor Msg Bus會判斷接收者是否在本機內

• Thread會不斷輪詢自己的LocalMsgQueue，取出對應的消息找到對應的Actor去處理該消息。如果LocalMsgQueue為空，則嘗試去從MsgChannel中取消息放到LocalMsgQueue中。

ActorMsg

運行時Actor的消息稱之為ActorMsg。ActorMsg有幾種類型（ActorMsgType）：

• kRegstMsg：表示這個ActorMsg包含了一個Regst。這是運行時Actor之間通信最主要的消息，生產(chǎn)者生產(chǎn)一個Regst通知下游消費者的消息、消費者使用完Regst返還給生產(chǎn)者說我用完了，都是RegstMsg?？梢詮腁ctorMsg的regst()接口中拿到該Regst。

  需要注意的是，無論是生產(chǎn)者通知消費者的消息，還是消費者用完的Ack消息，都是同一種消息。OneFlow的Actor通信中是不需要指明“Ack”的。各個Actor在處理ActorMsg的時候都可以從Regst中得知是不是Ack。

• kCmdMsg：一些控制指令信號。不包含數(shù)據(jù)。如kConstructActor（Thread直接處理的消息，用于Thread創(chuàng)建Actor）；kStart，Actor啟動并開始工作。運行時靠著Start消息的傳染，整個計算圖開始工作。

• kEordMsg：表示訓練結束，Actor可以切換到Zombie狀態(tài)。運行時靠著Eord消息的傳染，整個計算圖中的Actor均切換到Zombie狀態(tài)，等待銷毀和RunTime下線。運行時的結束不是一下子就結束的，有可能計算圖的源節(jié)點已經(jīng)發(fā)出了Eord的信號，并將自己切換成Zombie狀態(tài)，而計算圖中的后半部分還在工作中。

通過數(shù)據(jù)流以去中心化的方式控制整個計算圖的工作，是OneFlow區(qū)別于其他框架的一大特色。

3.5.4?Register

oneflow/core/register/路徑

在初始化全局對象時，會創(chuàng)建Global對象RegstMgr。每臺機器上的RegstMgr管理了所有的Regst。

RegstMgr

RegstMgr在初始化時就會根據(jù)Plan申請所有的本機上的內存：HostMemory、HostPinnedMemory（For CUDA CopyH2D）、DeviceMemory、LockedMemory（For RDMA）等。并根據(jù)Plan中的Regst配置信息分配相應的內存地址給Regst。Regst的內存地址是固定的，直到運行時結束Regst的內存地址和大小都不會變化。OneFlow的靜態(tài)內存管理是Runtime啟動時統(tǒng)一分配，Runtime結束時統(tǒng)一銷毀。運行時的內存調度開銷是0。

Regst

Regst是OneFlow運行時的基本內存單元，也是基本的消息單元，Actor之間的通信、所有的數(shù)據(jù)生產(chǎn)、消費、回收都是Regst。由于OneFlow是靜態(tài)內存分配，內存的分時復用調度是編譯期的內存復用算法已經(jīng)做好了（通過控制邊+offset方式），所以運行時僅需要按照編譯期生成的MemChunk、MemBlock、Regst的配置描述（RegstDescProto）信息一次性申請內存，并分配給對應的Regst即可。

Regst存儲了兩類信息：

• 生產(chǎn)者Actor id和消費者 Actor ids。一個Regst的生產(chǎn)者是唯一的，消費者可能有多個。

• Blob的信息

由于歷史原因（在介紹ExecGraph和ExecNode時也提到了），Actor內部可能會有一個執(zhí)行子圖（多個op/kernel），Actor的產(chǎn)出消費Regst均可能包含多個Blob（Tensor）。Regst需要管理blob name in op -> logical blob id -> blob的映射（blob name in op -> logical blob id 是op自己管理的），使得Kernel在執(zhí)行時可以直接根據(jù)blob name拿到對應的blob指針。

未來會精簡Regst的設計，一個Regst只包含一個blob；合并Tensor和Blob概念。

Regst相關概念

• RegstDesc 編譯期的Regst描述類（C++），提供元信息，關聯(lián)Task，包含mem block，包含regst_num。RegstDesc 與Regst是一對多的關系（相鄰Actor流水并行執(zhí)行的關鍵）。TaskNode的Build過程中的Produce/ConsumeRegst就是在創(chuàng)建和消費RegstDesc

• RegstDescProto 配置文件的proto描述，存儲在Plan中

• RtRegstDesc 運行時的Regst描述類（C++），關聯(lián)Actor，提供計算Size的接口

• Regst 運行時的Regst，存儲真正的blob內存，被Actor所管理。

關系：RegstDesc（Compiler）-> RegstDescProto（Plan）-> RtRegstDesc（Runtime）-> Regst（Runtime， 1 to n）

Blob相關概念

• BlobDesc 編譯期的Blob描述類（C++），提供元信息：Shape、DataType；Op的InferBlobDesc就是在推導BlobDesc。

• BlobDescProto Blob配置文件的Proto描述，存儲在Plan中（RegstDescProto中）

• RtBlobDesc 運行時的Blob描述類，跟BlobDesc的區(qū)別是提供Header和Body的Size/CudaAlignedSize

• Blob 運行時Kernel操作的基本數(shù)據(jù)對象。存儲在Regst中。

需要注意的是，由于RegstNum > 1時，同一個RegstDesc會有多個Regst，多個Regst會存儲多個相同BlobDesc的Blob，所以Kernel每次運行拿到的Blob指針、Blob中的數(shù)據(jù)地址也可能是不同的。

Tensor相關概念

• user_op::TensorDesc ?BlobDesc在UserOp框架下的代稱

• user_op::Tensor Blob在UserOp框架下的代稱

3.5.5 Actor

Actor是OneFlow運行時的基本單元。編譯期的主要工作就是把用戶定義的邏輯計算圖和分布式集群環(huán)境編譯成Plan。Plan由Actor的描述信息TaskProto組成。所以運行時就是根據(jù)Plan中的所有TaskProto創(chuàng)建所有的Actor。

3.5.5.1 創(chuàng)建流程

1） Runtime對象通過ThreadMgr->Thread->AddTask的方式新增一個Actor （HandoutTasks）

2） Runtime對象通過ActorMsgBus給每個Actor發(fā)送 ConstructActor的指令消息（ActorCmd::kConstructActor）

3）每個Actor所在的Thread收到構造Actor的消息后調用ConstructActor接口構造Actor，其中是使用了NewActor ，傳入ThreadCtx，調用Actor的Init方法初始化該Actor。

Actor::Init(JobDesc, TaskProto, ThreadCtx)

我們看看Actor的初始化過程中做了哪些事情：

1） 根據(jù)ThreadCtx創(chuàng)建DeviceCtx 。運行時的Context有三級：ThreadCtx->DeviceCtx->KernelCtx 。對于Context的解釋我們放在Device部分詳細介紹。

2） 構造Kernel（ConstructKernel）

3） 創(chuàng)建Regst（NewRegsts）

在調用RegstMgr->NewRegsts之前，RegstMgr已經(jīng)給所有的Regst都申請好了內存，NewRegsts更應該像是GetRegsts。對于同一個RegstDesc，根據(jù)其regst_num會有多個Regst實例

4） 處理消費的RegstDescId以及Regst之間的Inplace

5） 虛接口VirtualActorInit，供各個子類Actor自己重載自定義的初始化內容

3.5.5.2 Actor狀態(tài)機

當Actor初始化完畢以后，Actor就進入了等待狀態(tài)。在Actor收到Eord信號并銷毀之前，Actor一直都在等待狀態(tài)和執(zhí)行狀態(tài)之間切換。Actor的狀態(tài)機我在之前的知乎文章中簡要介紹過：

?

Actor所有的邏輯都通過ProcessMsg來實現(xiàn)。Thread將收到的消息交給Actor處理，Actor處理消息過程中可能會觸發(fā)執(zhí)行（Act），執(zhí)行會Launch其內部的Kernel。執(zhí)行結束會向上下游Actor發(fā)消息。運行時的去中心化調度就是靠著Actor之間的消息通信所實現(xiàn)的。

Actor內部有多種MsgHandler來處理消息（HandlerNormal和HandlerZombie）。在Actor正常運行過程中都使用HandlerNormal來處理消息。HandlerZombie用于Actor在有序退出時的消息管理。

HandlerNormal

Actor正常運行過程中主要處理的消息是RegstMsg，其中包含了上游發(fā)來的可供該Actor消費的Regst 或者 下游使用完該Actor產(chǎn)出的某個Regst。在HandlerNormal中，Actor會解析RegstMsg并更新自己的狀態(tài)，然后觸發(fā)ActUntilFail。

ActUntilFail

在ActUntilFail中，Actor會判斷執(zhí)行條件是否滿足，如果滿足就一直執(zhí)行，直到失敗。執(zhí)行條件是否滿足需要兩個條件：IsReadReady和IsWriteReady，通常Actor需要判斷其消費的Regst都到齊了，且有空閑塊可寫時，才會觸發(fā)執(zhí)行。每次執(zhí)行都會觸發(fā)消息的發(fā)送：包括給上游和下游Actor發(fā)消息。

Act

每次Actor執(zhí)行稱之為一次Act。Act 是一個虛方法，需要子類具體實現(xiàn)。我們可以參考一個最常見的Actor：NormalForwardActor，我們所有的用戶級別的Op都使用NormalForward類型的Actor。這種Actor的Act方法里調用了AsyncLaunchKernel，去Launch內部的Kernel執(zhí)行。

3.5.5.3 異步執(zhí)行 與 異步消息發(fā)送

AsyncLaunchKernel

Actor內部的Kernel是異步調用的。每次Launch Kernel，Actor都要給該Kernel關聯(lián)此次執(zhí)行對應的Regst（流水并行對Kernel無感）。

Actor的所有消息發(fā)送都是異步的。見ActUntilFail。等Kernel異步執(zhí)行結束以后，相關的消息才會被發(fā)送出去。

Actor中需要對Inplace的Regst/Msg做特殊處理，因為Inplace會改變Regst的生命周期（延長ConsumedRegst的生命周期直到ProducedRegst生命周期結束）

Actor控制邏輯掩蓋

異步消息保存在Actor的async_msg_queue中。如果消息的接收者和本Actor在同一個WorkStream（Thread）中時，異步消息可以提前發(fā)送，不需要等待Kernel異步執(zhí)行完就可以通知其他Actor，由于在同一個Stream中，任務的執(zhí)行是有序的，該Actor的后繼Actor可以提前將任務也提交到相同的Stream中，等上一個任務執(zhí)行完，下一個任務一定可以滿足執(zhí)行條件并執(zhí)行。OneFlow的Actor機制通過相同Stream提前發(fā)消息就可以掩蓋GPU上絕大多數(shù)Actor的控制邏輯開銷。

3.5.5.4 Actor擴展性

Actor子類可以定制消息的處理方式，可以定制執(zhí)行條件。在NormalForward這類常見Actor中，Actor需要所有的輸入和輸出都滿足才會Act一次，且Act結束會將輸入還給上游、輸出發(fā)給下游。

同時OneFlow也擴展了多種特殊的Actor，如

• Repeat，收到一個輸入以后，就可以連續(xù)重復Act多次（只要有可寫的），重復Act結束以后才會返還輸入Regst。

• Unpack，收到一個輸入以后，會將該輸入解碼，拆成多個Regst分批次發(fā)給下游，發(fā)完以后才會返還輸入Regst。

• Input-wise，該Actor有多路輸入，每到達一個輸入，就可以Act，無需等待所有的輸入都到齊才去Act

等等。

需要指出的是：Repeat和Unpack分別對應時間上的Broadcast和Split。OneFlow的BatchAccumulate就是通過插入Repeat和Unpack op來實現(xiàn)的（反向梯度會插入Acc）。一個模型的Repeat num = 4 、數(shù)據(jù) Unpack num = 4的單卡訓練 跟 4卡數(shù)據(jù)并行 從數(shù)學上是完全等價的。

Actor通過子類的多種自定義行為使得整個系統(tǒng)很容易擴展，一些特殊的需求僅需要在OneFlow中新增一個類型的Actor就能完成。

3.5.5.5 流控機制

Actor天然支持流水線，運行時每個Actor自己通過判斷跟自己相關的消息就能得知自己能否執(zhí)行，不依賴中心調度結點，使用最簡單的FIFO原則就解決了流控問題（Control Flow）。我們用一個數(shù)據(jù)預處理的流水線時間線的例子介紹一下Actor的流水線和流控機制：

一個非常常見的數(shù)據(jù)預處理流程如下：

為了方便推演，我們假設DataLoding、Preprocessing、Copy、Training都是一個Actor（實際上Preprocessing和Training分別都是由多個Actor所組成的子圖）。當這4個Actor之間的RegstNum均為2時，如果訓練時間比較長（訓練是整個網(wǎng)絡的瓶頸），我們看到一種流水線的時間線如下圖：

當訓練到第3個batch時，4個Actor的執(zhí)行時間成一種反序遞進的方式規(guī)律執(zhí)行。圖中灰色表示奇數(shù)Batch的數(shù)據(jù)，藍色表示偶數(shù)Batch的數(shù)據(jù)，為了方便理解，其中標出了Batch 6 的數(shù)據(jù)隨著時間線的演進在整個pipeline中的流向。圖中相鄰兩條時間線中間的兩個小方塊表示RegstNum=2的Regst，白色表示空閑狀態(tài)，藍色表示被偶數(shù)Batch的數(shù)據(jù)占用，灰色表示被奇數(shù)Batch的數(shù)據(jù)占用。隨著時間線的演進，我們羅列出了Regst狀態(tài)變化時刻的新狀態(tài)。當訓練是瓶頸時，數(shù)據(jù)加載、預處理、拷貝傳輸?shù)拈_銷都被完美掩蓋在訓練時間中了。

那么當數(shù)據(jù)加載是瓶頸的時候呢？下面這個時間線更容易理解OneFlow流控是如何實現(xiàn)的：

由于Preprocessing是耗時最長的，所以預處理上游的Actor（DataLoading）的工作節(jié)奏（背壓機制， Back Pressure）以及預處理下游的Actor（Copy、Training）的工作節(jié)奏均被預處理這個Actor的執(zhí)行所控制。

OneFlow使用背壓機制解決流控問題。如上面兩張圖所示，雖然DataLoading的時間很短，但并不會無節(jié)制的加載數(shù)據(jù)，而是當Regst被填滿之后就會等待，當Training是瓶頸時，Batch 3的數(shù)據(jù)在訓練時，DataLoading提前準備了Batch 7和Batch 8的數(shù)據(jù)，然后就等著；當Preprocessing是瓶頸時，DataLoading永遠都比Preprocessing提前處理了兩個Batch的數(shù)據(jù)。

3.5.6 Opertor 與 Kernel

Operator是OneFlow計算圖（ComputeGraph）的基本單元，是計算圖中的節(jié)點，Tensor是計算圖上的邊。Operator是編譯期概念，對應的運行時概念就是Kernel。

oneflow/core/operator路徑下包含了Operator的基類以及一些系統(tǒng)Op；

oneflow/core/kernel路徑下包含了Kernel的基類以及一些系統(tǒng)Kernel；

UserOp（UserKernel）的描述在oneflow/core/framework下定義，而具體的每個UserOp/UserKernel都放在了oneflow/user路徑下。

Operator

bn_in_op表示blob name in op。一般的Op都會定義"in"作為輸入的Tensor名字，“out”是輸出的Tensor名字。

lbi/lbn：分別是LogicalBlobId和LogicalBlobName的縮寫，LogicalBlobName是一個Op的輸入輸出Tensor在邏輯圖上的唯一字符串，通常以"op_name/bn_in_op"的方式描述。如Conv1 Op的輸出Tensor的LogicalBlobName就是"Conv1/out"，LogicalBlobId是LogicalBlobName的結構化表達。

一個具體Op的主要行為是規(guī)定輸入輸出、提供Tensor的推導方法、提供合法的SBP Signature等。

Kernel

Actor中通常都會有一個Kernel，每次執(zhí)行就是異步Launch一次Kernel。Kernel的Forward函數(shù)（對應UserKernel的Compute函數(shù)）就是在做實際的數(shù)學計算，讀輸入的Tensor數(shù)據(jù)，將計算完的數(shù)據(jù)寫到輸出Tensor的內存/顯存上。對于GPU的Kernel，Kernel的計算實際上是向CUDA Stream提交異步任務（對應Kernel中的cuda kernel的定義和調用）。

向GPU提交異步的計算任務使用到了KernelCtx，KernelCtx由DeviceCtx構造而來，而DeviceCtx又由ThreadCtx構造而來。其中最重要的結構就是CudaStreamHandle。

3.5.7 Device

oneflow/core/device路徑

Kernel向Device（GPU）提交計算任務，使用到了cudaStream_t，這個cuda stream是哪里來的呢？

當一個Thread是GPUThread時，創(chuàng)建GPUThread會創(chuàng)建相應的ThreadCtx，其中包含了一個cudaStream以及cuda callback event的Channel。GPUThread除了自己輪詢ActorMsgQueue的線程以外，還會有一個callback的輪詢線程：cuda callback event poller 。Actor（Kernel）異步執(zhí)行計算任務結束后，cuda callback event poller線程會拿到相應的callback event，并會執(zhí)行該event。

Kernel執(zhí)行結束的callback是什么時候被插入的呢？在Actor每次Act（AsyncLaunchKernel）結束后，都會將發(fā)消息的動作作為一個CallBack（Actor::AsyncSendQueuedMsg）通過DeviceCtx->AddCallBack接口壓入cuda stream。

DeviceCtx（KernelCtx、ThreadCtx）的主要作用就是提供一個CudaStream供Kernel提交計算任務，同時提供一個callback的channel用于執(zhí)行Actor發(fā)消息的邏輯。通過這種設計，OneFlow的運行時就實現(xiàn)了Actor的異步執(zhí)行和異步消息機制。

Actor通過DeviceCtx異步LaunchKernel的示意圖如下：

由于相同Stream下的多個Actor，可以不用等Kernel的異步計算任務執(zhí)行完就發(fā)消息，所以可以將Actor通信開銷、Kernel的CPU代碼執(zhí)行開銷完全掩蓋在CudaStream的計算開銷中。其時間線如下圖所示：

其中WithoutOverlap對比了如果沒有Actor在同一個Stream中可以提前發(fā)消息的優(yōu)化，CudaStream中的計算會因為消息通信、LauchKernel的開銷導致GPU計算資源沒有被充分利用。

3.5.8 內存管理

oneflow/core/memory路徑

MemCase

Regst通過MemCase標記了自己所屬的內存類型，如果是GPU上的顯存，還需要標記自己所屬的DeviceId。如果是CPU上的主存，會標記該Regst是否是被CopyHD或CommNet所使用的。Regst通過MemBlockId和MemBlockOffset標記了自己所屬于哪個MemBlock以及對應的偏移量。

MemoryAllocator

根據(jù)MemCase和Size申請對應大小和類型的內存塊，返回內存塊首地址；根據(jù)內存地址回收內存。在Lazy情況下，僅在Runtime的啟動/結束時（RegstMgr的構造函數(shù)和析構函數(shù)里）才會申請/釋放內存。

MemBlock與Chunk

這是OneFlow的多級內存設計：Chunk -> MemBlock -> Regst。

• MemBlock：同一個Chain（MemChain，通常是GPU上的前后向的所有activation regsts在一個MemChain中，Optimizer子圖部分的Regst在各自的MemChain中）內的Regst根據(jù)分時復用的原則共用一個MemBlock的不同段，通過size和offset標記。內存復用算法會盡可能讓MemBlock的Size小，同時滿足互斥的Regst（生命周期有重疊的）不會有內存區(qū)域的重疊。

• Chunk：一個Job內在同一塊GPU上的MemBlock的合集稱為一個Chunk。Chunk的Size是所有內部MemBlock的Size之和。（即同一個Chunk內部的MemBlock之間沒有復用內存）

• 多個Job在同一個塊GPU上的Chunk，會根據(jù)Job之間的互斥關系，完整復用一個大的Chunk（取最大值）作為最終的Chunk。如TrainJob和EvalJob互斥，所以TrainJob的所有可復用的Regst的總Chunk跟Eval的總Chunk合并復用一塊內存。通常情況下，Eval只有前向，比TrainJob計算圖要小，可以完全被TrainJob的Chunk所包含。即新增一個EvalJob不會新增任何內存。

3.5.9 網(wǎng)絡模塊

oneflow/core/comm_network路徑

Global對象CommNet提供運行時多機之間收發(fā)ActorMsg、傳輸Regst數(shù)據(jù)功能。分為Epoll（基于Socket）實現(xiàn)和Ibverbs（基于RDMA）實現(xiàn)。其中RDMA需要注冊內存（鎖頁內存），會將對應的Regst內存注冊。

3.5.10 IO模塊

oneflow/core/persistence路徑

提供一系列磁盤讀寫操作的接口，主要用于跟IO相關的Kernel實現(xiàn)。

• FileSystem 提供文件系統(tǒng)相關操作，如創(chuàng)建文件、查詢目錄等。（文件、目錄的增刪改查操作）。目前支持POSIX文件系統(tǒng)和Hadoop文件系統(tǒng)。

• Snapshot 分為SnapshotReader和SnapshotWriter，對應Checkpoint的Load和Save操作。

• In/OutStream 提供文件讀寫流操作。如DataReader的Kernel讀取OFRecord就使用了PersistentInStream。

oneflow/core/record 路徑

除了OFRecord以外其余內容都是即將過時的老版本的decoder/encoder接口。新版本的易于擴展的data reader設計見 oneflow/user/data路徑下的各個文件。主要分為了data reader、dataset、parser等抽象。新的DataReader設計后面會專門出一篇文章介紹。

3.5.11 ID manager

Global對象IDMgr（ID編址系統(tǒng)）嚴格意義上不能稱之為一個模塊，在OneFlow中是一個很小的單元。負責id的壓縮和映射、負責編譯期Task、運行時Actor的唯一標識符TaskId/ActorId的編碼和解碼。64位的task id包含了10位的machine id、11為的thread id、21位的local work stream id、21位的task id。

IDMgr還提供各種映射接口，主要是GPU相關的各個線程id映射（計算、copy、nccl等）。

4

總結

OneFlow的運行時是一套非常簡潔、高效的Actor系統(tǒng)，通過簡單的消息機制就解決了分布式訓練中的復雜調度問題、流控問題，流水線的實現(xiàn)等。相比于其他框架的運行時，OneFlow的Actor實際上是對Kernel的一層很簡單很淺層的封裝，但是這一套抽象解決了運行時眾多Kernel對各種資源的管理、分布式并行引入的Kernel間復雜的時序依賴、狀態(tài)依賴等問題。Actor系統(tǒng)還非常的模塊化，同時易于擴展和組合，可以支持各種復雜的分布式深度學習訓練需求。

點擊“閱讀原文”，前往OneFlow代碼倉庫。