深度解讀：讓你掌握OneFlow框架的系統(tǒng)設(shè)計(jì)(中篇)
本文是OneFlow系統(tǒng)設(shè)計(jì)分享系列文章的中篇，主要介紹OneFlow的編譯期Compiler如何將Job編譯為Plan的。其中最精華的部分是OneFlow的Boxing模塊，負(fù)責(zé)構(gòu)建兩個(gè)邏輯上的Op對(duì)應(yīng)的兩組物理上的Op在任意情形下的物理子圖，完成了分布式訓(xùn)練中各個(gè)機(jī)器各個(gè)設(shè)備之間的數(shù)據(jù)拷貝、切分、傳輸、通信的子圖搭建。值得一提的是，Boxing模塊的代碼實(shí)現(xiàn)是非常直觀且易擴(kuò)展的，使用了設(shè)計(jì)模式中的責(zé)任鏈模式（Chain of Responsibility），未來我們會(huì)結(jié)合OneFlow的代碼實(shí)現(xiàn)分享一些C++編程技巧的文章，以及為什么OneFlow要使用這些編程技巧，解決了哪些問題，敬請(qǐng)期待~
如果你對(duì)OneFlow這套致簡致快的框架設(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ì)。
本篇為系列內(nèi)容中的中篇，閱讀上篇，請(qǐng)點(diǎn)擊本次推送第一條。
深度學(xué)習(xí)框架原理
OneFlow系統(tǒng)架構(gòu)設(shè)計(jì)（簡略版）
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)
3.4 編譯期：Compiler(Job)->Plan
為了方便理解，我們?cè)俸喴枋鲆恍┲匾母拍詈统橄螅?/span>
Job：用戶定義的邏輯上的計(jì)算圖，由邏輯上的Op組成。
Plan：編譯生成的物理上的計(jì)算圖，由物理上的Task組成。
Task：運(yùn)行時(shí)Actor的配置描述，一個(gè)Actor與一個(gè)Task一一對(duì)應(yīng)，Task內(nèi)部有Op的運(yùn)行時(shí)描述Kernel的配置。Task并不一定關(guān)聯(lián)某個(gè)用戶計(jì)算圖Job中的邏輯上的Op，因?yàn)椴⒕幾g期會(huì)增加很多物理上的Op用于數(shù)據(jù)搬運(yùn)、網(wǎng)絡(luò)傳輸、切分/拼接等操作。Task中標(biāo)記了自己是在哪臺(tái)機(jī)器哪個(gè)設(shè)備上，并使用哪個(gè)線程工作等。Task還需要指定自己產(chǎn)出Regst的regst num、內(nèi)存類型、所屬內(nèi)存塊的偏移量等等信息。
Regst：運(yùn)行時(shí)的數(shù)據(jù)存儲(chǔ)、Actor之間通信的基本單元。存儲(chǔ)某個(gè)具體的Tensor。
編譯期（Compiler）的設(shè)計(jì)體現(xiàn)了OneFlow作為一個(gè)分布式深度學(xué)習(xí)訓(xùn)練框架的很多重要的設(shè)計(jì)原則：
1）一致性視角（Consistent View）
OneFlow把整個(gè)分布式集群抽象成為一個(gè)超級(jí)設(shè)備，用戶使用OneFlow做分布式訓(xùn)練跟做單機(jī)單卡的訓(xùn)練沒有任何區(qū)別。體現(xiàn)在：
編譯期Compiler僅在Master機(jī)器上編譯整個(gè)Plan，其他Worker機(jī)器等待獲取Plan啟動(dòng)運(yùn)行時(shí)即可。Master上編譯的Plan就包含了所有機(jī)器所有設(shè)備上的基礎(chǔ)單元——Task（Actor）。
所有的分布式訓(xùn)練過程中各個(gè)機(jī)器各個(gè)設(shè)備之間的數(shù)據(jù)通信、同步操作均被Compiler自動(dòng)生成，無需用戶關(guān)心和編寫分布式訓(xùn)練中的數(shù)據(jù)同步。
2）數(shù)據(jù)搬運(yùn)是一等公民
OneFlow將所有的數(shù)據(jù)加載、預(yù)處理、拷貝、網(wǎng)絡(luò)傳輸、Tensor的自動(dòng)切分/拼接/廣播/求和操作都抽象成了跟計(jì)算Op一樣的運(yùn)行時(shí)執(zhí)行體——Actor，即在分布式的物理計(jì)算圖上顯式表示了數(shù)據(jù)搬運(yùn)的操作。這樣做的好處是OneFlow可以感知到所有的數(shù)據(jù)搬運(yùn)、同步操作，因此編譯期Compiler可以更好的在整個(gè)物理計(jì)算圖上做全局調(diào)度，使得這些數(shù)據(jù)搬運(yùn)操作盡可能被計(jì)算操作所掩蓋，對(duì)數(shù)據(jù)搬運(yùn)操作的性能優(yōu)化轉(zhuǎn)而變成了圖分析與圖優(yōu)化。
3）編譯期全局調(diào)度，運(yùn)行時(shí)去中心化調(diào)度
OneFlow的運(yùn)行時(shí)是一個(gè)極其簡單的抽象——Actor，每個(gè)Actor僅需要關(guān)心和自己相關(guān)的上下游Actor的消息就可以知道自己能否工作，這樣做的好處是運(yùn)行時(shí)系統(tǒng)不會(huì)因?yàn)橛兄行恼{(diào)度節(jié)點(diǎn)導(dǎo)致性能瓶頸（在計(jì)算圖非常大的情況下）。為了做到這一點(diǎn)，OneFlow的調(diào)度工作大多都是在編譯期完成的，Compiler會(huì)做好全局的內(nèi)存調(diào)度、Op執(zhí)行調(diào)度、通信調(diào)度等工作，使得運(yùn)行時(shí)的調(diào)度開銷盡可能的低，從而達(dá)到更快的訓(xùn)練速度。
4）天然支持流水線，解決流控問題
編譯期Compiler通過推導(dǎo)和設(shè)置Task產(chǎn)出Regst的regst num，可以使得運(yùn)行時(shí)相鄰Actor之間可以流水并行起來。同時(shí)還可以通過背壓機(jī)制（Back Pressure）解決流控問題（Control Flow）。具體的Actor機(jī)制如何解決流水線和流控問題的討論我放在下篇中介紹。
在原生的OneFlow設(shè)計(jì)中，Compiler輸入是一個(gè)Job（用戶定義的op list），經(jīng)過編譯生成OneFlow的中間表示IR（Intermediate Representation）——Plan，Plan是一個(gè)被Runtime直接讀取就能生成運(yùn)行時(shí)執(zhí)行圖的描述。而上面介紹的OneFlow(JobSet)->MergedPlan是為了支持Python前端交互 + 多Job（Train/Eval同時(shí)做）而后設(shè)計(jì)出來的。我們下面介紹OneFlow的Compiler做了哪些事。
3.4.1 JobCompleter().Complete(job)
第一步，經(jīng)過JobCompleter將Job不斷重寫。經(jīng)過多個(gè)Pass以生成最終的Job。中間借助OpGraph抽象不斷推導(dǎo)新的Job對(duì)應(yīng)的邏輯圖。這些Pass包括一些優(yōu)化如插入KeepHeaderOnly節(jié)點(diǎn)；增加Source/Sink的Tick節(jié)點(diǎn)使得圖成為一個(gè)單源節(jié)點(diǎn)和單匯節(jié)點(diǎn)；增加控制邊；計(jì)算臨界區(qū)；以及使用XRT框架重新構(gòu)建Job。
XRT框架會(huì)將Job中的OpGraph進(jìn)行有選擇的合并，并選取使用XLA或者TensorRT來進(jìn)行編譯生成優(yōu)化后的Kernel。對(duì)于OneFlow而言，這些都是XrtLaunchOpConf，其Kernel都是XrtLaunchKernel。OneFlow系統(tǒng)并不關(guān)心其實(shí)現(xiàn)細(xì)節(jié)，實(shí)際上，經(jīng)過XRT優(yōu)化后的Kernel實(shí)現(xiàn)都是在其框架內(nèi)定義的頂層抽象：Executable 中存儲(chǔ)的，在XrtLaunchKernel的計(jì)算過程中調(diào)用executable->Run()去執(zhí)行。
3.4.2 生成OpGraph
Graph是OneFlow中的一個(gè)重要基礎(chǔ)抽象，各個(gè)重要的圖相關(guān)的概念（OpGraph、LogicalGraph、TaskGraph、ChainGraph、ExecGraph...）都繼承自Graph。Graph表示一個(gè)圖，里面保存著這個(gè)圖中的所有的節(jié)點(diǎn)Node和節(jié)點(diǎn)之間的連邊Edge。Graph上面提供一系列共用的遍歷方法（普通遍歷、拓?fù)浔闅v、BFS、DFS...），以及圖改寫（插入、刪除 節(jié)點(diǎn)/邊）圖查詢方法。
其實(shí)在第一階段JobCompleter在修改Job的過程中就需要多次Build OpGraph，在最終版本的Job生成以后，我們還需要在全局創(chuàng)建一個(gè)OpGraph，用于后續(xù)編譯過程中對(duì)各個(gè)邏輯Op和邏輯Tensor的查詢。
生成OpGraph分為幾步：（核心邏輯：OpGraph::InferLogicalBlobDesc https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/graph/op_graph.cpp#L563）
按照拓?fù)湫虮闅v每個(gè)Op（OpNode）
1） 推導(dǎo)ParallelSignature （Eager所需）
2） 推導(dǎo)BatchAxis（將要被廢棄，描述了哪一個(gè)維度是batch維，或者沒有batch維，如Variable那一支路上的op）
3） 推導(dǎo)MirroredSignature （推導(dǎo)每個(gè)Tensor是否是Mirrored，我認(rèn)為這個(gè)應(yīng)該跟Sbp成為同一級(jí)的東西：SBPM）
4） 推導(dǎo)SbpSignature
SBP是oneflow非常重要的概念，我在知乎文章——《都2020年了，為什么我們相信OneFlow會(huì)成功》 中有初步解釋了SbpParallel的語義：一種邏輯上的Tensor跟物理上的多個(gè)Tensor的映射關(guān)系。SbpSignature是一個(gè)SbpParallel的集合，在OneFlow的設(shè)計(jì)里是Op的屬性，它描繪了一個(gè)邏輯上的Op被映射成各個(gè)設(shè)備上的多個(gè)物理上的Op以后，這些物理上的Op是如何看待他們輸入輸出Tensor在邏輯上和物理上的映射關(guān)系的。
這里的推導(dǎo)SbpSignature，就是在每個(gè)Op多個(gè)合法的SbpSignature中搜索到一個(gè)最優(yōu)的（傳輸代價(jià)最低的）作為本次訓(xùn)練實(shí)際采用的SbpSIgnature。
在自動(dòng)并行（by @Yipeng1994 ）完成以后，推導(dǎo)SbpSignature就不再是按照拓?fù)湫蜇澬乃惴ㄍ茖?dǎo)，而是在全局搜索一個(gè)近似次優(yōu)解。
5） 推導(dǎo)Logical BlobDesc
此處是推導(dǎo)每個(gè)邏輯Op的邏輯Tensor的Shape、DType、is_dynamic等信息。
Op最重要的概念就是推導(dǎo)SBP，并根據(jù)SBP來推導(dǎo)Tensor的Shape。編譯期僅需要靜態(tài)推導(dǎo)出每個(gè)Tensor的形狀，以及特殊Op需要推導(dǎo)其Op/Kernel的特殊屬性：Inplace、TempBufferSize...
OpGraph是邏輯上的概念，當(dāng)OpGraph構(gòu)建完成后，每個(gè)（邏輯上的）Op、每個(gè)（邏輯上的）Tensor的描述信息都被推導(dǎo)、創(chuàng)建完成了。
3.4.3 生成LogicalGraph 【即將過時(shí)】
LogicalGraph是OneFlow的歷史遺留產(chǎn)物，在遠(yuǎn)古時(shí)期負(fù)責(zé)邏輯圖展開、后向生成、Model IO等工作。后面隨著OneFlow系統(tǒng)設(shè)計(jì)的演化，其功能逐步被OpGraph + JobCompleter + Pass所替代。之所以目前還保留，是因?yàn)镺p與TaskType的映射關(guān)系還保留在LogicalNode的不同子類中。在未來一段時(shí)間內(nèi)會(huì)移除掉LogicalGraph抽象，完全由OpGraph所取代。
3.4.4 生成TaskGraph
TaskGraph的生成過程是OneFlow編譯期最重要也是最精華的一部分。Task是Actor的編譯期抽象，Actor是Task的運(yùn)行時(shí)抽象。所以TaskGraph就描繪了整個(gè)運(yùn)行時(shí)計(jì)算圖的全貌。TaskGraph的生成過程分為兩部分：
構(gòu)圖部分 （https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/graph/task_graph.cpp#L156）
Build/Infer部分 （https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/compiler.cpp#L77）
3.4.4.1 構(gòu)圖部分
如何根據(jù)邏輯圖生成物理計(jì)算圖？
1） 遍歷LogicalGraph的每個(gè)LogicalNode，根據(jù)每個(gè)LogicalNode的placement：生成有序的ComputeTaskNode（專用于計(jì)算的TaskNode）
2） 遍歷LogicalGraph的每個(gè)LogicalEdge，根據(jù)前后LogicalNode的類型，找到對(duì)應(yīng)的生成這部分SubTaskGraph的方法：BuildSubTaskGraphMethod，執(zhí)行該方法給這兩個(gè)LogicalNode對(duì)應(yīng)的ComputeTaskNode連邊、新增節(jié)點(diǎn)構(gòu)圖。
在遠(yuǎn)古時(shí)期的OneFlow設(shè)計(jì)中，生成SubTaskGraph的方法是跟前后LogicalNode的類型相關(guān)。而在后面的boxing重構(gòu)中（by: @liujuncheng, 見：Oneflow-Inc/oneflow#2248 和 Oneflow-Inc/oneflow#2846 等），生成SubTaskGraph的方法被SubTskGphBuilder所推導(dǎo)，根據(jù)情況構(gòu)建紛繁多樣的SubTaskGraph。下面會(huì)粗略介紹一下其中的設(shè)計(jì)。
下圖展示了一種可能的SubTaskGraph構(gòu)建方式：在LogicalGraph中，邏輯上的Op_a產(chǎn)出一個(gè)Tensor X 供Op_b消費(fèi)，其中Op_a和Op_b的Placement分別是4，3，而Op_a和Op_b對(duì)X的SBP parallel根據(jù)各自O(shè)p的屬性、用戶指定/推導(dǎo)/自動(dòng)SBP的結(jié)果確定。Tensor X就是一條LogicalEdge。第一步：分別生成所有LogicalNode對(duì)應(yīng)的有序的ComputeTaskNode，Op_a的LogicalNode展開成4個(gè)CompTaskNode，Op_b的LogicalNode展開成3個(gè)ComputeTaskNode。第二步：這些ComputeTaskNode在Boxing架構(gòu)中會(huì)根據(jù)實(shí)際情況新增節(jié)點(diǎn)，并連邊，使得下面Op_b的3個(gè)TaskNode可以拿到其想要的那部分X的數(shù)據(jù)。
每個(gè)LogicalEdge就是一個(gè)邏輯上的Tensor，前后兩個(gè)邏輯上的Op對(duì)同一個(gè)Tensor的SBP、Placement看待可能一致也可能不一致。如何構(gòu)建這部分SubTaskGraph對(duì)應(yīng)的子圖呢？OneFlow提供了一系列SubTskGphBuilder，根據(jù)各種情況生成不同的子圖。
SubTskGphBuilder
構(gòu)建該子圖需要的全部信息是：源節(jié)點(diǎn)的CompTaskNode列表，匯節(jié)點(diǎn)對(duì)應(yīng)的CompTaskNode列表，源節(jié)點(diǎn)與匯節(jié)點(diǎn)的并行屬性（ParallelDesc，SBP），傳輸?shù)倪壿婽ensor的信息（Shape、Dtype、LogicalBlobId...）
目前OneFlow內(nèi)部有7種SubTskGphBuilder，每種Builder下面都可以根據(jù)SBP、Placement等信息自定義多種實(shí)際的構(gòu)圖方案，如SliceBoxingSubTskGphBuilder下面就有5種不同的構(gòu)圖情況，CollectiveBoxingSubTskGphBuilder下面又有7種集合通信的Builder。我們這里簡單介紹幾個(gè)常見的子圖構(gòu)建方式：
1） one to one
這是最常見的連接方式，即LogicalEdge的兩端節(jié)點(diǎn)在ParallelNum、SBP上對(duì)中間邏輯Tensor的看待方式完全一致，可以一對(duì)一的直連。在常見的數(shù)據(jù)并行情況下（如Mirror的方式），前后向Op都是一對(duì)一直連的。下圖展示了兩種一對(duì)一直連情況。
左邊是GPU內(nèi)部的一對(duì)一直連，右邊是當(dāng)兩個(gè)ComputeTaskNode不在同一個(gè)設(shè)備上時(shí)，我們會(huì)插入傳輸節(jié)點(diǎn)：CopyD2H（Device to Host），CopyCommNet（網(wǎng)絡(luò)傳輸），CopyH2D（Host to Device），使得一對(duì)一直連的匯節(jié)點(diǎn)CompTaskNode可以拿到對(duì)應(yīng)的Tensor。
2） collective boxing
集合通信（Collective Communication）大多采用NCCL的實(shí)現(xiàn)，包含了：AllReduce、ReduceScatter、AllGather、Reduce、Broadcast等操作。需要注意的是：
由于NCCL多個(gè)設(shè)備上的通信是在NCCL內(nèi)部實(shí)現(xiàn)的，在OneFlow的TaskGraph上，這些NcclTaskNode之間沒有顯式的連邊，但其實(shí)中間有隱含的同步操作。這樣如果在NCCL結(jié)點(diǎn)前后連控制邊不當(dāng)，可能會(huì)造成死鎖，所以系統(tǒng)中對(duì)NCCL附近的順序化連邊需要非常小心。
使用NCCL進(jìn)行集合通信操作，在構(gòu)圖上是one to one連接的。
下圖展示了OneFlow中使用NCCL進(jìn)行集合通信的collective boxing操作構(gòu)圖。
3） slice boxing
這種boxing涵蓋了oneflow中遇到的大多數(shù)跟SBP相關(guān)的Boxing。在SliceBoxingSubTskGphBuilder中提供了支持S2B、S2S、P2S、P2B、B2S等5種不同的SBP情形。
slice boxing會(huì)根據(jù)上下游兩組CompTaskNode的ParallelDesc、SBP的不同，把上面一組物理上的Tensor按照下游期望的SBP的方式分配給下游的一組CompTaskNode，同時(shí)考慮Machine id、CPU/GPU的不同，同時(shí)希望傳輸開銷、構(gòu)圖開銷盡可能少。
下圖展示了一種可能的S2S的slice boxing 情形。邏輯圖上SrcOp產(chǎn)出Tensor X供DstOp消費(fèi)，其中SrcOp在Machine0的GPU0、1以及Machine1的GPU2、3上產(chǎn)出SBP Parallel = Split(0) 的Tensor X，而DstOp在Machine3的GPU4、5以及Machine 4的GPu6上消費(fèi)SBP Parallel = Split(0)的Tensor X。故需要把已經(jīng)分成4份的Tensor X 先concat起來，然后再split成3份分發(fā)給各個(gè)DstOp。
我們?cè)倥e一個(gè)可能的P2B的例子。Src Op產(chǎn)出的兩個(gè)物理Tensor X分別是邏輯Tensor X的一部分值，經(jīng)過Add和Clone操作發(fā)給后面以Broadcast消費(fèi)X的兩個(gè)DstOp。
?
OneFlow中的Boxing設(shè)計(jì)是其分布式易用性以及分布式性能上最精華的一部分，這里僅介紹了其概況，后續(xù)會(huì)單獨(dú)出一篇文章分享其中的設(shè)計(jì)。
我們通過SubTaskGraphBuilder給每個(gè)LogicalEdge對(duì)應(yīng)的物理子圖構(gòu)圖，這樣就搭建起了整個(gè)TaskGraph，完成了邏輯圖到物理圖的映射。構(gòu)圖過程中，根據(jù)節(jié)點(diǎn)類型等信息可以給每個(gè)TaskNode分配Thread id、Area id等屬性。
Thread id 標(biāo)記了每個(gè)TaskNode（即Actor）工作在哪個(gè)線程上。由于分布式環(huán)境下每個(gè)機(jī)器上是一個(gè)進(jìn)程，所以每個(gè)TaskNode都會(huì)設(shè)置Machine id和Thread id。線程id分配的方式：CPU上是平均分配各個(gè)thread id；GPU上，同一個(gè)GPU的所有計(jì)算Task在同一個(gè)計(jì)算線程中；所有集合通信的Task在同一個(gè)NCCL線程中。這樣分配線程id的方式是因?yàn)榻?jīng)過實(shí)驗(yàn)驗(yàn)證，計(jì)算Task在相同線程中速度最快（最小切換開銷）。
Area id 【即將過時(shí)】標(biāo)記了不同類型的Op、TaskNode分別從屬于整個(gè)TaskGraph上的哪一個(gè)區(qū)域。有一些特殊的Area如kMdUpdtArea 標(biāo)記了這些Task是在Optimizer子圖部分的。然而Area id是一個(gè)過時(shí)設(shè)計(jì)。應(yīng)該被完備的Scope概念所取代，同時(shí)一個(gè)Op從屬于哪個(gè)Area也不是Op的類型決定的，而是Op在編譯期圖重寫的哪一個(gè)階段被插入所決定的。后續(xù)會(huì)把Area id移除。
ChainGraph?【即將過時(shí)】
在目前的設(shè)計(jì)中，TaskGraph還會(huì)生成ChainGraph，進(jìn)行Chain的合并，給每個(gè)Task上新增Chain id的屬性，用于將一組Task子圖標(biāo)記出來（方便做內(nèi)存復(fù)用）。
被合并到一個(gè)Chain中的這組Task有一個(gè)共性：在相同的Thread/Stream中執(zhí)行，當(dāng)Chain子圖中的源節(jié)點(diǎn)可以執(zhí)行以后，Chain子圖的所有后繼節(jié)點(diǎn)可以一股腦的執(zhí)行完，不需要依賴或者再等其他的節(jié)點(diǎn)。
Chain的合并算法在遠(yuǎn)古時(shí)期以Layer為單位進(jìn)行遍歷合并時(shí)是可以較好工作的，但是目前以O(shè)p為單位就顯得有些過時(shí)，尤其是在一些特殊的網(wǎng)絡(luò)（如包含where op）中會(huì)因?yàn)閳D的拓?fù)浔闅v順序的不同而有較大的合并效果差異，甚至是成環(huán)的BUG。
目前僅在內(nèi)存復(fù)用算法中依賴了Chain的合并結(jié)果。后續(xù)會(huì)重構(gòu)掉這塊，將Chain的概念從TaskNode中去掉。
3.4.4.2 TaskGraph的Build/Infer階段
在TaskGraph的構(gòu)圖完畢之后，Compiler會(huì)按照TaskGraph中TaskNode的拓?fù)湫虮闅v，依次構(gòu)建每個(gè)TaskNode對(duì)應(yīng)的各種信息：
1） 生成每個(gè)TaskNode的所有Regst，并把Regst綁定到TaskNode的出邊TaskEdge上。TaskNode::ProduceAllRegstsAndBindEdges （https://github.com/Oneflow-Inc/oneflow/blob/64c20462f245b5cbef4230a62fa06edff85411b3/oneflow/core/job/compiler.cpp#L77）
2） 將每個(gè)TaskNode的入邊TaskEdge中的Regst關(guān)聯(lián)到TaskNode中
3） 執(zhí)行每個(gè)TaskNode的Build過程
TaskNode
TaskNode根據(jù)其不同的TaskType 有對(duì)應(yīng)的TaskNode子類特化。每種類型的TaskNode其構(gòu)建過程都不同。最常見的是NormalForwardCompTaskNode，對(duì)應(yīng)了所有用戶定義的計(jì)算Op的Actor。每種TaskNode對(duì)應(yīng)一種Actor，其Actor內(nèi)部執(zhí)行的狀態(tài)機(jī)也不同。oneflow/core/graph/路徑下列出了目前所有種類的TaskNode子類及實(shí)現(xiàn)。
TaskNode的構(gòu)建過程中，內(nèi)部需要構(gòu)建Regst。
Regst
Regst是OneFlow中數(shù)據(jù)存儲(chǔ)、傳遞的基本單元。運(yùn)行時(shí)Actor之間的消息通信，數(shù)據(jù)傳遞都使用Regst。在目前的Regst設(shè)計(jì)中，一個(gè)Regst會(huì)包含多個(gè)Blob（廣義上的Tensor概念），但越來越多的需求是需要一個(gè)Regst僅包含一個(gè)Blob，后續(xù)的重構(gòu)中，會(huì)把Blob概念整合進(jìn)Tensor中，精簡這里的概念。
Tensor是用戶級(jí)別的概念，是獨(dú)立的的一塊數(shù)據(jù)，而Regst是Actor級(jí)別的概念，記錄了這個(gè)Regst是由哪個(gè)Actor生產(chǎn)的，并被哪些Actor所消費(fèi)的。
TaskNode的Build過程
TaskNode內(nèi)部會(huì)有一個(gè)ExecGraph（執(zhí)行子圖），執(zhí)行子圖上的節(jié)點(diǎn)稱之為ExecNode，邊稱之為ExecEdge。在遠(yuǎn)古的OneFlow設(shè)計(jì)中，每個(gè)TaskNode里是由多個(gè)Op組成的執(zhí)行子圖構(gòu)成的，每個(gè)Op對(duì)應(yīng)一個(gè)ExecNode，后面隨著性能優(yōu)化變成了一個(gè)TaskNode對(duì)應(yīng)一個(gè)Op。我們?nèi)匀槐Ａ袅薊xecGraph的設(shè)計(jì)，雖然在目前的絕大多數(shù)場景中ExecGraph里只有一個(gè)ExecNode，沒有ExecNode。
ExecNode 和 Op 的區(qū)別：（雖然我不止一次希望把ExecNode和Op合并）
在OneFlow最初的設(shè)計(jì)中，Op是一個(gè)描述概念，并不關(guān)心具體的某個(gè)Blob/Tensor，僅提供一系列方法用于推導(dǎo)，Op是無狀態(tài)的。而ExecNode是在某個(gè)具體的TaskNode內(nèi)部，同時(shí)要關(guān)聯(lián)具體的Regst，是有狀態(tài)的。
1） ExecGraph：絕大多數(shù)TaskNode的Build過程都是根據(jù)LogicalNode中的Op（CompTaskNode）/ 新建Op （CopyTaskNode），先構(gòu)建ExecGraph。
2） ExecNode：bind regst。在TaskNode中，入邊消費(fèi)的Regst和出邊生產(chǎn)的Regst內(nèi)部都維護(hù)了一個(gè)或多個(gè)lbi（logical blob id），用于標(biāo)識(shí)一個(gè)Blob（Tensor）。TaskNode的構(gòu)建過程中需要把這些Regst里的lbi跟Op內(nèi)部的BnInOp綁定起來。
3） ExecNode：InferBlobDesc。推導(dǎo)每個(gè)Blob/Tensor的Shape等信息（存儲(chǔ)在BlobDesc中）
3.4.4.3 TaskGraph & Plan 優(yōu)化
在TaskGraph Build結(jié)束以后，原本的Compiler還會(huì)對(duì)整個(gè)TaskGraph進(jìn)行一些優(yōu)化。在前后向分離、python前端的重構(gòu)中，Compiler這里的優(yōu)化被精簡成了幾步：
1） 移除空的Regst
2） 增加Chain內(nèi)的控制邊保證執(zhí)行順序
3） 推導(dǎo)Inplace的內(nèi)存共享
Inplace的推導(dǎo)使用了 InplaceLbiGraph 進(jìn)行推導(dǎo)。需要注意的是，我們?cè)贠p（UserKernel）里定義的SetInplaceProposalFn 僅是一種“建議”，而實(shí)際上這個(gè)Op的輸出和輸入能否Inplace，還需要經(jīng)過InplaceLbiGraph進(jìn)行推導(dǎo)以后才能決定。一些顯而易見的約束是，一個(gè)Tensor不能同時(shí)被兩個(gè)消費(fèi)它的Op進(jìn)行Inplace，因?yàn)镮nplace會(huì)改寫輸入的Tensor數(shù)據(jù)，是一種Mutable消費(fèi)。Inplace在一些情況下可以加速計(jì)算。
4） 推導(dǎo)時(shí)間形狀（time shape）
在OneFlow的Regst中，除了其中的數(shù)據(jù)有物理上的形狀，Regst本身也有時(shí)間形狀（time shape），表示整個(gè)網(wǎng)絡(luò)執(zhí)行一個(gè)Batch的數(shù)據(jù)，該Regst需要被生產(chǎn)幾次。time shape 有2維，最常見的是（1， 1），表示一個(gè)batch執(zhí)行一次。一些特殊的Op/Actor會(huì)修改時(shí)間形狀：Repeat/Acc、Unpack/Pack。由于這些特殊Op可能會(huì)嵌套，所以我們讓時(shí)間形狀有兩維，表示最多允許兩層嵌套。當(dāng)網(wǎng)絡(luò)中插入一個(gè)Repeat Op，會(huì)把該Tensor重復(fù)發(fā)送k次，其時(shí)間形狀就是（1，k）。當(dāng)網(wǎng)絡(luò)中插入U(xiǎn)npack Op，會(huì)把一個(gè)Tensor切分成k段，按k次分別發(fā)送給后面的Op（相當(dāng)于在時(shí)間上一種數(shù)據(jù)并行）。
如果網(wǎng)絡(luò)中連續(xù)插入多個(gè)RepeatOp，比如第一個(gè)Repeat將輸出的時(shí)間形狀修改為（1， k1），后續(xù)的Regst均為該時(shí)間形狀；再插入第二個(gè)Repeat，則輸出的時(shí)間形狀會(huì)被修改為（ k2， k1）。
3.4.4.4 生成Plan
最終TaskGraph中的每個(gè)TaskNode會(huì)生成Plan中的TaskProto，得到一個(gè)naive的Plan。
Plan里最重要的內(nèi)容就是所有的TaskProto，每個(gè)TaskProto就描述了運(yùn)行時(shí)的一個(gè)Actor所需的所有信息。
3.4.5 Improver(naive_plan) -> complete_plan
在naive的Plan生成之后，Improver會(huì)把Plan進(jìn)行改寫。
Improver的最初設(shè)計(jì)是為了推導(dǎo)RegstNum。
在我之前的兩篇知乎文章中，都提到了運(yùn)行時(shí)Actor機(jī)制的相鄰Actor流水線是通過RegstNum > 1來實(shí)現(xiàn)的。naive_plan中沒有推導(dǎo)RegstNum，所以所有的RegstNum均=1。而Improver中設(shè)計(jì)了一套算法，用于推導(dǎo)每個(gè)TaskNode對(duì)應(yīng)的RegstNum，但是算法依賴每個(gè)Actor的實(shí)際執(zhí)行時(shí)間。所以需要有試跑。
在TaskGraph中我們hack了代碼，使得所有的CopyHdTaskNode的MinRegstNum=2，也就是RegstNum=2，目的是為了讓數(shù)據(jù)預(yù)處理跟GPU計(jì)算可以流水并行起來，未來會(huì)刪除掉這個(gè)hack。TaskGraph上對(duì)于每個(gè)Regst都推導(dǎo)了其min、max的regst num，一般的數(shù)據(jù)Regst min = 1， max = inf。也有的regst，我們不希望有任何多余的備份，故讓這些Regst min = 1，max = 1。
由于試跑對(duì)于后續(xù)的OneFlow開發(fā)非常不友好，于是Improver這里的試跑一直都沒有被啟用。而且即使所有的RegstNum = 1，在非相鄰的兩個(gè)Actor之間也可以流水并行起來。
目前Improver中最重要的目的是為了推導(dǎo)內(nèi)存復(fù)用。內(nèi)存復(fù)用也經(jīng)歷了多個(gè)階段，一開始是使用一種染色算法對(duì)Regst進(jìn)行染色，相同顏色的共用一段內(nèi)存。后續(xù)我設(shè)計(jì)開發(fā)了內(nèi)存復(fù)用2.0（見 Oneflow-Inc/oneflow#2267、Oneflow-Inc/oneflow#2319），采用了Chunk、MemBlock、Regst三級(jí)內(nèi)存結(jié)構(gòu)，仍使用Improver作為入口。所以complete_plan中會(huì)比na?ve plan新增了Chunk和MemBlock的信息。OneFlow中的內(nèi)存復(fù)用設(shè)計(jì)后續(xù)會(huì)單獨(dú)出一篇文章進(jìn)行分享。
后續(xù)會(huì)將內(nèi)存復(fù)用算法放在Compiler中，使得Compiler的結(jié)果就是最終的Plan。
至此，我們就描述清楚了如何從一個(gè)Job編譯成一個(gè)Plan的全過程。
本文是OneFlow系統(tǒng)設(shè)計(jì)分享文章的中篇，主要介紹OneFlow完整運(yùn)行流程的中間部分：編譯期Compiler將Job編譯成Plan的過程。在下一篇《僅此一文讓您掌握OneFlow框架的系統(tǒng)設(shè)計(jì)（下篇）》中，我們會(huì)介紹OneFlow的運(yùn)行時(shí)（Runtime）以及倉庫源碼下的主要各目錄的模塊簡介，其中會(huì)包含Actor運(yùn)行時(shí)如何高效的調(diào)度，以及如何解決流水線和流控問題。閱讀下篇，請(qǐng)點(diǎn)擊第三條推送。
點(diǎn)擊“閱讀原文”，前往OneFlow代碼倉庫。