2020 年，最轟動的 AI 新聞莫過于 OpenAI 發布的 GPT-3 了。它的1750億參數量及其在眾多NLP任務上超過人類的出眾表現讓人們開始堅信：大模型才是未來。但與之帶來的問題是，訓練超大模型所需的算力、存儲已不再是單機就能搞定。

據 NVIDIA 估算，如果要訓練GPT-3 ，即使單個機器的顯存/內存能裝得下，用 8 張 V100 的顯卡，訓練時長預計要 36 年；即使用 512 張 V100 ，訓練也需要將近 7 個月；如果擁有 1024 張 80GB A100， 那么完整訓練 GPT-3 的時長可以縮減到 1 個月。

除去硬件資源這個經濟問題，在技術層面，意味著訓練大模型一定是一個分布式問題。因為算力需求還是一個相對容易解決的問題，畢竟擁有大集群的組織并不只 OpenAI 一家，而如何解決上千塊 GPU 的分布式訓練問題才是關鍵。

根據目前業界已有的分布式訓練方案，即便你是一位非常優秀的數據科學家，知曉并能解決 Transformer 相關的所有算法問題，但如果你不知道如何解決分布式訓練時上百臺服務器之間的通信、拓撲、模型并行、流水并行等問題，你甚至都無法啟動這次訓練。一定程度上，這解釋了GPT-3發布時隔一年，卻只有 NVIDIA 、微軟等大企業可以復現 GPT-3 。

目前，開源的 GPT 模型庫主要是 NVIDIA開發的 Megatron-LM 和經過微軟深度定制開發的 DeepSpeed，其中，DeepSpeed 的模型并行等內核取自 Megatron，它們都是專門為支持 PyTorch 分布式訓練 GPT 而設計。

不過在實際訓練中，PyTorch 、 Megatron、DeepSpeed 都走了一條非常長的彎路。不僅是彎路，你會發現 Megatron 的代碼只能被 NVIDIA 的分布式訓練專家所復用，它對于 PyTorch 的算法工程師而言門檻極高，以至于任何想要用 PyTorch 復現一個分布式大模型的算法工程師，都得先等 NVIDIA 開發完才能再使用 Megatron 提供的模型。

作為新一代深度學習開源框架，致力于“大模型分布式”高效開發的 OneFlow 框架用一套通用設計非常簡單清晰地解決了GPT模型的分布式訓練難題，同時還在已有的測試規模上性能超過 NVIDIA 的 Megatron，這為大規模分布式訓練框架提出了更優的設計理念和路徑。

一、PyTorch 分布式訓練GPT的痛點是什么？

此前，NVIDIA 放出了一篇重量級的論文：Efficient Large-Scale Language Model Training on GPU Clusters ，用 3072 張 80 GB A100 訓練 GPT，最大規模的模型參數量達到了 1T，這是 GPT-3 原版規模的 5 倍。

NVIDIA 訓練 GPT-3 最大到 1T 參數規模

論文里，NVIDIA 介紹了分布式訓練超大規模模型的三種必須的并行技術：

· 數據并行（Data Parallelism）

· 模型并行（Tensor Model Parallelism）

· 流水并行（Pipeline Model Parallelism）

其中，數據并行是最常見的并行方式。而模型并行是對某一層（如 Linear/Dense Layer 里的 Variable ）的模型 Tensor 切分，從而將大的模型 Tensor 分成多個相對較小的 Tensor 進行并行計算；流水并行，是將整個網絡分段（stage），不同段在不同的設備上，前后階段流水分批工作，通過一種“接力”的方式并行。

對于 1T 規模的模型，NVIDIA 一共使用了 384 臺 DGX-A100 機器（每臺裝有 8 張 80GB A100 GPU），機器內部各 GPU 間使用超高速 NVLink 和 NVSwitch 互聯，每臺機器裝有 8 個 200Gbps 的 InfiniBand （IB） 網卡，可以說是硬件集群頂配中的頂配。

