作者 | 一流科技工程師成誠
編者按：7月31日，一流科技在創業1300天後，他們宣布開源自研的深度學習框架OneFlow，此前，CSDN對CEO袁進輝進行了專訪。本文中，一流科技工程師成誠將從OneFlow設計思路、特點、開源趨勢等角度出發，詳細闡釋其在深度學習框架的競爭中勝出的可能性。
主要內容如下：
OneFlow的設計思路，它的獨特性OneFlow的特色一：Actor機制OneFlow的特色二：SBP機制對人工智慧/深度學習未來的看法對開源的感想和總結
我是OneFlow的一名工程師，在公司創業之初就加入了。研究生就讀於清華軟院，讀研期間在OneFlow實習了兩年，畢業之後正式入職一年。我對OneFlow的整體設計是比較熟悉的，對深度學習框架也有一點小小的見解。
時至今日，OneFlow即將開源，我相信OneFlow的設計和實現可以給大家一種眼前一亮的感覺，也相信在眾多深度學習框架的競爭中，OneFlow有機會勝出。
OneFlow是獨特的
時至今日，一個框架有沒有機會成功，要看它有沒有差異化的特點。OneFlow是有其獨特的設計理念和技術路線的。目前市面上已有的眾多開源框架，用戶最多的是PyTorch和TensorFlow，除此之外還有各大公司的自研框架：PaddlePaddle、MXNet、MindSpore、MegEngine等等。其中TensorFlow突出的特點是最完備，推理、serving、XLA、tensorboard等一應俱全，工業部署的主流還是TensorFlow；PyTorch的特點是易用性最好，eager執行、動態圖、跟python的交互方式等等，非常受科研人員喜歡。
可以說完備性和易用性這兩極分別被TF和PyTorch佔據，OneFlow作為一個小團隊打造的框架，如果單純的模仿別的框架，跟TF比完備性，跟PyTorch比易用性，那是絕對沒戲的。
但深度學習框架還有第三極：性能。OneFlow的特點就是追求極致的性能，而且是分布式多機多卡環境下的橫向擴展性。OneFlow的核心設計理念就是從分布式的性能角度出發，打造一個多機多卡體驗就像使用單卡一樣容易，而且速度最快的深度學習框架。
深度學習是吞沒算力的巨獸。多機多卡的理想很豐滿，現實很骨感，用戶會遇到多機多卡加速比不高，得不償失，用戶會遇到參數梁巨大時，現有框架不支持模型並行而無法訓練。為解決此類問題，業界不僅僅改進深度學習框架自身，還研發了多種第三方插件，譬如NCCL, Horovod, BytePS，HugeCTR，Mesh-tensorflow, Gpipe等等。但是，仍只滿足一小部分需求。
下面我就分別介紹兩點OneFlow相較於其他框架的獨特設計，來體現OneFlow是如何看待分布式場景下的深度學習訓練的。
Actor——用一套簡潔的機制解決所有令人頭禿的技術難題
（關鍵特性：去中心化調度流水線數據搬運是一等公民傳輸被計算掩蓋控制邏輯被執行邏輯掩蓋）
在OneFlow的設計中，分為Compiler和Runtime兩個時期，Compiler把用戶定義的網絡、分布式環境信息等編譯成一個靜態圖的中間表示（稱之為Plan）；Runtime時期，各個機器根據Plan裡的Actor描述信息真實地創建屬於自己機器的眾多Actor們，整個深度學習訓練期間，OneFlow的執行基本單元就是Actor，Actor之間通過消息機制通信，靜態執行圖上的節點就是Actor，節點之間的連邊是Register，Register存儲了Actor之間的生產者消費者信息，以及生產者Actor產生的數據。
1、Actor機制實現去中心化調度
OneFlow的運行時去中心化調度就是用Actor機制實現的。在整個由Actor構成的靜態圖中，沒有一個中心的調度節點，每個Actor都只需要關心自己消費的那些Actor，和消費自己生產出的數據的那些Actor即可。這樣在超大規模的分布式訓練場景下，完全的去中心化調度可以避免中心調度的單點性能瓶頸問題。
每個Actor內部都有一個狀態機，Actor之間的收發消息和Actor自己內部執行都會改變自己的狀態。當Actor收到了它執行所需要讀取的Register，且它當前有空閒的Register可以生產的時候，這個Actor就執行（Act）一次，生產出一個Register。生產完以後，該Actor就向需要消費這個Register的那些消費者Actor們發消息，表示你們可以來讀取我產生的數據了；同時該Actor還需要把它消費完的那些Register還給這些Regsiter的生產者Actor們，表示我用完了你們的數據，你可以回收了。Actor內部的狀態機如圖4所示。
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。
在上面的例子中，到了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機制就實現了流水並行。
3、數據搬運是一等公民
在多機多卡的分布式環境中，各個機器和各個設備之間的數據傳輸往往是影響系統橫向擴展性的最重要因素，如果傳輸開銷可以被計算開銷掩蓋，那麼分布式深度學習訓練就可以達到理想的線性加速比。相較於其他的框架，OneFlow把數據搬運視為跟數據計算同等地位的操作，提出「數據搬運是一等公民 」的思想。
已有框架在編譯期的關注焦點是數據計算，認為數據搬運是背後隱式發生的，因此在靜態分析計算圖時略過計算和搬運的重疊編排，OneFlow在計算圖中顯式表達了數據搬運，而且在靜態分析時同等對待數據搬運和數據計算，以最大化重疊搬運和計算。
在OneFlow最終的執行圖中，數據搬運操作也是一個個Actor。除了在設備上做數據計算用的Actor以外，還有計算機內存到GPU顯存之間的數據拷貝Actor，機器之間做網絡通信的網絡Actor，負責數據的切分、合併、複製的Actor，負責讀取磁碟數據的Actor，負責加載保存模型的Actor等等。很多其他框架都把數據加載、多卡模型梯度的同步、網絡、模型加載更新等分別做成一個單獨的模塊，而OneFlow的設計是所有的功能都在一張由Actor組成的靜態執行圖裡實現了。OneFlow這樣的設計不僅簡潔、優雅，還非常高效。
圖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例子：
2、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的簡單示例。
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：
Y: Split(0), A:Split(0), B:BroadcastY: 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。
在圖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這麼看重分布式的性能和易用性呢？這源於我們團隊對人工智慧和深度學習的未來的看法。
我們認為隨著深度學習的發展，模型會越來越大，訓練深度學習模型所需的算力會越來越高，同時模型增大的速度要大於GPU單卡顯存擴容的速度；訓練對算力的增長要求要大於GPU單卡算力增長的速度。所以分布式深度學習訓練會越來越常見、越來越重要。BERT、GPT-3等模型的出現印證了我們的判斷。單個晶片的能力總是受到物理約束，我們相信算力增長的解決之道在於橫向擴展，而這必須從系統軟體的角度去求解，一旦解決，即使是把性能「一般」的晶片互聯起來，只要系統軟體可以讓它們協同工作好，就可以滿足任何算力需求。
在分布式深度學習訓練中，性能至關重要。我們認為性能和橫向擴展性是深度學習框架的核心競爭力。我們在2017年就是這樣認為的，也是這樣做的，這就是OneFlow這個小團隊能存活至今的原因。
私以為，如果深度學習的未來是單機單卡就能搞得定的規模的話，那麼深度學習也就只能做做目前已知的一些簡單任務，深度學習的浪潮也會褪去。
開源，OneFlow還有很長的路要走
OneFlow團隊是一個小團隊，OneFlow作為一個開源框架還有很多不足。我們框架的易用性和完善性不如Pytorch和Tensorflow。我們團隊還在努力補上OneFlow的短板，同時也非常需要開源社區的支持。如果以爬山來打比方，OneFlow相當於從先難後易的南坡爬珠峰，其它框架相當於從先易後難的北坡爬珠峰。
OneFlow是獨特的，有區別於其他框架的技術路線，如果你也認同OneFlow的設計理念的話，歡迎您來幫助我們一起完善OneFlow。
我們相信，從技術上看，OneFlow 是更接近完美的那一個框架。
以上就是我對OneFlow這個深度學習框架的一點粗淺的理解，肯定有很多理解不到位的地方，希望大家批評指正。
OneFlow代碼倉庫：
https://github.com/Oneflow-Inc/oneflow
OneFlow官網：https://www.oneflow.org/
OneFlow文檔：https://docs.oneflow.org/
原文連結：
https://www.zhihu.com/question/409036335/answer/1373468192