OneFlow系統設計
本文的主要內容如下：
? OneFlow 的設計目標
? OneFlow 的特色一：Actor 機制
? OneFlow 的特色二：SBP 機制
? 總結
一、OneFlow 的設計目標
OneFlow 的設計目標是追求極致的性能，特別是分布式多機多卡環境下的橫向擴展性，希望能讓用戶使用多機多卡就像使用單機單卡一樣容易，且享受線性加速的運行效率。
為什么 OneFlow 要聚焦于分布式場景的性能和易用性呢？隨著深度學習的發展，模型越來越大，訓練深度學習模型所需的算力會越來越高，同時模型增大的速度要大于 GPU 單卡顯存擴容的速度；訓練對算力的增長要求要大于 GPU 單卡算力增長的速度。單卡的計算能力和顯存遠遠不能滿足深度學習模型訓練的需求，必須借助多機多卡并行加速。
若深度學習框架可以讓互聯的多個 GPU 協同工作好，實現線性加速比，即使每塊 GPU 性能“一般”，也可滿足任意規模的算力需求，這就是所謂的橫向擴展性，堅信這是算力增長的解決之道。
但是，已有框架都是首先聚焦于單卡的用戶體驗，僅對適合數據并行的多機多卡場景處理的較好，即把單卡的計算圖鏡像復制到多機多卡上，各個卡和機器之間輔助于模型同步的模塊。
對于 BERT/GPT-3 等參數量巨大的模型，用戶在使用已有深度學習框架時常常會遇到多機多卡不好用、效率低下或無法實現等問題。用戶做分布式訓練常常需要較高的學習成本，還需要關心多機多卡之間模型的同步問題。業界為解決分布式深度學習的痛點，不僅改進深度學習框架自身，還研發了多種第三方插件，譬如 NCCL，Horovod，BytePS，HugeCTR，Mesh-tensorflow，Gpipe 等等，但仍不能滿足用戶極致的性能需求。
OneFlow 的核心設計理念是，讓多機多卡分布式訓練高效地協同運轉，同時要讓用戶在多機多卡的訓練體驗上，就像單卡一樣簡單容易。下面來介紹OneFlow 實現此目標最核心的兩點想法，來說明 OneFlow 是如何看待分布式場景下的深度學習訓練的。
二、Actor：一套簡潔的機制解決幾乎所有技術難題
關鍵特性：
? 去中心化調度
? 流水線
? 數據搬運是一等公民
? 傳輸被計算掩蓋
? 控制邏輯被執行邏輯掩蓋
在 OneFlow 的設計中，分為 Compile 和 Runtime 兩個時期，Compile 時期把用戶定義的神經網絡、分布式環境信息等編譯成一個靜態圖的執行計劃 Plan，Plan 由執行單元 Actor 的描述信息組成；Runtime 時期，各個機器根據 Plan 里的 Actor 描述信息真實地創建屬于自己機器的眾多 Actor 實例，然后啟動 Actor 運行系統。整個深度學習訓練期間，OneFlow 的執行基本單元就是 Actor，Actor 對應靜態執行圖上的節點，Actor 之間生產、消費的數據存儲在 Register 中，Actor 之間通過消息傳遞來協作運行。

1. Actor 機制實現去中心化調度
   OneFlow 的運行時去中心化調度就是用 Actor 機制實現的。在整個由 Actor 構成的靜態圖中，沒有一個中心的調度節點，每個 Actor 都只需要關心自己所需數據的生產者（上游 Actor ）和自己生產的數據的消費者（下游 Actor）即可。這樣在超大規模的分布式訓練場景下， 完全的去中心化 調度可以避免中心調度的單點性能瓶頸問題。
   每個 Actor 內部都有一個 狀態機 ，Actor 收發的消息、執行的情況都會改變自己的狀態。需要注意的是，Register 是存儲塊，存放了 Actor 生產出來的數據，而消息是包含了 Register 存儲塊的內存地址的輕量級數據，Actor 之間傳遞的是消息，而不是 Register，這樣就實現了 zero-copy。
   當 Actor 收到了新消息，判斷它執行所需要消費的 Register 已經就緒，且它將要生產的數據有空閑的 Register 可以寫入時，這個 Actor 就執行（Act）一次，生產出一個 Register。
   生產完以后，該 Actor 就向需要消費這個 Register 的那些消費者 Actor 們發消息，表示 “可以來讀取生產的數據了” ；同時該 Actor 還需要把消費完的那些 Register 還給這些 Regsiter 的生產者 Actor ，表示 “用完了數據，可以回收了” 。Actor 內部的狀態機如圖1 所示。
   