那么，這些機器是如何協同工作的？GPT 網絡是由很多層 Transformer Layer 組成，每一層內部是一個由多層 MLP 和 attention 機制組成的子圖，對于參數規模 1T 的 GPT 而言就有 128 層的 Transformer Layer，這個超大超深的網絡被分割成了 64 個 stage （階段），每個 stage 跑在 6 臺 DGX-A100 上，其中 6 臺機器之間進行數據并行，每臺機器內部的 8 張卡之間做模型并行，整個集群的 3072 張 A100 按照機器拓撲被劃分成了 ［6 x 8 x 64］ 的矩陣，同時使用數據并行 & 模型并行 & 流水并行進行訓練。

3072 張 A100 集群拓撲

GPipe、梯度累加、重計算（Checkpointing）和 1F1B（One Forward pass followed by One Backward pass）是分布式訓練 GPT 的流水并行的核心技術。無論是 NVIDIA 的Megatron（PyTorch），還是 OneFlow、PaddlePaddle、MindSpore ，都是通過不同的設計實現了上述相同的功能。

基于 PyTorch 開發的 Megatron，本質上是一個專用于 GPT 的模型庫，所有的代碼都是 Python 腳本，NVIDIA 為 GPT 專門定制了分布式訓練所需的算子、流水并行調度器、模型并行所需的通信原語等功能，在 GPU 上的性能表現上，Megatron已經非常優異。可以說，NVIDIA 在使用 PyTorch 做分布式訓練上已經做到極致了。

但是，用 PyTorch 做分布式訓練，真的好用嗎？

具體來說，從PyTorch 在分布式并行上的設計以及開放給用戶的接口來看，PyTorch 分布式的有以下困境：

· PyTorch 只有物理視角（Physical View），沒有邏輯視角（Logical View）。PyTorch 的用戶想要做分布式并行，任何時候都需要自己推導深度學習模型中哪處需要跟其他的物理設備進行通信和數據同步操作，既要推導通信所在的位置，又要推導通信的操作類型，還要推導跟其他哪些設備通信。這個在簡單的數據并行下可以使用 DDP 或 Horovod 來實現，但是在復雜的模型并行、混合并行下，做并行的門檻非常高。

NVIDIA 模型并行通信推導

· PyTorch 沒有將模型網絡的算法邏輯和分布式并行訓練的通信邏輯解耦出來，導致用戶需要在算子的 kernel 實現中，搭網絡的腳本里到處插入通信原語。這些手寫通信原語的操作不僅繁瑣、易錯、而且沒法復用，是根據特定模型、特定腳本位置、特定算子特判得到的。

· PyTorch 在非對稱的并行方式里（如流水并行，PyTorch 需要人工排線和精細控制流水），各個設備的調度邏輯需要用戶自己手寫。用戶需要自己精細的控制每個設備上的啟動以及執行邏輯，且執行邏輯把前后向執行和send/recv通信操作糅合在一起，即使在最規整的 Transformer Layer 的流水并行下也很復雜，想要擴展到其他模型上的工作量也很大。

模型并行度和流水并行度對性能的影響

· PyTorch 沒有機制保證分布式并行訓練中的正確性和數學一致性。即使用戶寫錯了通信操作，插錯了位置， 跟錯誤的設備進行通信，PyTorch也檢查不出來。

上述困境使得普通算法工程師使用 PyTorch 開發復雜分布式訓練的腳本極為困難。其實，NVIDIA、 微軟、 PyTorch 都被繞進了一個大坑：在沒有一致性視角（ Consistent View ）的情況下做復雜的分布式并行非常困難，往往只能做一些具體網絡、具體場景、具體算子的特判和分析，通過簡單的通信原語來實現分布式。

那么，OneFlow如何解決這些困境？

二、OneFlow 用一致性視角輕松填平分布式訓練難的鴻溝

對于分布式集群環境（多機多卡訓練場景），OneFlow 會把整個分布式集群抽象成一個超級設備，用戶只需要在這個超級設備上搭建深度學習模型即可。這個虛擬出來的超級設備稱之為邏輯視角，而實際上的分布式集群的多機多卡就是物理視角，OneFlow維護邏輯視角和物理視角之間的數學上的正確性就稱之為一致性視角。

基于分布式訓練難的鴻溝，OneFlow通過一致性視角下的 Placement（流水并行） + SBP （數據和模型的混合并行），非常簡單的實現了通用的復雜并行支持。當然，這離不開 OneFlow 的兩大獨特設計：

