GPT-3難以復現(xiàn)，為什么說PyTorch走上了一條“大彎路”？
2020 年，最轟動的 AI 新聞莫過于 OpenAI 發(fā)布的 GPT-3 了。它的 1750 億參數(shù)量及其在眾多 NLP 任務上超過人類的出眾表現(xiàn)讓人們開始堅信：大模型才是未來。但與之帶來的問題是，訓練超大模型所需的算力、存儲已不再是單機就能搞定。
據(jù) NVIDIA 估算，如果要訓練 GPT-3 ，即使單個機器的顯存 / 內(nèi)存能裝得下，用 8 張 V100 的顯卡，訓練時長預計要 36 年；即使用 512 張 V100 ，訓練也需要將近 7 個月；如果擁有 1024 張 80GB A100， 那么完整訓練 GPT-3 的時長可以縮減到 1 個月。
除去硬件資源這個經(jīng)濟問題，在技術層面，意味著訓練大模型一定是一個分布式問題。因為算力需求還是一個相對容易解決的問題，畢竟擁有大集群的組織并不只 OpenAI 一家，而如何解決上千塊 GPU 的分布式訓練問題才是關鍵。
根據(jù)目前業(yè)界已有的分布式訓練方案，即便你是一位非常優(yōu)秀的數(shù)據(jù)科學家，知曉并能解決 Transformer 相關的所有算法問題，但如果你不知道如何解決分布式訓練時上百臺服務器之間的通信、拓撲、模型并行、流水并行等問題，你甚至都無法啟動這次訓練。一定程度上，這解釋了 GPT-3 發(fā)布時隔一年，卻只有 NVIDIA 、微軟等大企業(yè)可以復現(xiàn) GPT-3 。
目前，開源的 GPT 模型庫主要是 NVIDIA 開發(fā)的 Megatron-LM 和經(jīng)過微軟深度定制開發(fā)的 DeepSpeed，其中，DeepSpeed 的模型并行等內(nèi)核取自 Megatron，它們都是專門為支持 PyTorch 分布式訓練 GPT 而設計。
不過在實際訓練中，PyTorch 、 Megatron、DeepSpeed 都走了一條非常長的彎路。不僅是彎路，你會發(fā)現(xiàn) Megatron 的代碼只能被 NVIDIA 的分布式訓練專家所復用，它對于 PyTorch 的算法工程師而言門檻極高，以至于任何想要用 PyTorch 復現(xiàn)一個分布式大模型的算法工程師，都得先等 NVIDIA 開發(fā)完才能再使用 Megatron 提供的模型。
作為新一代深度學習開源框架，致力于“大模型分布式”高效開發(fā)的 OneFlow 框架用一套通用設計非常簡單清晰地解決了 GPT 模型的分布式訓練難題，同時還在已有的測試規(guī)模上性能超過 NVIDIA 的 Megatron，這為大規(guī)模分布式訓練框架提出了更優(yōu)的設計理念和路徑。
一、PyTorch 分布式訓練 GPT 的痛點是什么？
此前，NVIDIA 放出了一篇重量級的論文：Efficient Large-Scale Language Model Training on GPU Clusters ，用 3072 張 80 GB A100 訓練 GPT，最大規(guī)模的模型參數(shù)量達到了 1T，這是 GPT-3 原版規(guī)模的 5 倍。
NVIDIA 訓練 GPT-3 最大到 1T 參數(shù)規(guī)模
論文里，NVIDIA 介紹了分布式訓練超大規(guī)模模型的三種必須的并行技術：
數(shù)據(jù)并行（Data Parallelism）
模型并行（Tensor Model Parallelism）
流水并行（Pipeline Model Parallelism）
其中，數(shù)據(jù)并行是最常見的并行方式。而模型并行是對某一層（如 Linear/Dense Layer 里的 Variable ）的模型 Tensor 切分，從而將大的模型 Tensor 分成多個相對較小的 Tensor 進行并行計算；流水并行，是將整個網(wǎng)絡分段（stage），不同段在不同的設備上，前后階段流水分批工作，通過一種“接力”的方式并行。
對于 1T 規(guī)模的模型，NVIDIA 一共使用了 384 臺 DGX-A100 機器（每臺裝有 8 張 80GB A100 GPU），機器內(nèi)部各 GPU 間使用超高速 NVLink 和 NVSwitch 互聯(lián)，每臺機器裝有 8 個 200Gbps 的 InfiniBand (IB) 網(wǎng)卡，可以說是硬件集群頂配中的頂配。
那么，這些機器是如何協(xié)同工作的？GPT 網(wǎng)絡是由很多層 Transformer Layer 組成，每一層內(nèi)部是一個由多層 MLP 和 attention 機制組成的子圖，對于參數(shù)規(guī)模 1T 的 GPT 而言就有 128 層的 Transformer Layer，這個超大超深的網(wǎng)絡被分割成了 64 個 stage （階段），每個 stage 跑在 6 臺 DGX-A100 上，其中 6 臺機器之間進行數(shù)據(jù)并行，每臺機器內(nèi)部的 8 張卡之間做模型并行，整個集群的 3072 張 A100 按照機器拓撲被劃分成了 [6 x 8 x 64] 的矩陣，同時使用數(shù)據(jù)并行 & 模型并行 & 流水并行進行訓練。
3072 張 A100 集群拓撲
GPipe、梯度累加、重計算（Checkpointing）和 1F1B（One Forward pass followed by One Backward pass）是分布式訓練 GPT 的流水并行的核心技術。無論是 NVIDIA 的 Megatron（PyTorch），還是 OneFlow、PaddlePaddle、MindSpore ，都是通過不同的設計實現(xiàn)了上述相同的功能。
基于 PyTorch 開發(fā)的 Megatron，本質(zhì)上是一個專用于 GPT 的模型庫，所有的代碼都是 Python 腳本，NVIDIA 為 GPT 專門定制了分布式訓練所需的算子、流水并行調(diào)度器、模型并行所需的通信原語等功能，在 GPU 上的性能表現(xiàn)上，Megatron 已經(jīng)非常優(yōu)異?？梢哉f，NVIDIA 在使用 PyTorch 做分布式訓練上已經(jīng)做到極致了。
但是，用 PyTorch 做分布式訓練，真的好用嗎？
具體來說，從 PyTorch 在分布式并行上的設計以及開放給用戶的接口來看，PyTorch 分布式的有以下困境：
PyTorch 只有物理視角（Physical View），沒有邏輯視角（Logical View）。
PyTorch 的用戶想要做分布式并行，任何時候都需要自己推導深度學習模型中哪處需要跟其他的物理設備進行通信和數(shù)據(jù)同步操作，既要推導通信所在的位置，又要推導通信的操作類型，還要推導跟其他哪些設備通信。這個在簡單的數(shù)據(jù)并行下可以使用 DDP 或 Horovod 來實現(xiàn)，但是在復雜的模型并行、混合并行下，做并行的門檻非常高。
NVIDIA 模型并行通信推導
PyTorch 沒有將模型網(wǎng)絡的算法邏輯和分布式并行訓練的通信邏輯解耦出來，導致用戶需要在算子的 kernel 實現(xiàn)中，搭網(wǎng)絡的腳本里到處插入通信原語。這些手寫通信原語的操作不僅繁瑣、易錯、而且沒法復用，是根據(jù)特定模型、特定腳本位置、特定算子特判得到的。
PyTorch 在非對稱的并行方式里（如流水并行，PyTorch 需要人工排線和精細控制流水），各個設備的調(diào)度邏輯需要用戶自己手寫。用戶需要自己精細的控制每個設備上的啟動以及執(zhí)行邏輯，且執(zhí)行邏輯把前后向執(zhí)行和 send/recv 通信操作糅合在一起，即使在最規(guī)整的 Transformer Layer 的流水并行下也很復雜，想要擴展到其他模型上的工作量也很大。
模型并行度和流水并行度對性能的影響
PyTorch 沒有機制保證分布式并行訓練中的正確性和數(shù)學一致性。即使用戶寫錯了通信操作，插錯了位置， 跟錯誤的設備進行通信，PyTorch 也檢查不出來。
上述困境使得普通算法工程師使用 PyTorch 開發(fā)復雜分布式訓練的腳本極為困難。其實，NVIDIA、 微軟、 PyTorch 都被繞進了一個大坑：在沒有一致性視角（ Consistent View ）的情況下做復雜的分布式并行非常困難，往往只能做一些具體網(wǎng)絡、具體場景、具體算子的特判和分析，通過簡單的通信原語來實現(xiàn)分布式。
那么，OneFlow 如何解決這些困境？
二、OneFlow 用一致性視角輕松填平分布式訓練難的鴻溝
對于分布式集群環(huán)境（多機多卡訓練場景），OneFlow 會把整個分布式集群抽象成一個超級設備，用戶只需要在這個超級設備上搭建深度學習模型即可。這個虛擬出來的超級設備稱之為邏輯視角，而實際上的分布式集群的多機多卡就是物理視角，OneFlow 維護邏輯視角和物理視角之間的數(shù)學上的正確性就稱之為一致性視角。
基于分布式訓練難的鴻溝，OneFlow 通過一致性視角下的 Placement（流水并行） + SBP （數(shù)據(jù)和模型的混合并行），非常簡單的實現(xiàn)了通用的復雜并行支持。當然，這離不開 OneFlow 的兩大獨特設計：
運行時 Actor 機制
編譯期一致性視角，通過 Placement + SBP + Boxing 解決分布式易用性的問題。
一致性視角（Consistent View）抽象
理想情況下，抽象出來的超級設備（邏輯視角）的算力是所有物理視角下的設備算力之和（如果算力完全用滿，就是線性加速比）；邏輯視角下的顯存資源也是所有物理設備的顯存資源之和。
總體而言，基于一致性視角的 OneFlow 分布式有以下易用性體現(xiàn)：
OneFlow 的一致性視角將分布式訓練下的多機通信和算法邏輯解耦，使得用戶可以不用關心分布式訓練的細節(jié)，降低了分布式訓練的使用門檻。
相比其他框架和高級定制用戶在所有分布式并行上的努力，OneFlow 通過 Placement + SBP 機制解決了分布式訓練中任意并行場景的需求。用戶只需要配置 op 的 Placement 就可以完成流水并行，只需要配置 Tensor 的 SBP 就可以實現(xiàn)數(shù)據(jù)并行、模型并行和混合并行。并且，任何并行方式都是 Placement + SBP 的一種特例， OneFlow 從系統(tǒng)層面不需要做任何的特判，SBP 才是各種分布式并行的本質(zhì)。
上圖展示了一個 Placement 例子，用于 GPU0 和 GPU1 之間的流水并行。圖中負責在 CPU 和 GPU、GPU 與 GPU 之間進行數(shù)據(jù)搬運的 Op（CopyH2D、CopyD2D）是 OneFlow 系統(tǒng)自動添加的。
OneFlow 的通信邏輯可以復用，不需要為任何特定網(wǎng)絡和特定算子實現(xiàn)相應的通信邏輯。通信邏輯由 OneFlow 的 Boxing 機制完成，與具體的算子和模型無關。
OneFlow 的 SBP 還保證了數(shù)學上的一致性。相同的邏輯上的模型腳本，使用任意的并行方式（數(shù)據(jù)并行、模型并行、流水并行）、使用任意的集群拓撲，OneFlow 都從數(shù)學上保證了模型分布式訓練的正確性。
SBP 邏輯與物理 Tensor 的對應關系（SBP 描述了 邏輯上的 Tensor 和 物理上的 Tensor 的映射關系。SBP 全稱叫做 SbpParallel，是三種基礎映射的首字母組合：Split、Broadcast、Partial，其中 Partial 是一個 reduce 操作，包括 PartialSum、PartialMin、PartialMax 等）
采用這樣一套簡潔設計可解決分布式并行的各種難題，OneFlow 使得每一位算法工程師都有能力訓練 GPT 模型。它讓你不需要成為一位分布式訓練的專家也有能力做復雜的分布式訓練， 只要有硬件資源，任何一位算法工程師都可以訓練 GPT， 都可以開發(fā)一個新的大規(guī)模分布式訓練的模型。
三、為什么分布式深度學習框架要像 OneFlow 這樣設計？
上述內(nèi)容從用戶角度分析和比較了 OneFlow 和 PyTorch（Megatron）的分布式易用性，那么從框架設計和開發(fā)者的角度，它又是如何具體實現(xiàn)分布式并行的？為什么說 OneFlow 會是分布式訓練更為本質(zhì)的設計？
OneFlow 如何實現(xiàn)流水并行？
OneFlow 的運行時 Actor 機制有以下幾個特點：
天然支持流水線， Actor 通過內(nèi)部的狀態(tài)機和產(chǎn)出的 Regst 個數(shù)以及上下游的 Regst 消息機制解決了流控問題（Control Flow）。
Actor 狀態(tài)機
Actor 組成的計算圖運行時調(diào)度是去中心化的，每個 Actor 當前是否可以執(zhí)行都僅與自己的狀態(tài)、空閑 Regst 數(shù)量以及收到的消息有關。
所以使用 Actor 做流水并行，本身就不需要自己定制復雜的調(diào)度邏輯。以數(shù)據(jù)加載的 Pipeline 為例， 當一個由 Actor 組成的數(shù)據(jù)預處理流程如下圖所示：
數(shù)據(jù)預處理流程
當這 4 個 Actor 之間的 RegstNum 均為 2 時，如果訓練時間比較長（訓練是整個網(wǎng)絡的瓶頸），就會得到如下這種流水線的時間線：
數(shù)據(jù)預處理 pipeline 時間線
在執(zhí)行幾個 Batch 之后， 4 個階段的執(zhí)行節(jié)奏完全被最長的那個階段所控制，這就是 OneFlow 使用背壓機制（Back Pressure）解決流控問題。
所以流水并行問題，在 OneFlow 中就是 Regst 數(shù)量的問題。在實際實現(xiàn)中， OneFlow 采用了一個更通用的算法實現(xiàn)了 Megatron 的流水并行：插入 Buffer Op。在邏輯計算圖上， 會給后向消費前向的邊插入一個 Buffer Op， Buffer 的 Regst 數(shù)量 和 Stage 相關。由于后向?qū)η跋虻南M經(jīng)過 Checkpointing 優(yōu)化后，每個 Placement Group 下只會有非常少的幾條消費邊。
OneFlow 通過插入 Buffer Op 實現(xiàn)流水并行
與 Megatron 復雜的手寫調(diào)度器和手寫通信原語相比， OneFlow 系統(tǒng)層面只需要插入 Buffer 就可以實現(xiàn)流水并行。
OneFlow 如何實現(xiàn)數(shù)據(jù) + 模型的混合并行？
以 Linear Layer 的數(shù)據(jù) + 模型并行為例，來解釋所有的數(shù)據(jù)并行和模型并行的組合，本質(zhì)上都是被 SBP 所描述的 Signature。任何并行方式的設備間通信操作，該在整個網(wǎng)絡的哪里插入、該插入什么通信操作、每個設備該和誰通信，完全都是 SBP 自動推導得到的，而且還保證數(shù)學上的一致性。
可以說，OneFlow 的設計使得算法工程師告別了分布式并行中的通信原語。不僅如此，OneFlow 的框架開發(fā)者絕大多數(shù)時候也不需要關心分布式里的通信原語，SBP 這層抽象使得算子 / 網(wǎng)絡跟分布式通信解耦。
以 1-D SBP 為例，1-D SBP 下的數(shù)據(jù)并行，對于一個 Linear Layer 而言，主要是其中的 MatMul（矩陣乘法）計算。假設矩陣乘法計算在邏輯視角上是一個 (m, k) x (k, n) = (m, n) 的計算，m 表示一共有多少個樣例， k 和 n 分別是 Linear Layer 中的隱藏層神經(jīng)元數(shù)量以及輸出神經(jīng)元數(shù)量。
數(shù)據(jù)并行的邏輯計算圖 -> 物理計算圖 的映射關系如下圖所示：
數(shù)據(jù)并行下邏輯計算圖轉(zhuǎn)物理計算圖
數(shù)據(jù)并行下，每個設備上都有全部的模型（Tensor b， Shape = (k, n)），假設共有兩張卡，則 GPU0 上有前一半的數(shù)據(jù) （Tensor a，Shape = (m/2, k)），GPU1 上有后一半的數(shù)據(jù)， 則 Tensor a 的 SBP Parallel = Split(0)。同時可以看到矩陣乘的輸出 Tensor out，也是按照第 0 維切分的。
模型并行對于 Linear Layer 而言，有兩種，分別是切模型 Tensor 的第 0 維（行切分，對應 Megatron 里的 RowParallelLinear）和 第 1 維（列切分，對應 Megatron 里的 ColumnParallelLinear）。
第一種行切分（RowParallelLinear）模型并行的 邏輯計算圖 -> 物理計算圖 的映射關系如下圖所示：
模型并行（行切分） 邏輯圖轉(zhuǎn)物理圖
模型并行下，每個設備都只有一部分的模型，在這個例子中， GPU 0 上有前一半的模型， GPU 1 上有后一半的模型，每個設備上的模型大小 Tensor b 的 Shape = (k/2, n)。在這種情況下， 每個設備輸出的 Tensor out 都是完整的數(shù)據(jù)大小， Shape = (m, n)， 但每個位置上的元素的值，都是邏輯上的輸出 out 對應位置的值的一部分，即 out 的 SBP Parallel = PartialSum 。
第二種列切分（ColumnParallelLinear）模型并行的 邏輯計算圖 -> 物理計算圖 的映射關系如下圖所示：
模型并行（列切分）邏輯圖轉(zhuǎn)物理圖
這個例子中，模型 Tensor b 是按照 Split(1) 切分的，輸出 Tensor out 也是按照 Split(1) 切分的，每個設備都需要全部的數(shù)據(jù)。
在 GPT 網(wǎng)絡中，實際上的模型并行是組合使用 RowParallelLinear 和 ColumnParallelLinear 實現(xiàn)的（ColumnParallelLinear 后面接了 RowParallelLinear）。
因為 Column 的輸出 Tensor SBP 是 Split(1)， Row 的輸入數(shù)據(jù) Tensor SBP 也是 Split(1)， 所以當 Column 后接 Row 時，兩者之間是不需要插入任何通信的。但由于 Row 的輸出是 PartialSum， 當后面消費該 Tensor （在網(wǎng)絡中是 Add 操作）的 Op 需要全部的數(shù)據(jù)時（Broadcast）， 此處就需要插入 AllReduce 實現(xiàn)通信了。
這在 OneFlow 中稱之為 Boxing。當兩個邏輯上的 Op 對于同一個邏輯上的 Tensor 看待的 SBP Parallel 不一致時， OneFlow 系統(tǒng)會自動插入通信節(jié)點以完成數(shù)據(jù)的切分 / 傳輸 / 拼接等操作，使得下游 Op 總能拿到按照自己期望 SBP 切分的 Tensor。
Boxing：通過 AllReduce 實現(xiàn) PartialSum 轉(zhuǎn) Broadcast
Boxing：通過 AllGather 實現(xiàn) Split(1) 轉(zhuǎn) Broadcast
在 OneFlow 中， 所有的分布式通信操作都是基于 SBP 的推導結(jié)果，按照需要插入。OneFlow 通過 Boxing 機制，就實現(xiàn)了任意的數(shù)據(jù)并行和模型并行。
2-D SBP 其實就是將兩組 1-D SBP 按照設備拓撲的維度拼起來就可以得到。其實 GPT 中用到的 2-D SBP 只是最簡單情形的特例， 分布式下的并行經(jīng)過 2-D SBP 可以拓展出非常多復雜、靈活多邊的組合出來。而針對復雜的組合， 再想用 Megatron 的設計就非常難做，但對于 OneFlow 而言，二者的難度是一樣的，因為本質(zhì)上是用 Boxing 完成一組 2-D SBP 的變換。
四、GPT 分布式訓練性能對比：OneFlow vs Megatron
與 Megatron 相比，OneFlow 除了在用戶接口（分布式易用性） 和框架設計上更簡潔、更易用，在 4 機 32 卡 16GB V100 的測試規(guī)模上性能也超過 Megatron。值得一提的是，經(jīng)過 NVIDIA 的深度優(yōu)化， Megatron 在 GPU 上的分布式訓練性能已經(jīng)接近極致，DeepSpeed 也無法與之相比。
以下的所有實驗數(shù)據(jù)均在相同的硬件環(huán)境、相同的第三方依賴（CUDA、 cuDNN 等）、使用相同的參數(shù)和網(wǎng)絡結(jié)構(gòu)下， 對比了 OneFlow 和 Megatron 在 GPT 模型下的性能表現(xiàn)。所有的性能結(jié)果均公開且可復現(xiàn)。（GPT 模型腳本在 Oneflow-Inc/OneFlow-Benchmark 倉庫， 公開的評測報告、復現(xiàn)方式稍后在 Oneflow-Inc/DLPerf 倉庫中可查看。)
數(shù)據(jù)并行性能對比
注：每組參數(shù)的縮略版含義：· DP 數(shù)據(jù)并行；MP 模型并行；2D 數(shù)據(jù) & 模型 的 混合并行；PP 流水并行· dxmxp_B_hxl 其中：· d = 數(shù)據(jù)并行度（data-parallel-size）· m = 模型并行度（tensor-model-parallel-size）· p = 流水并行度（pipeline-model-parallel-size）· B = 總的 BatchSize（global-batch-size）· h = 隱藏層大?。╤idden-size）影響每層 Transformer Layer 的模型大小· l = Transformer Layer 層數(shù)（num-layers）
模型并行數(shù)據(jù)對比
注：由于單卡 GPU 顯存限制，各組參數(shù)里的模型大小是不同的，所以整體不像數(shù)據(jù)并行那樣呈線性增加的關系。如第 4 組參數(shù)（MP_1x32x1_16_3072x32）的模型大小是第 2 組參數(shù)（MP_1x8x1_16_1536x16）的 8 倍以上。NVIDIA 論文中有模型規(guī)模跟各個參數(shù)的計算公式：
其中 l 表示 num-layers ，h 表示 hidden-size， V 表示詞表大?。╲ocabulary size = 51200）， S 表示句子長度（a sequence length = 2048）， P 表示參數(shù)規(guī)模。
數(shù)據(jù) + 模型并行性能對比（注：其中前 4 組的模型規(guī)模一致；后 2 組的模型規(guī)模一致。）
數(shù)據(jù) + 模型 + 流水并行性能對比（注：第 1 組參數(shù)的模型比后 3 組的都要小，因為機器內(nèi)的數(shù)據(jù)并行限制了參數(shù)規(guī)模。）
五、小結(jié)
在分布式訓練領域擁有獨特的設計和視角，OneFlow 解決了分布式訓練中的各種并行難題，因此在大規(guī)模預訓練模型場景下用 OneFlow 做分布式訓練更易用也更高效。
同時，OneFlow 團隊正在全力提升框架的單卡使用體驗。據(jù)悉，OneFlow 即將在 5 月發(fā)布的大版本 OneFlow v0.4.0 起，將提供兼容 PyTorch 的全新接口以及動態(tài)圖等特性。而在 v0.5.0 版本，OneFlow 預計全面兼容 PyTorch， 屆時用戶可將 PyTorch 的模型訓練腳本一鍵遷移為 OneFlow 的訓練腳本。此外， OneFlow 還會提供 Consistent 視角的分布式 Eager，用戶可以既享受動態(tài)圖的易用性，又可以非常方便的進行各種分布式并行訓練。
OneFlow GitHub 鏈接：https://github.com/Oneflow-Inc/oneflow
內(nèi)容詳版請查看：https://mp.weixin.qq.com/s/YQshbDYgfYCO9POH7Ls3cw
你也「在看」嗎？??