深度解讀：讓你掌握OneFlow框架的系統(tǒng)設(shè)計(jì)(上篇)
OneFlow開源5個(gè)月后發(fā)布的v0.3.2版本，相較于上個(gè)大版本，又新增了眾多算子和功能（如亞線性內(nèi)存優(yōu)化、Partial FC、足夠靈活易用的新版Checkpoint...）。另外，跟各位朋友預(yù)告一下：完備性（如Serving）、易用性（全新的API）也在快速推進(jìn)中，敬請(qǐng)期待。
兩個(gè)月前我們通過(guò)DLPerf項(xiàng)目證明了OneFlow的性能優(yōu)秀，其中PK了經(jīng)NVIDIA深度優(yōu)化后的各個(gè)主流框架實(shí)現(xiàn)以及官方實(shí)現(xiàn)，對(duì)比了ResNet50和BERT這兩個(gè)CV和NLP領(lǐng)域應(yīng)用最廣的模型在數(shù)據(jù)并行下的吞吐率和加速比。這兩個(gè)月我們又做了大量的實(shí)驗(yàn)，對(duì)比了其他各個(gè)應(yīng)用場(chǎng)景下的第三方框架，如：
超大規(guī)模人臉識(shí)別案例 vs InsightFace
Wide&Deep vs HugeCTR
GPT-2 vs Megatron-LM
SSP vs PipeDream
Optimizer-Placement-Optimization vs DeepSpeed ZeRO
Auto Parallelism vs FlexFlow
實(shí)驗(yàn)證明了OneFlow這一套簡(jiǎn)潔的抽象（SBP + Actor）在支持各種模型并行、混合并行、流水并行、自動(dòng)并行、ZeRO等方面是如此的簡(jiǎn)單高效。上述的每個(gè)特性，都有一個(gè)專門的第三方的框架項(xiàng)目對(duì)主流框架進(jìn)行魔改，至少涉及了數(shù)千行的項(xiàng)目代碼，而這些特性在OneFlow中要么是原生支持的特性，要么是新增一個(gè)Actor類型或者是一個(gè)圖優(yōu)化的Pass（幾十行的代碼）就可以支持的非常好。我們認(rèn)為OneFlow這套設(shè)計(jì)不僅是性能最快的框架設(shè)計(jì)，同時(shí)也是分布式深度學(xué)習(xí)訓(xùn)練框架中最簡(jiǎn)單、最易擴(kuò)展的框架設(shè)計(jì)。上述的每個(gè)實(shí)驗(yàn)，近期都會(huì)有專門的技術(shù)文章分享，感興趣的小伙伴可以期待一下~
如果你對(duì)OneFlow這套致簡(jiǎn)致快的框架設(shè)計(jì)感興趣，或者對(duì)深度學(xué)習(xí)框架、分布式系統(tǒng)感興趣的話，本文就會(huì)讓你全面掌握OneFlow的系統(tǒng)設(shè)計(jì)。相信讀完這篇文章，你就會(huì)理解我們是如何看待分布式深度學(xué)習(xí)訓(xùn)練的，我們?yōu)槭裁匆@樣設(shè)計(jì)，這樣設(shè)計(jì)的好處是什么，以及我們?yōu)槭裁聪嘈臤neFlow這套設(shè)計(jì)是分布式深度學(xué)習(xí)訓(xùn)練框架的最優(yōu)設(shè)計(jì)。
深度學(xué)習(xí)框架原理
OneFlow系統(tǒng)架構(gòu)設(shè)計(jì)（簡(jiǎn)略版）
OneFlow完整運(yùn)行流程與各模塊的交互方式
3.1 分布式集群環(huán)境初始化
3.2 Python端搭建計(jì)算圖
3.3 編譯期：OneFlow(JobSet) -> MergedPlan
3.4 編譯期：Compiler(Job)->Plan
3.5 運(yùn)行時(shí)：Runtime(Plan)
深度學(xué)習(xí)框架原理
深度學(xué)習(xí)框架是人工智能領(lǐng)域的“操作系統(tǒng)”，為深度學(xué)習(xí)相關(guān)的算法工程師提供一套簡(jiǎn)潔易用的用戶接口，使之能方便的搭建深度學(xué)習(xí)模型，進(jìn)行深度學(xué)習(xí)模型的訓(xùn)練、驗(yàn)證、測(cè)試、調(diào)參、遷移、部署、迭代開發(fā)等工作。同時(shí)深度學(xué)習(xí)框架作為底層硬件跟算法工程師之間的中間件，要做到設(shè)備無(wú)關(guān)，使得算法工程師可以不用關(guān)心具體的計(jì)算設(shè)備、存儲(chǔ)設(shè)備的細(xì)節(jié)就能方便的開發(fā)模型。
深度學(xué)習(xí)框架本質(zhì)上是一個(gè)基于張量（Tensor）之間的計(jì)算（Operator）表達(dá)式所組成的計(jì)算圖（Graph）編譯執(zhí)行引擎，提供了一系列張量的定義、一元操作、二元操作等數(shù)學(xué)原語(yǔ)，并根據(jù)反向傳播算法（Back Propagation）進(jìn)行梯度自動(dòng)求導(dǎo)以及模型更新。在大量數(shù)據(jù)分批次流入計(jì)算圖進(jìn)行模型訓(xùn)練之后，使得模型學(xué)習(xí)到數(shù)據(jù)中的內(nèi)在關(guān)聯(lián)關(guān)系，從而獲得對(duì)應(yīng)場(chǎng)景中的“智能”感知與判斷能力。
OneFlow系統(tǒng)架構(gòu)設(shè)計(jì)
OneFlow總體分為3個(gè)層次：Python前端、編譯期（Compiler）、運(yùn)行時(shí)（Runtime）。
Python端是用戶接口，是OneFlow啟動(dòng)、編譯、運(yùn)行的入口，負(fù)責(zé)構(gòu)建邏輯圖（Job），且負(fù)責(zé)運(yùn)行時(shí)跟底層計(jì)算圖執(zhí)行引擎交互，包括發(fā)送控制指令（運(yùn)行一個(gè)global_function / job）、喂數(shù)據(jù)（input）、處理輸出（output，callback）。
編譯期（Compiler）負(fù)責(zé)將前端用戶的定義的邏輯上的計(jì)算圖進(jìn)行編譯，產(chǎn)出實(shí)際上的物理計(jì)算圖 （Plan）
運(yùn)行時(shí)（Runtime）負(fù)責(zé)根據(jù)Plan創(chuàng)建真正的執(zhí)行圖——即一個(gè)由Actor組成的去中心化流式計(jì)算圖，每個(gè)Actor各司其職，有的Actor負(fù)責(zé)接收Python端的控制信號(hào)，有的Actor負(fù)責(zé)加載數(shù)據(jù)，有的Actor負(fù)責(zé)初始化模型、計(jì)算、更新、存儲(chǔ)、傳輸...，有的Actor負(fù)責(zé)返還給Python端數(shù)據(jù)，數(shù)據(jù)在計(jì)算圖中流動(dòng)，實(shí)現(xiàn)深度學(xué)習(xí)的模型訓(xùn)練功能。
OneFlow的設(shè)計(jì)原則是編譯期做大量的調(diào)度優(yōu)化、圖優(yōu)化、通信優(yōu)化、內(nèi)存優(yōu)化將用戶定義的邏輯計(jì)算圖編譯成分布式的物理計(jì)算圖，而運(yùn)行時(shí)通過(guò)一套極其簡(jiǎn)單的Actor系統(tǒng)就完成了去中心化調(diào)度，每個(gè)Actor僅需要關(guān)心自己的上下游就能知道自己什么時(shí)候該工作，什么時(shí)候該等待，省去了運(yùn)行時(shí)分布式訓(xùn)練中大量的調(diào)度開銷；同時(shí)這套機(jī)制還非常的高效和易擴(kuò)展，解決了分布式訓(xùn)練中各種復(fù)雜的并行難題、時(shí)序依賴、控制依賴難題，做到了將控制、傳輸盡可能掩蓋在計(jì)算任務(wù)中，使得分布式訓(xùn)練速度最大化。
總體架構(gòu)圖如下圖所示：
OneFlow完整運(yùn)行流程 & 各個(gè)模塊之間交互方式
我們通過(guò)介紹一次OneFlow完整運(yùn)行的流程來(lái)了解系統(tǒng)中的各個(gè)主要模塊是如何協(xié)同工作的。
3.1. 初始化環(huán)境（Env）
OneFlow是一個(gè)分布式計(jì)算系統(tǒng)，在Python前端啟動(dòng)時(shí)，第一件要做的就是初始化整個(gè)集群環(huán)境（Env）。環(huán)境由一個(gè)配置文件（EnvProto）所描述，里面包含了有多少臺(tái)機(jī)器，每臺(tái)機(jī)器的id、ip地址、控制端口號(hào)、數(shù)據(jù)傳輸端口號(hào)等信息。（Resource、MachineCtx是相近的概念，因?yàn)闅v史遺留原因目前還保留，未來(lái)會(huì)合并進(jìn)Env里）
OneFlow分布式環(huán)境啟動(dòng)目前有兩種方式，第一種是類MPI的方式啟動(dòng)，第二種是Master-Worker的方式啟動(dòng)。無(wú)論是哪種啟動(dòng)方式，目前Lazy執(zhí)行分布式都是Single-Client的模式（而PyTorch、TensorFlow都是Multi-Client的模式），未來(lái)Eager完善以后，也會(huì)支持Multi-Client模式啟動(dòng)。
為什么OneFlow的Lazy分布式要使用Single-Client方式啟動(dòng)？
這其實(shí)是因?yàn)镺neFlow有一套一致性視角（Consistent View）的抽象。OneFlow想把整個(gè)分布式集群抽象成一個(gè)“超級(jí)設(shè)備”，這個(gè)超級(jí)設(shè)備的計(jì)算資源、存儲(chǔ)資源都是實(shí)際整個(gè)分布式集群的總和。這個(gè)超級(jí)設(shè)備我們稱之為邏輯層面（logical plane），而實(shí)際分布式集群上每個(gè)機(jī)器/設(shè)備上的計(jì)算、集群內(nèi)的網(wǎng)絡(luò)通信稱之為物理層面（physical plane）。用戶僅需要在這個(gè)抽象后的超級(jí)設(shè)備上搭建深度學(xué)習(xí)模型并進(jìn)行訓(xùn)練就可以了。那么用戶的單機(jī)單卡訓(xùn)練腳本和分布式訓(xùn)練腳本就完全一致。而對(duì)于其他的所有框架，都是直接讓用戶在物理層面編程，如PyTorch的分布式訓(xùn)練數(shù)據(jù)并行需要在Optimizer里手寫跟其他設(shè)備的通信同步操作；如果是更加復(fù)雜的模型并行、流水并行、混合并行，用戶就要寫非常多的多機(jī)多卡上的通信和同步腳本，為此也有很多第三方框架如DeepSpeed、Megatron-LM幫助PyTorch支持復(fù)雜的并行方式。
由于OneFlow有一致性視角這個(gè)抽象，那么用戶的分布式訓(xùn)練腳本就仿佛是單機(jī)單卡的，所以僅需要由master節(jié)點(diǎn)去觸發(fā)計(jì)算圖的構(gòu)建、編譯和執(zhí)行即可，其他機(jī)器是worker節(jié)點(diǎn)，等待master發(fā)送執(zhí)行計(jì)劃（Plan）并啟動(dòng)運(yùn)行時(shí)執(zhí)行。
下面是OneFlow分布式啟動(dòng)的代碼細(xì)節(jié)：
如果是類MPI方式啟動(dòng)，各個(gè)機(jī)器會(huì)執(zhí)行相同的Python腳本，每個(gè)機(jī)器在執(zhí)行腳本時(shí)會(huì)判斷得知自己的machine_id，從而知道自己是不是master：
如果是master，則真正執(zhí)行python腳本，啟動(dòng)session、進(jìn)入global function、構(gòu)圖...
如果不是master，則在Python腳本的入口(oneflow.env.init())就卡住，進(jìn)入cluster的WorkerLoop()中循環(huán)、等待、執(zhí)行集群中master發(fā)來(lái)的指令（Eager::Instruction）、邏輯圖集合（Lazy::JobSet）。
如果是以ssh & worker的方式啟動(dòng)（目前主要使用這種方式，未來(lái)會(huì)替換成類MPI方式），則僅在master機(jī)器上啟動(dòng)了python進(jìn)程，master會(huì)把oneflow_worker可執(zhí)行程序通過(guò)ssh的scp命令拷貝到各個(gè)worker機(jī)器上，并根據(jù)配置執(zhí)行oneflow_worker程序，進(jìn)入WorkerLoop()的循環(huán)。
二者的區(qū)別：如果是類MPI的方式啟動(dòng)，各個(gè)機(jī)器上都需要安裝oneflow的python包，每個(gè)機(jī)器上僅需要一份python腳本即可；而以ssh & worker的方式啟動(dòng)，需要把oneflow_worker的二進(jìn)制文件臨時(shí)拷貝到各個(gè)機(jī)器上，不需要python腳本。
環(huán)境啟動(dòng)時(shí)做了什么事呢？
各個(gè)機(jī)器上啟動(dòng)了oneflow的進(jìn)程
創(chuàng)建CtrlServer和CtrlClient，并互相監(jiān)聽(tīng)直到每臺(tái)機(jī)器跟其他所有機(jī)器（包括自己）都建立了連接
Ctrl就是oneflow的控制平面（control plane），負(fù)責(zé)發(fā)送控制消息和數(shù)據(jù)，如master向worker發(fā)送JobSet、Plan等。
在OneFlow的Runtime階段，每個(gè)機(jī)器都會(huì)創(chuàng)建CommNet全局對(duì)象，這是OneFlow的數(shù)據(jù)平面（data plane），運(yùn)行時(shí)各個(gè)機(jī)器上的Actor之間的消息通信、數(shù)據(jù)傳輸均通過(guò)數(shù)據(jù)平面發(fā)送。
控制平面使用rpc方式通信和傳輸數(shù)據(jù)，簡(jiǎn)單直接；數(shù)據(jù)平面通過(guò)高性能的網(wǎng)絡(luò)（epoll，或者ibverbs）通信和傳輸數(shù)據(jù)，效率更高。這里補(bǔ)充一句，在使用ibverbs（RDMA）構(gòu)建數(shù)據(jù)平面的過(guò)程中，RDMA的數(shù)據(jù)傳輸需要使用注冊(cè)內(nèi)存（pinned memory，又稱鎖頁(yè)內(nèi)存， page-locked memory）。而各個(gè)機(jī)器之間需要通信知曉各自的注冊(cè)內(nèi)存地址，這是通過(guò)控制平面rpc的方式傳輸注冊(cè)內(nèi)存的元信息的。見(jiàn)：IBVerbsCommNet::RegisterMemoryDone (地址https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/comm_network/ibverbs/ibverbs_comm_network.cpp#L59)
3.2. Python端Job構(gòu)圖
在初始化環(huán)境之后，master上的python進(jìn)程會(huì)執(zhí)行用戶在global function中的構(gòu)圖代碼，生成Job。
Job是對(duì)整個(gè)邏輯圖的基本描述，有兩個(gè)主要部分：net和placement。
net是一個(gè)op list，表述了整個(gè)網(wǎng)絡(luò)是由哪些op以哪種方式連接起來(lái)的。net可以轉(zhuǎn)化成一個(gè)DAG，圖上的點(diǎn)表示一個(gè)op，圖上的邊表示op之間的產(chǎn)出和消費(fèi)的tensor。
placement表示了每個(gè)op放置在哪些設(shè)備哪些卡上。對(duì)于env里的所有機(jī)器以及所有設(shè)備，任何一個(gè)op都可以以任何方式分布在這些機(jī)器上。placement表示了邏輯上的op跟物理上的op的映射關(guān)系。
Python端通過(guò)C++（oneflow_internal_helper.h -> c_api_util.py）暴露出來(lái)的接口，實(shí)際上使用JobBuildAndInferCtx的AddAndInferOp接口進(jìn)行推導(dǎo)。JobBuildAndInferCtx會(huì)保存已經(jīng)加入的Op及其相關(guān)狀態(tài)（SBP、shape等），并根據(jù)新加入的OpConf推導(dǎo)生成新的Op及其相關(guān)狀態(tài)。同時(shí)JobBuildAndInferCtx會(huì)給Python端提供一系列查詢接口，這樣在Python的global function中的構(gòu)圖邏輯，后一個(gè)op的python代碼在執(zhí)行前，之前所有的op和tensor（的描述，TensorDesc）都已經(jīng)構(gòu)建好了，這樣就實(shí)現(xiàn)了在global function中“類似eager的方式構(gòu)圖。
在整個(gè)global function中的代碼都執(zhí)行完之后，JobBuildAndInferCtx會(huì)被調(diào)用Complete，生成最終的用戶定義的Job。
在Complete過(guò)程中，會(huì)調(diào)用執(zhí)行多個(gè)JobPass，每個(gè)pass是對(duì)Job進(jìn)行一次圖修改、重寫。其中最重要的pass就是生成后向op以及Optimizer（GenerateBackwardAndOptimizerOpConfs）。每個(gè)pass輸入是一個(gè)job，輸出是重寫后的job。很多性能優(yōu)化的pass也是這個(gè)時(shí)期做，比如“FuseAddToOutputPass”、自動(dòng)混合精度"AutoMixedPrecision"等。
用戶可能會(huì)定義多個(gè)global function（如cnn的train job和eval job），所有用戶定義的Job構(gòu)成一個(gè)集合（JobSet）。而OneFlow的C++主體對(duì)象Oneflow就只接收一個(gè)JobSet對(duì)象啟動(dòng)Complier和Runtime。
3.3. 編譯期：OneFlow(JobSet) -> MergedPlan
這部分我們介紹OneFlow支持多子圖編譯執(zhí)行的設(shè)計(jì)。用戶定義的多個(gè)計(jì)算圖（稱之為user job）會(huì)合并成為一個(gè)JobSet集合，同時(shí)OneFlow編譯期還會(huì)根據(jù)用戶定義的多個(gè)user job構(gòu)建出多個(gè)系統(tǒng)Job（稱之為SysJob），這些SysJob負(fù)責(zé)user job的輸入/輸出、模型加載/保存等操作。除了所有的user job和SysJob以外，OneFlow編譯期還會(huì)構(gòu)建一個(gè)MainJob，負(fù)責(zé)將這些計(jì)算圖連接起來(lái)，作為和Python端用戶的控制邏輯交互的主體。最終整個(gè)JobSet會(huì)編譯生成一個(gè)MergedPlan，交給運(yùn)行時(shí)啟動(dòng)。在構(gòu)建MainJob的過(guò)程中，一個(gè)非常重要的概念就是臨界區(qū)，一個(gè)計(jì)算圖會(huì)劃分多個(gè)臨界區(qū)，臨界區(qū)之間的互斥關(guān)系會(huì)保存在一個(gè)可重入鎖里，OneFlow使用臨界區(qū)的概念最大化重疊多Job的執(zhí)行和計(jì)算過(guò)程。OneFlow運(yùn)行時(shí)Python端的控制邏輯跟實(shí)際的計(jì)算圖執(zhí)行是異步的、多Job之間的執(zhí)行也是異步的，這樣就能將所有Python代碼的執(zhí)行時(shí)間完全掩蓋，如果兩個(gè)臨界區(qū)之間沒(méi)有互斥關(guān)系，那么他們也能并行執(zhí)行?？梢哉f(shuō)OneFlow這套JobSet->MergedPlan的設(shè)計(jì)就是為了異步并行執(zhí)行，使得運(yùn)行時(shí)訓(xùn)練盡可能快。
由于歷史原因，Oneflow的Complier僅編譯單個(gè)Job，多Job的編譯、Job間內(nèi)存復(fù)用、MainPlan等均在Oneflow的CompileAndMergePlanOnMaster接口中執(zhí)行。我們先假定Complier已經(jīng)將Job編譯成對(duì)應(yīng)的Plan了（Compiler的編譯過(guò)程我們放在后面講）。
Oneflow生成最終的MergedPlan的流程：
輸入是用戶定義的多個(gè)job（已經(jīng)過(guò)前后向展開以及各種圖優(yōu)化），我們稱之為user job。
3.3.1 構(gòu)建Model IO Job
Model IO Job中的每個(gè)Variable由全部user job中的Variable op name唯一確定。如果多個(gè)user job中有完全相同的Variable，則這兩個(gè)Variable是內(nèi)存共享的。即，Variable op的name是全局唯一的，是一個(gè)全局變量。舉例：train job和eval job中的同名Variable共享同一份內(nèi)存。
另外，Model IO Job（包含了三類 Model Init Job、Model Save Job、Model Load Job）中的Variable op，跟多個(gè)user job中的同名Variable也是內(nèi)存共享的。
目前，OneFlow里有兩種構(gòu)建Model IO Job的方式（MakeModelIoJobs， MakeModelIoV2Jobs），分別表示這些Variable是共用一個(gè)Init/Load/Save op去處理，還是每個(gè)Variable單獨(dú)一個(gè)Op去處理。下圖展示了Model IO Job的幾種形式：
請(qǐng)注意，Model IO中真正存放各個(gè)模型的Op類型是Output，是InterfaceOp的一種。而不是Variable。Output1產(chǎn)出的Tensor::Var1跟其他某幾個(gè)user job中的VariableOp::Var1內(nèi)存共享。以此類推。
OneFlow中有幾種類型的InterfaceOp:
Input（Python端的global function輸入Tensor）
Output
Return（Python端的global function的return Tensor）
Job之間的數(shù)據(jù)傳遞和綁定均通過(guò)InterfaceOp來(lái)實(shí)現(xiàn)。
InterfaceOp產(chǎn)出的Tensor的“RegstNum”恒為一，即僅有一份內(nèi)存塊，不支持流水；同時(shí)這塊內(nèi)存是被這個(gè)Tensor所獨(dú)占的，不會(huì)跟系統(tǒng)中的其他Op產(chǎn)出的Tensor內(nèi)存進(jìn)行內(nèi)存復(fù)用。
注：目前的Model IO是非常靜態(tài)的，非常不利于用戶對(duì)Checkpoint靈活使用的需求。@daquexian的新Model IO會(huì)徹底解決這個(gè)靈活性的問(wèn)題。PR見(jiàn)：Oneflow-Inc/oneflow#3540（地址：https://github.com/Oneflow-Inc/oneflow/pull/3540）
3.3.2 構(gòu)建Push/Pull Job
遍歷所有user job中的Input Op和Return Op，針對(duì)每個(gè)Input Op，分別構(gòu)建一個(gè)對(duì)應(yīng)的Push Job，針對(duì)每個(gè)Return Op，分別構(gòu)建一個(gè)對(duì)應(yīng)的Pull Job。Push/Pull的原理見(jiàn)下圖：
其中ForeignInput Op內(nèi)部維護(hù)一個(gè)buffer，該buffer等待Python端喂數(shù)據(jù)，當(dāng)數(shù)據(jù)喂完時(shí)該Op/Kernel執(zhí)行完畢。ForeignOutput Op內(nèi)部也有一個(gè)buffer，當(dāng)往該buffer內(nèi)填完數(shù)據(jù)以后，python端對(duì)應(yīng)的of blob對(duì)象中的numpy就拷貝了對(duì)應(yīng)的數(shù)據(jù)。
參見(jiàn)：
ForeignInputKernel 與 ForeignOutputKernel
Python端：OfBlob._CopyBodyFromNdarray(). OfBlob._CopyToNdarrayListAndIsNewSliceStartMask()
C++與Python端連接處：?Dtype_GetOfBlobCurTensorCopyToBufferFuncName
為什么OneFlow與Python端的數(shù)據(jù)交換需要通過(guò)兩種獨(dú)立的Job子圖實(shí)現(xiàn)？
有兩個(gè)目的：
新增Push/Pull Job，并使用內(nèi)存共享的方式，對(duì)原有的Job沒(méi)有構(gòu)圖上的破壞。
（主要目的）為了盡可能重疊Python與C++數(shù)據(jù)交換的過(guò)程。如何重疊？需要依賴OneFlow構(gòu)圖上的重要設(shè)計(jì)：MainJob和CriticalSection。我們放在下一節(jié)講。
3.3.3 編譯所有的job
順序編譯所有的user job和Model IO Job、Push/Pull Job。每個(gè)Job編譯時(shí)，都是用Compiler完整編譯至plan。且各個(gè)job之間不知道彼此的存在（歷史原因）。
Compiler將一個(gè)Job編譯成Plan的過(guò)程放在下一章節(jié)講。
3.3.4 生成MainJob并得到最終的MergedPlan
這個(gè)過(guò)程分為幾步。
1） 將每個(gè)Job生成的Plan（SubPlan）合并到一個(gè)大的MergedPlan中
2） Job之間的內(nèi)存復(fù)用和內(nèi)存共享 （OneFlow中的內(nèi)存共享和內(nèi)存復(fù)用是一個(gè)很大的話題，我們后面會(huì)專門單獨(dú)寫一篇文章分享其中的設(shè)計(jì)）
3） 計(jì)算CriticalSection
4） 生成MainJob
5） 編譯MainJob得到MainPlan
6） 將MainPlan和MergedPlan中每個(gè)Job生成的SubPlan進(jìn)行l(wèi)ink，得到最終的MergedPlan
CriticalSection
CriticalSection是OneFlow構(gòu)圖中一個(gè)非常重要的概念——臨界區(qū)。多個(gè)Job編譯的多個(gè)Plan分布在各個(gè)臨界區(qū)中。每個(gè)Job都關(guān)聯(lián)多個(gè)臨界區(qū)，臨界區(qū)有兩種類型：InputOutput 和 Total。其中InputOutput是根據(jù)這個(gè)Job的Input、Output、Return等特殊類型的Op專門設(shè)立的臨界區(qū)，Total是每個(gè)Job必有的臨界區(qū)，Job內(nèi)的所有Op都被包含在Total臨界區(qū)里。
整個(gè)JobSet會(huì)劃分成眾多臨界區(qū)，臨界區(qū)之間最重要的關(guān)系就是——互斥。如果兩個(gè)臨界區(qū)互斥，則其中一個(gè)臨界區(qū)在執(zhí)行的時(shí)候，另一個(gè)臨界區(qū)必須等待。如果兩個(gè)臨界區(qū)不互斥，則可以并行同時(shí)執(zhí)行。如何判斷兩個(gè)臨界區(qū)是否互斥？借助全局概念的Op——InterfaceOp和VariableOp，如果兩個(gè)臨界區(qū)中的全局Op有同名，則這兩個(gè)臨界區(qū)在執(zhí)行的時(shí)候會(huì)訪問(wèn)同一個(gè)全局的Op，則這兩個(gè)臨界區(qū)必然互斥，無(wú)法同時(shí)訪問(wèn)同一個(gè)全局Op。
臨界區(qū)是比Job更細(xì)粒度的概念（但跟Op相比，仍然是粗粒度的）。
為什么要把Job分成多個(gè)臨界區(qū)？
為什么要有InputOutput和Total兩種類型的臨界區(qū)？
原因是想讓不同的Job之間盡可能流水并行起來(lái)。如何使得相鄰的兩個(gè)有消費(fèi)關(guān)系的Job（Job A -> Job B, A的output被B的input消費(fèi)/共享）同時(shí)執(zhí)行？借助CriticalSection以及MainJob里的幾個(gè)重要組件，我們就能實(shí)現(xiàn)多Job之間盡可能的并行執(zhí)行。
Idea by @Ldpe2G
臨界區(qū)的互斥可以區(qū)分讀寫互斥，這樣多個(gè)只讀的臨界區(qū)可以并行執(zhí)行，在某些場(chǎng)景下可以更好的流水并行。
MainJob
MainJob的結(jié)構(gòu)圖如下：
Main Job 的結(jié)構(gòu)大體上反映了運(yùn)行時(shí)Python端跟OneFlow系統(tǒng)的交互情況：
1) Python端每調(diào)用一次global_function，都會(huì)向WaitAndSendIds op發(fā)送一個(gè)job id，WaitAndSendIds會(huì)把收到的job id對(duì)應(yīng)的多個(gè)CriticalSection id發(fā)送給ReentrantLock op。
2) ReentrantLock——可重入鎖，里面維護(hù)了所有臨界區(qū)之間的互斥情況，并且會(huì)維護(hù)一個(gè)等待隊(duì)列。其輸入有兩個(gè)：
一個(gè)是python端發(fā)來(lái)的控制指令說(shuō)要執(zhí)行哪個(gè)Job對(duì)應(yīng)的多個(gè)CriticalSection id，稱之為“start”輸入
另一個(gè)是esac返還回來(lái)的CriticalSection id，稱之為“end”輸入
start表示需要執(zhí)行哪個(gè)CriticalSection，end表示哪個(gè)CriticalSection已經(jīng)執(zhí)行完了。每個(gè)輸入都會(huì)更新可重入鎖中的等待隊(duì)列。由可重入鎖來(lái)判斷哪個(gè)CriticalSection可以被執(zhí)行。
舉例：start來(lái)了一個(gè)CriticalSection id 3。我們假設(shè)CriticalSection 3 與 0 互斥，且當(dāng)前CriticalSection 0 正在被執(zhí)行中，所以可重入鎖會(huì)讓3進(jìn)入等待隊(duì)列，直到0的執(zhí)行完畢信號(hào)還回來(lái)時(shí)（end 來(lái)了 0），3可以執(zhí)行了，那么才放3執(zhí)行。
3） ReentrantLock會(huì)根據(jù)內(nèi)部的臨界區(qū)互斥情況和等待隊(duì)列來(lái)判斷要向下發(fā)送真正可以立即執(zhí)行的CriticalSection id，發(fā)給Case op，Case Op執(zhí)行的就是一個(gè)switch的操作，觸發(fā)下面對(duì)應(yīng)id的CriticalSection去執(zhí)行。
4） MainJob的主體部分是所有的CriticalSection，注意在MainJob里每個(gè)CriticalSection用一個(gè)identity的tick op來(lái)標(biāo)識(shí)。當(dāng)整個(gè)MainJob編譯成MainPlan后，會(huì)執(zhí)行Link操作，將每個(gè)SubPlan連接替換MainPlan中的identity tick op。
5） 當(dāng)某個(gè)CriticalSection執(zhí)行完畢后，會(huì)給Esac op發(fā)消息。“Esac” 的命名是“Case”的字母逆序，因?yàn)槠涔δ芫褪歉鶦ase完全對(duì)稱相反的。Esac會(huì)把執(zhí)行結(jié)束的CriticalSection id發(fā)給ReentrantLock op（end輸入）用于更新?tīng)顟B(tài)。
另外圖中還有另外一個(gè)Esac op，僅連接了各個(gè)job對(duì)應(yīng)的Total Critical Section，該op接收某個(gè)Job執(zhí)行完畢的消息，并通過(guò)CallbackNotify Op發(fā)送給Python用于通知Python某個(gè)Job執(zhí)行結(jié)束了，可以執(zhí)行對(duì)應(yīng)的Callback了（如loss收集、acc統(tǒng)計(jì)等）。
ReentrantLock + CriticalSection 實(shí)現(xiàn)Job之間的流水并行
我們假設(shè)一種Job之間的消費(fèi)情況：Job A -> Job B，B消費(fèi)A的輸出，A和B均對(duì)應(yīng)了多個(gè)CriticalSection，A對(duì)應(yīng)0，1，2；B對(duì)應(yīng)3，4，5，其中1，4是TotalCriticalSection類型，其余是InputOutputCriticalSection類型。由于兩個(gè)Job僅在輸入輸出之間有消費(fèi)關(guān)系，所以僅有2，3互斥，其余均不互斥?；コ怅P(guān)系如下圖：
所以在ReentrantLock那里，僅會(huì)把2，3互斥相互block住。而兩個(gè)Job的主體：(A, 1), （B, 4）是不互斥的。故當(dāng)條件允許時(shí)，Job A和Job B可以流水并行執(zhí)行。如果沒(méi)有InputOutputCriticalSection，則A和B是一定會(huì)串行執(zhí)行的。
在OneFlow中，Push/Pull Job跟對(duì)應(yīng)的UserJob就是通過(guò)上述方式進(jìn)行流水并行的。通過(guò)這種設(shè)計(jì)，用戶定義的Python端喂數(shù)據(jù)的邏輯（Push Job）可以跟上一個(gè)Batch的計(jì)算任務(wù)完全重疊起來(lái)；Python端對(duì)每個(gè)Batch計(jì)算完的返回值（如Loss、Accuracy等）的處理（Pull Job）也可以被計(jì)算完全掩蓋。
由于Loss的返回一般發(fā)生在計(jì)算圖的中間（Forward->Loss->Backward->Optimizer），通過(guò)Return Op對(duì)應(yīng)的Pull Job可以不用等計(jì)算圖結(jié)束就可以執(zhí)行，Python端拿到Loss數(shù)據(jù)之后的統(tǒng)計(jì)處理操作可以被后向計(jì)算完全掩蓋。
在下一篇《僅此一文讓你掌握OneFlow框架的系統(tǒng)設(shè)計(jì)（中篇）》中，我們會(huì)介紹編譯期Complier如何將Job編譯成Plan的過(guò)程，其中會(huì)簡(jiǎn)要介紹OneFlow編譯期最精華的Boxing章節(jié)，敬請(qǐng)點(diǎn)擊本次推送二條閱讀。
點(diǎn)擊下方“閱讀原文”，前往OneFlow代碼倉(cāng)庫(kù)。