1. 運行時 Actor 機制

2. 編譯期一致性視角，通過 Placement + SBP + Boxing 解決分布式易用性的問題。

一致性視角（Consistent View）抽象

理想情況下，抽象出來的超級設備（邏輯視角）的算力是所有物理視角下的設備算力之和（如果算力完全用滿，就是線性加速比）；邏輯視角下的顯存資源也是所有物理設備的顯存資源之和。

總體而言，基于一致性視角的 OneFlow 分布式有以下易用性體現：

· OneFlow 的一致性視角將分布式訓練下的多機通信和算法邏輯解耦，使得用戶可以不用關心分布式訓練的細節，降低了分布式訓練的使用門檻。

· 相比其他框架和高級定制用戶在所有分布式并行上的努力，OneFlow 通過 Placement + SBP 機制解決了分布式訓練中任意并行場景的需求。用戶只需要配置 op 的 Placement 就可以完成流水并行，只需要配置 Tensor 的 SBP 就可以實現數據并行、模型并行和混合并行。并且，任何并行方式都是 Placement + SBP 的一種特例， OneFlow 從系統層面不需要做任何的特判，SBP 才是各種分布式并行的本質。

上圖展示了一個 Placement 例子，用于 GPU0 和 GPU1 之間的流水并行。圖中負責在 CPU 和 GPU、GPU 與 GPU 之間進行數據搬運的Op（CopyH2D、CopyD2D）是 OneFlow 系統自動添加的。

· OneFlow 的通信邏輯可以復用，不需要為任何特定網絡和特定算子實現相應的通信邏輯。通信邏輯由 OneFlow 的 Boxing 機制完成，與具體的算子和模型無關。

· OneFlow 的 SBP 還保證了數學上的一致性。 相同的邏輯上的模型腳本，使用任意的并行方式（數據并行、模型并行、流水并行）、使用任意的集群拓撲，OneFlow 都從數學上保證了模型分布式訓練的正確性。

SBP 邏輯與物理 Tensor 的對應關系（SBP 描述了 邏輯上的 Tensor 和 物理上的 Tensor 的映射關系。SBP 全稱叫做 SbpParallel，是三種基礎映射的首字母組合：Split、Broadcast、Partial，其中 Partial 是一個 reduce 操作，包括 PartialSum、PartialMin、PartialMax等）

采用這樣一套簡潔設計可解決分布式并行的各種難題，OneFlow 使得每一位算法工程師都有能力訓練 GPT模型。它讓你不需要成為一位分布式訓練的專家也有能力做復雜的分布式訓練， 只要有硬件資源，任何一位算法工程師都可以訓練 GPT， 都可以開發一個新的大規模分布式訓練的模型。

三、為什么分布式深度學習框架要像 OneFlow 這樣設計？

上述內容從用戶角度分析和比較了 OneFlow 和 PyTorch（Megatron）的分布式易用性，

那么從框架設計和開發者的角度，它又是如何具體實現分布式并行的？為什么說 OneFlow 會是分布式訓練更為本質的設計？

1. OneFlow 如何實現流水并行？

OneFlow 的運行時 Actor 機制有以下幾個特點：

· 天然支持流水線， Actor 通過內部的狀態機和產出的 Regst 個數以及上下游的 Regst 消息機制解決了流控問題（Control Flow）。

Actor 狀態機

· Actor 組成的計算圖運行時調度是去中心化的，每個 Actor 當前是否可以執行都僅與自己的狀態、空閑 Regst 數量以及收到的消息有關。

所以使用 Actor 做流水并行，本身就不需要自己定制復雜的調度邏輯。以數據加載的 Pipeline 為例， 當一個由 Actor 組成的數據預處理流程如下圖所示：

數據預處理流程

當這4個Actor之間的 RegstNum 均為2時，如果訓練時間比較長（訓練是整個網絡的瓶頸），就會得到如下這種流水線的時間線：

數據預處理 pipeline 時間線

在執行幾個 Batch 之后， 4 個階段的執行節奏完全被最長的那個階段所控制，這就是 OneFlow 使用背壓機制（Back Pressure）解決流控問題。