圖1 Actor 內部狀態機
在 Actor 啟動之后，會根據與其它 Actor 之間收發消息來切換自己的兩個狀態：等待狀態 和 執行狀態 。
一個 Actor 收到的消息一般分為幾個類型：
? 上游的生產者 Actor 發來消息說：可以來讀生產的數據了；
? 下游的消費者 Actor 發來消息說：用完生產的數據了。
當這個數據被所有消費者都用完以后，就可以回收成為空閑塊等待下一次被該 Actor 重新生產一份新的數據。
一個 Actor 收到消息以后都會去嘗試判斷當前是否滿足執行條件，執行條件一般有兩個：
? 需要讀取的數據是否都到齊了；
? 是否有空閑塊可以拿來被生產。當滿足執行狀態以后 Actor 就開始調用自己內部的 Kernel 真實的去讀寫數據。
執行完畢后 Actor 會向上下游發消息：
? 向下游的消費者 Actor 發消息說：剛生產了一塊數據，可以來讀了；
? 向上游的生產者 Actor 發消息說：剛用完了之前發給的數據了。
Actor 只需要關心上下游的消息就能判斷自己能不能執行。每個 Actor 都通過自己內部的狀態機和收發消息機制，實現了完全去中心化 的分布式協同工作。
2. Actor 機制實現流水線
上面介紹了 Actor 的內部狀態機，Actor 之間的消息傳遞和數據傳遞是依賴 Register 實現的。一個 Actor 是否能執行，只跟兩個條件相關：
? 自己消費的那些 Register 是否可讀；
? 自己生產的那些 Register 是否有空閑塊可寫。
對于一個 Register，如果運行時給它分配多個空閑塊，那么相鄰的兩個 Actor 就可以同時工作，工作時間重疊起來，這樣就實現了各個 Actor 之間的流水線。理想狀態下整個靜態執行圖的執行時間就是整個系統中是性能瓶頸的那個 Actor 運行的總時間，其余 Actor 的執行時間都被流水線掩蓋起來了。
舉一個例子來解釋 Actor 機制下的流水線是如何運轉起來的。圖2是一個由3個 Actor（a, b, c）組成的計算圖的執行時序圖。其中深綠色的 Regst方塊表示正在被使用的 Register 塊，白色的 Regst 方塊表示同一個 Register 的備用空閑塊。
? 1）在 Time0 時刻，Actor a 產出了一個 Regst_a_0，Actor b 和 Actor c 由于沒有可讀的 Register，所以處在等待狀態。假設每個 Actor的執行時間都是單位時間。
? 2）到 Time1 時刻，Actor a 給 Actor b 發消息說可以來讀產出的 Regst_a_0 了，Actor b 收到了消息，并檢查自己生產的 Register b 是否有空閑 Regst 塊可用，發現有可用的 Regst_b_0，于是 Time1 時刻Actor b 執行，讀取 Regst_a_0，寫 Regst_b_0；同時 Actor a 還會去看自己是否有空閑塊可寫，發現有，Time1 時刻 Actor a 也在執行，寫 Regst_a_1（這里需要說明的是，Regst_a_0 和 Regst_a_1 邏輯上是屬于同一個 Register，只是空間上分成了不同的空閑塊備份而已。在深度學習訓練任務中，Regst_a_0 和 Regst_a_1 里存放的是同一個 operator 產出的不同batch的數據）。于是 Actor a 和 Actor b 就并行工作起來了。Actor c 由于沒有數據可讀，仍在等待。
? 3）到 Time2 時刻，Actor b 生產出了 Regst_b_0，于是給下游的消費者Actor c 發消息說可以來讀生產的 Regst_b_0，同時給上游的生產者Actor a 發消息說用完了的 Regst_a_0。此時 Actor a 已經把剛剛生產的 Regst_a_1 又發給了 Actor b，Actor b 檢查自己仍有 Regst_b_1 空閑，于是 Actor b 開始讀 Regst_a_1，寫 Regst_b_1；Actor c 收到 Regst_b_0，發現自己有 Regst_c_0 空閑，于是 Actor c 開始執行，讀 Regst_b_0，寫 Regst_c_0；Actor a 收到了 Actor b 用完還回來的 Regst_a_0，檢查 Regst_a_0 所有的消費者都用完了，于是將 Regst_a_0 回收，標記為空閑塊，同時 Actor a 還可以繼續執行，寫 Regst_a_2。

