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

發(fā)布人：AI科技大本營時間：2021-05-15 來源：工程師

加入技術交流群
- 掃碼加入
  和技術大咖面對面交流
  海量資料庫查詢

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

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

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

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

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

Actor 狀態(tài)機

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

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

數據預處理流程

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

數據預處理 pipeline 時間線

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

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

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

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

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

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

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

我們先以 1-D SBP 為例，之后再擴展到 2-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)），假設共有兩張卡，則 GPU 0 上有前一半的數據（Tensor a，Shape = (m/2, k)），GPU 1 上有后一半的數據，則我們說 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）。

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

因為 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 系統會自動插入通信節(jié)點以完成數據的切分/傳輸/拼接等操作，使得下游 Op 總能拿到按照自己期望 SBP 切分的 Tensor。

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

Boxing：通過 AllReduce 實現 PartialSum 轉 Broadcast

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

Boxing：通過 AllGather 實現 Split(1) 轉 Broadcast

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

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

對于數據并行的 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 同時做數據并行和模型并行，其 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 是按照設備拓撲的第 0 維分組進行的。（如設備拓撲是 (4, 8) 時，所有的 AllReduce 會分為 4 組，每組內的 8 個設備會互相 AllReduce，組間沒有通信。這就實現了組內（機器內）模型并行，組間（機器間）數據并行）

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

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

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

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

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

1.數據并行性能對比

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

· 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）

數據并行性能對比

2.模型并行性能對比

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

其中 l 表示 num-layers ，h 表示 hidden-size， V 表示詞表大?。╲ocabulary size = 51200）， S 表示句子長度（a sequence length = 2048）， P 表示參數規(guī)模。