所以流水并行問題，在 OneFlow 中就是 Regst 數量的問題。在實際實現中， OneFlow 采用了一個更通用的算法實現了 Megatron 的流水并行：插入 Buffer Op。在邏輯計算圖上， 會給后向消費前向的邊插入一個 Buffer Op， Buffer 的 Regst 數量 和 Stage 相關。由于后向對前向的消費經過 Checkpointing 優化后，每個 Placement Group 下只會有非常少的幾條消費邊。

OneFlow 通過插入 Buffer Op 實現流水并行

與 Megatron 復雜的手寫調度器和手寫通信原語相比， OneFlow 系統層面只需要插入 Buffer 就可以實現流水并行。

2. OneFlow 如何實現數據+模型的混合并行？

以 Linear Layer 的數據 + 模型并行為例，來解釋所有的數據并行和模型并行的組合，本質上都是被 SBP 所描述的 Signature。任何并行方式的設備間通信操作，該在整個網絡的哪里插入、該插入什么通信操作、每個設備該和誰通信，完全都是 SBP 自動推導得到的，而且還保證數學上的一致性。

可以說，OneFlow的設計使得算法工程師告別了分布式并行中的通信原語。不僅如此，OneFlow 的框架開發者絕大多數時候也不需要關心分布式里的通信原語，SBP 這層抽象使得算子/網絡跟分布式通信解耦。

以 1-D SBP 為例，1-D SBP 下的數據并行，對于一個 Linear Layer 而言，主要是其中的 MatMul（矩陣乘法）計算。假設矩陣乘法計算在邏輯視角上是一個 （m， k） x （k， n） = （m， n） 的計算，m 表示一共有多少個樣例， k 和 n 分別是 Linear Layer 中的隱藏層神經元數量以及輸出神經元數量。

數據并行的邏輯計算圖 -》 物理計算圖 的映射關系如下圖所示：

數據并行下邏輯計算圖轉物理計算圖

數據并行下，每個設備上都有全部的模型（Tensor b， Shape = （k， n）），假設共有兩張卡，則 GPU0 上有前一半的數據 （Tensor a，Shape = （m/2， k）），GPU1 上有后一半的數據， 則Tensor a 的 SBP Parallel = Split（0）。同時可以看到矩陣乘的輸出 Tensor out，也是按照第 0 維切分的。

模型并行對于 Linear Layer 而言，有兩種，分別是切模型 Tensor 的第0維（行切分，對應 Megatron 里的 RowParallelLinear）和 第1維（列切分，對應 Megatron 里的 ColumnParallelLinear）。

第一種行切分（RowParallelLinear）模型并行的 邏輯計算圖 -》 物理計算圖 的映射關系如下圖所示：

模型并行（行切分） 邏輯圖轉物理圖

模型并行下，每個設備都只有一部分的模型，在這個例子中， GPU 0 上有前一半的模型， GPU 1上有后一半的模型，每個設備上的模型大小 Tensor b 的 Shape = （k/2， n）。在這種情況下， 每個設備輸出的 Tensor out 都是完整的數據大小， Shape = （m， n）， 但每個位置上的元素的值，都是邏輯上的輸出 out 對應位置的值的一部分，即 out 的 SBP Parallel = PartialSum 。

第二種列切分（ColumnParallelLinear）模型并行的 邏輯計算圖 -》 物理計算圖 的映射關系如下圖所示：

模型并行（列切分）邏輯圖轉物理圖

這個例子中，模型 Tensor b 是按照 Split（1） 切分的，輸出 Tensor out 也是按照 Split（1） 切分的，每個設備都需要全部的數據。

在 GPT 網絡中，實際上的模型并行是組合使用 RowParallelLinear 和 ColumnParallelLinear 實現的（ColumnParallelLinear 后面接了 RowParallelLinear）。

因為 Column 的輸出 Tensor SBP 是 Split（1）， Row 的輸入數據 Tensor SBP 也是 Split（1）， 所以當 Column 后接 Row 時，兩者之間是不需要插入任何通信的。但由于 Row 的輸出是 PartialSum， 當后面消費該 Tensor （在網絡中是 Add 操作）的 Op 需要全部的數據時（Broadcast）， 此處就需要插入 AllReduce 實現通信了。