圖2 Actor 生產消費關系和執行時序圖
在上面的例子中，到了 Time2 時刻，其實 Actor a、b、c 都在工作，在深度學習訓練任務中，Time2 時刻 Regst_b_0、Regst_c_0 存放的是 Batch 0 的數據，Regst_a_1、Regst_b_1 存放的是 Batch 1 的數據，Regst_a_2 存放的是 Batch 2 的數據。通過一個 Register 有多個空閑塊的設計，Actor 機制就實現了流水并行。
在這里拋出一個更進一步深入的問題：整個數據流的執行像一個網絡，數據在網絡中的流動就完成了計算，如何避免生產者生產太快，消費者消費不及，以及如何避免生產者生產太慢，消費者感到饑餓的問題，這涉及到對計算、內存、傳輸帶寬的規劃，盡可能使系統的瓶頸之處最寬，需要解決流控（flow control）的問題以及資源分配問題（如每個 Actor 的 Register 到底分配幾個內存塊配額），這非常關鍵，也是 OneFlow 系統已解決的問題。
3. 數據搬運是一等公民
在多機多卡的分布式環境中，各個機器和各個設備之間的數據傳輸往往是影響系統的橫向擴展性的最重要因素，如果傳輸開銷可以被計算開銷掩蓋，那么分布式深度學習訓練就可以達到理想的線性加速比。相較于其它的框架，OneFlow 把數據搬運視為跟數據計算同等地位的操作，提出 數據搬運是一等公民 的思想。
已有框架在編譯期的關注焦點是數據計算，認為數據搬運是背后隱式發生的，因此在靜態分析計算圖時略過計算和搬運的重疊編排，OneFlow 在計算圖中顯式表達了數據搬運，而且在靜態分析時同等對待數據搬運和數據計算，以最大化重疊搬運和計算。
在最終的執行圖中，數據搬運操作也是一個個 Actor。除了在設備上做數據計算用的 Actor 以外，還有計算機內存到 GPU 顯存之間的數據拷貝 Actor，機器之間做網絡通信的網絡 Actor，負責數據的切分、合并、復制的Actor，負責讀取磁盤數據的 Actor，負責加載保存模型的 Actor 等等。很多其它框架都把數據加載、多卡模型梯度的同步、網絡、模型加載更新等分別做成一個單獨的模塊，而 OneFlow 的設計是所有的功能都在一張由Actor組成的靜態執行圖里實現了。OneFlow 這樣的設計不僅簡潔、優雅，還非常高效。

