GPT-3模型為何難以復(fù)現(xiàn)？這也許是分布式AI框架的最優(yōu)設(shè)計（4）

為什么分布式深度學習框架要像 OneFlow 這樣設(shè)計？

上一個章節(jié)，我們從用戶角度分析和比較了 OneFlow 和 PyTorch（Megatron）的分布式易用性。這一章節(jié)我們會從框架設(shè)計和開發(fā)者的角度解釋，為什么我們這一套 Placement + SBP 的設(shè)計是分布式訓練更本質(zhì)的設(shè)計。其他框架和高級定制用戶的在所有分布式并行上的努力，其實都是 SBP 下的一個特例而已。

1. OneFlow 如何實現(xiàn)流水并行？

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

天然支持流水線， Actor 通過內(nèi)部的狀態(tài)機 和 產(chǎn)出的 Regst 個數(shù) 以及上下游的 Regst 消息機制解決了流控問題（Control Flow）。Actor 的狀態(tài)機 如下圖所示：

Actor 狀態(tài)機

Actor 組成的計算圖運行時調(diào)度是去中心化的，每個 Actor 當前是否可以執(zhí)行都僅與自己的狀態(tài)、空閑 Regst 數(shù)量以及收到的消息有關(guān)。

所以使用 Actor 做流水并行，本身就不需要自己定制復(fù)雜的調(diào)度邏輯。我們可以先舉一個數(shù)據(jù)加載的 Pipeline 示例， 當一個由 Actor 組成的數(shù)據(jù)預(yù)處理流程如下圖所示時（我們可以將各個階段約減為一個 Actor）：

數(shù)據(jù)預(yù)處理流程

當這4個Actor之間的 RegstNum 均為2時，如果訓練時間比較長（訓練是整個網(wǎng)絡(luò)的瓶頸），我們就會得到下面這種流水線的時間線：

數(shù)據(jù)預(yù)處理 pipeline 時間線

在執(zhí)行幾個 Batch 之后， 4 個階段的執(zhí)行節(jié)奏完全被最長的那個階段所控制。這就是 OneFlow 使用背壓機制（Back Pressure）解決流控問題。（比如 Preprocessing 階段， 該 Actor 是否執(zhí)行不僅僅跟有沒有數(shù)據(jù)相關(guān)，也要考慮自己有沒有“空閑”的 Regst 塊可寫）

所以流水并行問題，在 OneFlow 中就是 Regst 數(shù)量的問題。在實際實現(xiàn)中， OneFlow 采用了一個更通用的算法實現(xiàn)了 Megatron 的流水并行：插入 Buffer Op。在邏輯計算圖上， 我們會給后向消費前向的邊插入一個 Buffer Op， Buffer 的 Regst 數(shù)量 和 Stage 相關(guān)。由于后向?qū)η跋虻南M經(jīng)過 Checkpointing 優(yōu)化后，每個 Placement Group 下只會有非常少的幾條消費邊。整體的算法實現(xiàn)可以通過下面這個示意圖來解釋：

OneFlow 通過插入 Buffer Op 實現(xiàn)流水并行

假設(shè)整個網(wǎng)絡(luò)分為 4 個 stage， 共有 8 個 Transformer Layer， 則我們需要在前 3 個 （stage_num - 1）stage 的前后向計算圖中插入 Buffer Op。最后一個 stage 由于每做完一個 micro-batch 的前向，立馬做該 micro-batch 的反向，則不需要插入 Buffer。buffer 的 regst_num 跟 stage_num 相關(guān)。（圖中是理想情況下，假設(shè) stage 之間的傳輸開銷可以忽略不計，則至少需要 stage_num - 1 的 buffer_size）由于我們對每一個 Transformer Layer 做了 Checkpointing，則每個 Layer 僅有一條前向到后向的數(shù)據(jù)邊， 則只需要插入一個 Buffer Op。

跟 Megatron 復(fù)雜的手寫調(diào)度器 和 手寫通信原語相比， OneFlow 系統(tǒng)層面只需要插入 Buffer 就可以實現(xiàn)流水并行。

2. OneFlow 如何實現(xiàn)數(shù)據(jù)+模型的混合并行？