這在 OneFlow 中稱之為 Boxing。 當兩個邏輯上的 Op 對于同一個邏輯上的 Tensor 看待的 SBP Parallel 不一致時， OneFlow 系統會自動插入通信節點以完成數據的切分/傳輸/拼接等操作，使得下游 Op 總能拿到按照自己期望 SBP 切分的 Tensor。

?

Boxing：通過 AllGather 實現 Split（1） 轉 Broadcast

在 OneFlow 中， 所有的分布式通信操作都是基于 SBP 的推導結果，按照需要插入。OneFlow 通過 Boxing 機制，就實現了任意的數據并行和模型并行。

2-D SBP 其實就是將兩組 1-D SBP 按照設備拓撲的維度拼起來就可以得到。其實 GPT 中用到的 2-D SBP 只是最簡單情形的特例， 分布式下的并行經過 2-D SBP 可以拓展出非常多復雜、靈活多邊的組合出來。而針對復雜的組合， 再想用 Megatron 的設計就非常難做，但對于 OneFlow 而言，二者的難度是一樣的，因為本質上是用 Boxing 完成一組 2-D SBP 的變換。

四、GPT 分布式訓練性能對比：OneFlow vs Megatron

與 Megatron 相比，OneFlow 除了在用戶接口（分布式易用性） 和框架設計上更簡潔、更易用，在 4 機 32卡 16GB V100 的測試規模上性能也超過 Megatron。值得一提的是，經過 NVIDIA 的深度優化， Megatron 在 GPU 上的分布式訓練性能已經接近極致，DeepSpeed 也無法與之相比。

以下的所有實驗數據均在相同的硬件環境、相同的第三方依賴（CUDA、 cuDNN等）、使用相同的參數和網絡結構下， 對比了 OneFlow 和 Megatron 在 GPT 模型下的性能表現。所有的性能結果均公開且可復現。（GPT 模型腳本在Oneflow-Inc/OneFlow-Benchmark 倉庫， 公開的評測報告、復現方式稍后在Oneflow-Inc/DLPerf 倉庫中可查看。）

數據并行性能對比

注：每組參數的縮略版含義：

· DP 數據并行；MP 模型并行；2D 數據 & 模型 的 混合并行；PP 流水并行

· dxmxp_B_hxl 其中：

· d = 數據并行度（data-parallel-size）

· m = 模型并行度（tensor-model-parallel-size）

· p = 流水并行度（pipeline-model-parallel-size）

· B = 總的BatchSize（global-batch-size）

· h = 隱藏層大小（hidden-size）影響每層 Transformer Layer 的模型大小

· l = Transformer Layer 層數（num-layers）

模型并行數據對比

注：由于單卡 GPU 顯存限制，各組參數里的模型大小是不同的，所以整體不像數據并行那樣呈線性增加的關系。如第 4 組參數（MP_1x32x1_16_3072x32）的模型大小是第 2 組參數（MP_1x8x1_16_1536x16）的 8 倍以上。NVIDIA 論文中有模型規模跟各個參數的計算公式：

其中 l 表示 num-layers ，h 表示 hidden-size， V 表示詞表大小（vocabulary size = 51200）， S 表示句子長度（a sequence length = 2048）， P 表示參數規模。

數據+模型+流水并行性能對比（注：第 1 組參數的模型比后 3 組的都要小，因為機器內的數據并行限制了參數規模。）

五、小結

在分布式訓練領域擁有獨特的設計和視角，OneFlow 解決了分布式訓練中的各種并行難題，因此在大規模預訓練模型場景下用 OneFlow 做分布式訓練更易用也更高效。

同時，OneFlow 團隊正在全力提升框架的單卡使用體驗。據悉，OneFlow 即將在 5 月發布的大版本 OneFlow v0.4.0 起，將提供兼容 PyTorch 的全新接口以及動態圖等特性。而在 v0.5.0 版本，OneFlow 預計全面兼容 PyTorch， 屆時用戶可將 PyTorch 的模型訓練腳本一鍵遷移為 OneFlow 的訓練腳本。此外， OneFlow 還會提供 Consistent 視角的分布式 Eager，用戶可以既享受動態圖的易用性，又可以非常方便的進行各種分布式并行訓練。