圖 3 數據是如何從一個設備搬運到另一個設備上的
圖3表示了沒有 GPU-Direct 的況下，在 OneFlow 的 Runtime 階段，一個設備上的計算節點如果消費了另一個設備的計算節點，數據是如何搬運過去的。
4. 盡可能并行
在 OneFlow 的設計中，所有的出發點都是希望可以盡可能并行，從而達到最優的分布式性能。比如考慮到分布式訓練模型梯度同步時，顯存到內存的傳輸帶寬高于機器之間的網絡傳輸帶寬，OneFlow 會做兩級的 scatter 和 gather 操作（本機的和各個機器之間的），用于增加 locality，提高整體性能。
又比如在異步啟動深度學習訓練時，Python 端用戶的控制邏輯跟 OneFlow 運行時的執行圖是并行執行的，同時 OneFlow 有一套互斥臨界區的設計保證執行的高效性和正確性。
數據加載部分無論是從磁盤讀數據還是從 Python 端喂數據，OneFlow 都能保證盡可能并行，使得計算設備不會因為要等數據而導致性能下降。
已有框架如果想要盡可能重疊數據搬運和計算，一般借助多層回調（Callback）函數，當嵌套層次過多時，會遇到所謂的 Callback Hell 麻煩，正確性和可讀性都可能下降。但在 OneFlow 中，以上的這些并行并發特性，都是在這一套簡潔的 Actor 機制下實現的，解決了令人頭禿的 Callback Hell 問題。
此外，在多機的網絡通信部分，OneFlow 底層的網絡通信庫原生支持 RDMA 的高性能通信，也有一套基于 epoll 的高效通信設計。而目前最流行的 Pytorch，多機還需要通過 RPC 來做數據同步。
三、OneFlow 如何做到分布式最易用
OneFlow 是目前分布式場景中支持數據并行、模型并行、流水并行等最易用的深度學習框架。用戶只需要像單卡一樣去搭建網絡模型，并告訴 OneFlow 有哪些機器哪些卡，OneFlow 就會用最高效的方式把這些機器和設備使用起來。
這源于 OneFlow 的一套獨特的設計：ConsistentView（一致性視角）。對于多機多卡，OneFlow 會 把它抽象成一個超級大的設備 ，稱之為邏輯上的設備，這個邏輯設備的顯存是實際多個物理設備的顯存之和，這個邏輯設備的算力也是實際多個物理設備的算力之和。
用戶只需要在這個邏輯上的超級設備里，定義深度學習模型是如何構建的，其余的便不需要用戶來操作，由 OneFlow 來完成邏輯上的設備到物理上的設備的映射。
這里先明確兩個概念：“邏輯上的”和“物理上的”。“邏輯上的”表示 OneFlow 把分布式集群抽象成一個超級計算機之后的計算和數據，“物理上的”表示那些真實的部署到各個機器和設備上的計算和數據。
深度學習網絡是由 Op 構成的計算圖，Op 之間生產消費 Tensor 數據。在多機多卡的環境下，一個邏輯上的 Op 會對應多個真實的物理上的 Op，每個物理上的 Op 實際執行的計算都是這個邏輯 Op 計算的一部分，一個邏輯上的 Tensor 也會對應多個物理上的 Tensor，每個物理上的 Tensor 都是邏輯 Tensor 的一部分。
對于其它的框架定義的分布式訓練，每張卡是一個“world”，多卡之間根據暴露出來的接口來同步模型梯度；而對于 OneFlow 而言，多機多卡也都是一個“world”，使用一套 Placement+SBP 的方式做全局的統籌管理。
Placement
在 OneFlow 的計算圖搭建過程中，每個計算 Op 都有一個屬性叫做 Placement，表示了該邏輯上的 Op，是要部署到哪些機器哪些設備上的。對于常見的數據并行，就是所有的 Op 都部署到所有的設備上。但 OneFlow 也支持用戶指定 Op 的 Placement，比如當網絡過大單卡根本放不下的時候，在 OneFlow 可以讓網絡的前一部分在一張卡上，后一部分在另一張卡上，用一種“接力”的方式工作，實現流水并行。
圖4展示了一種可能的 Placement 例子。用戶定義了一個由3個 Op 組成的網絡：Op_0 -> Op_1 -> Op_2。
其中 Op_0 和 Op_1 的 Placement 是 Device 0，Op_2 的 Placement 是 Device 1，這就是一個流水并行的例子，Oneflow 會自動在 Op_1 和 Op_2 之間插入需要的數據搬運的 Copy Op。

圖4 一個流水并行的Placement示例圖
SBP
SBP 是 OneFlow 獨有的概念，他是三個單詞的首字母組合：Split、Broadcast、PartialSum（以 PartialSum 為例，實際上還可以是PartialMin、 PartialMax 等 reduce 操作），全稱叫 SbpParallel，表示一種邏輯上的 Tensor 跟物理上的多個 Tensor 的映射關系。
其中 Split 表示物理上的 Tensor 是邏輯 Tensor 按照某一維度切分后得到的， Split 有個參數 axis，表示切分的維度，如果把多個物理上的 Tensor 按照 Split 的維度進行拼接，就能還原出邏輯 Tensor。
Broadcast 表示物理上的 Tensor 是跟邏輯上的 Tensor 完全相同的。
PartialSum 表示物理上的 Tensor 雖然跟邏輯上的 Tensor 形狀一致，但是物理上的 Tensor 里的值是邏輯 Tensor 里對應位置的一部分，如果把物理上的多個 Tensor 按照對應位置相加，即可還原出邏輯上的 Tensor。
圖5展示了 SBP 的簡單示例。

