都2020年了,開源AI框架OneFlow還有機會嗎?| 技術深度解讀

7月31日,一流科技在創業1300天后,宣佈開源自研深度學習框架 OneFlow。與目前業界主流的開源 AI 框架,以及國內已經開源的三款 AI 框架相比,OneFlow 似乎開發的不太早、開源的更晚,它要憑藉什麼在鐵板一塊的市場格局情況下殺出重圍?此前,InfoQ專訪了OneFlow 創建者、一流科技創始人袁進輝,深入解讀 OneFlow 背後的故事以及他對於 AI 框架技術、產品和商業化的思考。本文中,一流科技工程師成誠將詳細解釋OneFlow的設計思路及特色,從技術角度闡釋OneFlow在深度學習框架競爭中勝出的可能性。

本文主要內容如下:

  1. OneFlow的設計目標
  2. OneFlow的特色一:Actor機制
  3. OneFlow的特色二:SBP機制
  4. 總結

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收到的消息一般分爲幾個類型:a) 上游的生產者Actor發來消息說你可以來讀我生產的數據了;b) 下游的消費者Actor發來消息說我用完你生產的數據了。當這個數據被所有消費者都用完以後,就可以回收成爲空閒塊等待下一次被該Actor重新生產一份新的數據。

一個Actor收到消息以後都會去嘗試判斷當前是否滿足執行條件,執行條件一般有兩個:a) 需要讀取的數據是否都到齊了;b) 是否有空閒塊可以拿來被生產。當滿足執行狀態以後Actor就開始調用自己內部的Kernel真實的去讀寫數據。

執行完畢後Actor會向上下游發消息:a) 向下遊的消費者Actor發消息說我剛生產了一塊數據,你可以來讀了;b) 向上遊的生產者Actor發消息說我剛用完了你之前發給我的數據了。

Actor只需要關心上下游的消息就能判斷自己能不能執行。每個Actor都通過自己內部的狀態機和收發消息機制實現了完全去中心化的分佈式協同工作。

2 Actor機制實現流水線

上面我們介紹了Actor的內部狀態機,Actor之間的消息傳遞和數據傳遞是依賴Register實現的。一個Actor是否能執行,只跟兩個條件相關:1)自己消費的那些Register是否可讀;2)自己生產的那些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來做數據同步。

Placement+SBP : 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的方式做全局的統籌管理。

1 Placement

在OneFlow的計算圖搭建過程中,每個計算Op都有一個屬性叫做Placement,表示了該邏輯上的Op,是要部署到哪些機器哪些設備上的。對於常見的數據並行,就是所有的Op都部署到所有的設備上。但OneFlow也支持用戶指定Op的Placement,比如當網絡過大單卡根本放不下的時候,在OneFlow可以讓網絡的前一部分在一張卡上,後一部分在另一張卡上,用一種“接力”的方式工作,實現流水並行。

圖 4 一個流水並行的Placement示例圖

圖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。

2 SBP

SBP是OneFlow獨有的概念,他是三個單詞的首字母組合:Split、Broadcast、PartiaSum(以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) X B(10, 50) -> Y(64, 50)

且該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。

圖6 MatMul的兩種合法SbpSignature

圖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機制進行高效的梯度同步工作。

3 最易用的分佈式並行框架

OneFlow的這套Placement + SBP + Boxing的機制,可以使得用戶定義的計算圖中的Op、Tensor以任意的方式分佈在各個機器和各個設備上,無論是數據並行、模型並行還是流水並行,對於OneFlow而言,都只是一個特定Placement下的特定SbpSignature的組合而已,用戶可以方便的配置,也可以交給OneFlow來做自動的處理。

另外,早在微軟推出ZeRO-2框架之前,OneFlow就已經支持了類似的功能,多機多卡情況下,每個模型Tensor都只保存在其中一個設備上,降低梯度計算中的內存佔用。

總結

綜上,在編譯期,OneFlow通過設計了一套數學上嚴謹的形式系統來表示所有合法的並行模式,並支持編譯器較方便地自動搜索最優並行方案;在運行期,則通過Actor系統最優地、靈活地支持並行、併發執行。OneFlow的內核具有簡潔、高效和高擴展性的優點。基於此設計,OneFlow使得分佈式訓練的性能達到極致,且分佈式訓練跟單卡一樣簡單易用。

附上OneFlow相關鏈接,感興趣的同學可以關注~

OneFlow的代碼倉庫:
https://github.com/Oneflow-Inc/oneflow

OneFlow官網:
https://www.oneflow.org/

OneFlow文檔:
https://docs.oneflow.org/

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章