我們以 Linear Layer 的數(shù)據(jù) + 模型并行為例，來解釋所有的數(shù)據(jù)并行和模型并行 的組合，本質(zhì)上都是被 SBP 所描述的 Signature 而已。任何并行方式的設(shè)備間通信操作，該在整個網(wǎng)絡(luò)的哪里插入、該插入什么通信操作、每個設(shè)備該和誰通信，完全都是 SBP 自動推導得到的，而且還保證數(shù)學上的一致性。有了 OneFlow， 算法工程師就告別了分布式并行中的通信原語了。不僅如此，OneFlow 的框架開發(fā)者絕大多數(shù)時候也不需要關(guān)心分布式里的通信原語，SBP 這層抽象使得算子/網(wǎng)絡(luò)跟分布式通信解耦。

我們先以 1-D SBP 為例，之后再擴展到 2-D SBP。1-D SBP 下的數(shù)據(jù)并行，對于一個 Linear Layer 而言，主要是其中的 MatMul（矩陣乘法）計算。我們假設(shè)矩陣乘法計算在邏輯視角上 是一個 (m, k) x (k, n) = (m, n) 的計算，m 表示一共有多少個樣例， k 和 n 分別是 Linear Layer 中的隱藏層神經(jīng)元數(shù)量 以及 輸出神經(jīng)元數(shù)量。

數(shù)據(jù)并行的 邏輯計算圖 -> 物理計算圖 的映射關(guān)系如下圖所示：

數(shù)據(jù)并行下邏輯計算圖轉(zhuǎn)物理計算圖

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

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

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

模型并行（行切分） 邏輯圖轉(zhuǎn)物理圖

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

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

模型并行（列切分）邏輯圖轉(zhuǎn)物理圖

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

在 GPT 網(wǎng)絡(luò)中，實際上的模型并行是組合使用 RowParallelLinear 和 ColumnParallelLinear 實現(xiàn)的（ColumnParallelLinear 后面接了 RowParallelLinear）。

注：這里我沒有列出整個 Transformer Layer 的 SBP 推導信息，只說了“交替使用 Row 和 Linear，插入 AllReduce ”。后續(xù)我會補充實際 GPT 的網(wǎng)絡(luò)結(jié)構(gòu) + 模型并行 SBP 信息的圖和說明。

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

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

下面展示兩個 Boxing 的示例。第一個是 PartialSum -> Broadcast ，為了得到完整的數(shù)據(jù)我們需要將 PartialSum 的 Tensor 對應(yīng)位置加起來，這時候 OneFlow 會自動插入 AllReduce 操作（這里類比 Megatron 在模型腳本里手寫 AllReduce 通信原語）。

Boxing：通過 AllReduce 實現(xiàn) PartialSum 轉(zhuǎn) Broadcast

第二個是 Split(1) -> Broadcast， 此時我們需要將按照第1維切分的 Tensor 拼接起來，這時候 OneFlow 會自動插入 AllGather 操作：

Boxing：通過 AllGather 實現(xiàn) Split(1) 轉(zhuǎn) Broadcast

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

2-D SBP 其實就是將兩組 1-D SBP 按照設(shè)備拓撲的維度拼起來就可以得到。

對于 數(shù)據(jù)并行的 MatMul （a x b = out）操作， SBP Signature是：{a : Split(0), b : Broadcast, out : Split(0)}， 模型并行（行切分， RowParallelLinear） 的 SBP Signature 是 ：{a : Split(1), b : Split(0), out : PartialSum}， 那如果邏輯上的一個 MatMul Op 同時做數(shù)據(jù)并行和模型并行， 其 2-D SBP Signature 就是：{a : [Split(0), Split(1)], b : [Broadcast, Split(0)], out : [Split(0), PartialSum] } （其實這個 out 的 [Split(0), PartialSum] 就是 GPT 中 RowParallelLinear 輸出 Tensor 的 SBP），當下游消費這個 out Tensor 的 期望 SBP 是 [Split(0), Broadcast]時， OneFlow 會自動在這兩組 Op 之間 插入一組 AllReduce 通信 Op， 且 這組 AllReduce 通信 Op 是按照設(shè)備拓撲的 第 0 維 分組進行的。（如 設(shè)備拓撲是 (4, 8) 時， 所有的 AllReduce 會分為 4 組，每組內(nèi)的 8 個設(shè)備會互相 AllReduce， 組間沒有通信。這就實現(xiàn)了 組內(nèi)（機器內(nèi)）模型并行， 組間（機器間）數(shù)據(jù)并行）