圖5 幾種 SbpParallel 的簡單情形
SbpSignature 是一個 SbpParallel 的集合，在 OneFlow 的設計里是 Op 的屬性，它描繪了一個邏輯上的 Op 被映射成各個設備上的多個物理上的Op以后，這些物理上的 Op 是如何看待他們輸入輸出Tensor在邏輯上和物理上的映射關系的。一個 Op 會有多個合法的 SbpSignature，一個最簡單的合法 signature 就是輸入輸出都是 Broadcast，這表示了這個 Op 需要整個邏輯上的 Tensor 數據。
當用戶構建的邏輯上的計算圖確定以后，OneFlow 在 Compiler 生成分布式的物理上的執行圖時，會考慮每個 Op 的 Placement 和該 Op 允許的合法 SbpSignature 列表，在其中選擇一個傳輸開銷最小的 SbpSignature 作為本次訓練的 SbpSignature，用于指導 Compiler 生成最高效的執行圖。
關于 Op 的合法 SbpSignature 的列表，舉一個矩陣乘法（matmul）的Op的例子。
定義: Y = matmul(A, B) , A, B, Y 都是 Tensor，表示 Y = AB。那么至少存在兩種合法的 SbpSignature：
? 1) Y: Split(0), A: Split(0) , B: Broadcast
? 2) Y: Split(1), A: Broadcast, B: Split(1)
兩種合法的 signature 在兩個設備上的示意圖如圖6所示。假設邏輯上的 MatMul 的輸入輸出 Tensor 的形狀是：
A(64, 10) × B(10, 50) -> Y(64, 50)

圖6 MatMul的兩種合法SbpSignature
且該 Op 分布在兩個設備上。在第一種 SbpSignature 下，0號設備上的A是邏輯上 A 的前一半，1號設備上的 A 是邏輯 A 的后一半（按照第0維切分），而兩個設備上的 B 跟邏輯上的 B 完全一致，兩個設備輸出的 Y 分別是邏輯上的 Y 的前一半和后一半。同樣可以分析第二種 SbpSignature。
值得一提的是，當 A 是數據，B 是模型的時候，第一種 SbpSignature 就是 數據并行 ，第二種 SbpSignature 就是 模型并行 。如果兩個相鄰的 MatMul op，前一個使用第一種 SbpSignature，后一個使用第二種 SbpSignature，整個網絡就實現了 混合并行 。
圖7是一個混合并行的示例，定義了 Y0 = MatMul_0(A0, B0) , Y1 = MatMul_1(Y0, B1) 這樣一個由兩個op組成的計算圖，其中A0, Y0, Y1是數據Tensor，B0, B1 是模型Tensor。

圖7 混合并行示例
在圖7中 MatMul_0 產出的 Y0 被 MatMul_1 消費，但是這兩個 op 對同一個 Tensor 的 SBP 看待方式是不一樣的，MatMul_0 認為 Y0 是 Split(axis=0) 切分，但是 MatMul_1 需要一個 Broadcast 的 Y0 輸入。這時候OneFlow會自動插入一個“萬能”的 Boxing Op 做必要的數據裁剪、拼接、搬運和求和等操作，使得所有的Op都可以在分布式環境下高效的拿到自己想要的數據。
另外在數據并行的時候，訓練的前向模型 Tensor 的是 Broadcast，對應反向傳播的梯度就是PartialSum，當 Optimizer 需要全部的梯度來更新模型時，就會觸發 OneFlow 的 Boxing 機制進行高效的梯度同步工作。
最易用的分布式并行框架
OneFlow 的這套 Placement + SBP + Boxing 的機制，可以使得用戶定義的計算圖中的 Op、Tensor 以任意的方式分布在各個機器和各個設備上，無論是數據并行、模型并行還是流水并行，對于 OneFlow 而言，都只是一個特定 Placement 下的特定 SbpSignature 的組合而已，用戶可以方便的配置，也可以交給 OneFlow 來做自動的處理。
另外，早在微軟推出 ZeRO-2 框架之前，OneFlow 就已經支持了類似的功能，多機多卡情況下，每個模型 Tensor 都只保存在其中一個設備上，降低梯度計算中的內存占用。
四、總結
綜上，在編譯期，OneFlow 通過設計一套數學上嚴謹的形式系統來表示所有合法的并行模式，并支持編譯器較方便地自動搜索最優并行方案。
在運行期，OneFlow 通過 Actor 系統最優地、靈活地支持并行、并發執行。OneFlow 的內核具有簡潔、高效和高擴展性的優點。
基于此設計，OneFlow 使得分布式訓練的性能達到極致，且分布式訓練跟單卡一樣簡單易用。

總結

以上是生活随笔為你收集整理的OneFlow系统设计的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。