其實 GPT 中用到的 2-D SBP 只是最簡單情形的特例， 分布式下的并行經(jīng)過 2-D SBP 可以拓展出非常多復(fù)雜、靈活多邊的組合出來。而針對復(fù)雜的組合， 再想用 Megatron 那套就非常難做了，但是對于 OneFlow 而言，二者的難度是一樣的，因為本質(zhì)上是用 Boxing 完成 一組 2-D SBP 的變換而已。

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

OneFlow 跟 Megatron 相比，除了用戶接口（分布式易用性） 和框架設(shè)計上更簡潔、更易用，我們在已有的測試規(guī)模上性能也略微領(lǐng)先 Megatron（其實經(jīng)過 NVIDIA 的深度優(yōu)化， Megatron 在 GPU 上的分布式訓練性能已經(jīng)接近極致了，DeepSpeed 也比不上， 而我們在 Megatron 的基礎(chǔ)上 易用性、效率都更進一步）。

注：由于我們自己擁有的集群規(guī)模限制，我們只測試了 4 機 32卡 16GB V100 規(guī)模內(nèi)的性能結(jié)果。我們非常歡迎有大規(guī)模集群硬件的伙伴一起通過 OneFlow 合作研究、訓練大規(guī)模預(yù)訓練模型。未來我們會公布更大規(guī)模集群上 OneFlow 的優(yōu)異表現(xiàn)。

以下的所有實驗數(shù)據(jù)均在相同的硬件環(huán)境、相同的第三方依賴（CUDA、 cuDNN等）、使用相同的參數(shù)和網(wǎng)絡(luò)結(jié)構(gòu)下， 對比了 OneFlow 和 Megatron 在 GPT 模型下的性能表現(xiàn)。所有的性能結(jié)果均保證公開且可復(fù)現(xiàn)。我們的 GPT 模型腳本在 ：Oneflow-Inc/OneFlow-Benchmark 倉庫， 公開的評測報告、復(fù)現(xiàn)方式稍后在：Oneflow-Inc/DLPerf 倉庫中可以查看。

1.數(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 = 隱藏層大小（hidden-size）影響每層 Transformer Layer 的模型大小

· l = Transformer Layer 層數(shù)（num-layers）

數(shù)據(jù)并行性能對比

2.模型并行性能對比

注：由于單卡 GPU 顯存限制，我們對于各組參數(shù)里的模型大小是不同的，所以整體不是像數(shù)據(jù)并行那樣呈一個線性增加的關(guān)系。如第 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ù)對比

3.混合并行（數(shù)據(jù)&模型）性能對比

數(shù)據(jù) + 模型并行性能對比（注：其中前 4 組的模型規(guī)模一致；后 2 組的模型規(guī)模一致。）

4.流水并行 + 混合并行（數(shù)據(jù)&模型） 性能對比

數(shù)據(jù)+模型+流水并行性能對比（注：第 1 組參數(shù)的模型比后 3 組的都要小，因為機器內(nèi)的數(shù)據(jù)并行限制了參數(shù)規(guī)模。）

小結(jié)

OneFlow 在分布式訓練領(lǐng)域擁有獨特的設(shè)計和視角，解決了分布式訓練中的各種并行難題，因此在大規(guī)模預(yù)訓練模型場景下用 OneFlow 做分布式訓練更易用也更高效。但相比 PyTorch 在單機單卡視角下的極致易用性，OneFlow 的前端用戶接口還有明顯的差距。

OneFlow 研發(fā)團隊正在全力提升框架的單卡使用體驗， 并從即將在 5 月發(fā)布的下個大版本 OneFlow v0.4.0 起， OneFlow 開始提供兼容 PyTorch 的全新接口以及動態(tài)圖等特性。