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

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

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

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

OneFlow 的運(yùn)行時(shí) Actor 機(jī)制有以下幾個(gè)特點(diǎn)：

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

Actor 狀態(tài)機(jī)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

第二個(gè)是 Split(1) -> Broadcast， 此時(shí)我們需要將按照第1維切分的 Tensor 拼接起來(lái)，這時(shí)候 OneFlow 會(huì)自動(dòng)插入 AllGather 操作：

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

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

2-D SBP 其實(shí)就是將兩組 1-D SBP 按照設(shè)備拓?fù)涞木S度拼起來(lái)就可以得到。

對(duì)于 數(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}， 那如果邏輯上的一個(gè) MatMul Op 同時(shí)做數(shù)據(jù)并行和模型并行， 其 2-D SBP Signature 就是：{a : [Split(0), Split(1)], b : [Broadcast, Split(0)], out : [Split(0), PartialSum] } （其實(shí)這個(gè) out 的 [Split(0), PartialSum] 就是 GPT 中 RowParallelLinear 輸出 Tensor 的 SBP），當(dāng)下游消費(fèi)這個(gè) out Tensor 的 期望 SBP 是 [Split(0), Broadcast]時(shí)， OneFlow 會(huì)自動(dòng)在這兩組 Op 之間 插入一組 AllReduce 通信 Op， 且 這組 AllReduce 通信 Op 是按照設(shè)備拓?fù)涞?第 0 維 分組進(jìn)行的。（如 設(shè)備拓?fù)涫?(4, 8) 時(shí)， 所有的 AllReduce 會(huì)分為 4 組，每組內(nèi)的 8 個(gè)設(shè)備會(huì)互相 AllReduce， 組間沒(méi)有通信。這就實(shí)現(xiàn)了 組內(nèi)（機(jī)器內(nèi)）模型并行， 組間（機(jī)器間）數(shù)據(jù)并行）

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

GPT 分布式訓(xùn)練性能對(duì)比：OneFlow vs Megatron

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

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

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

1.數(shù)據(jù)并行性能對(duì)比

注：每組參數(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ù)并行性能對(duì)比

2.模型并行性能對(duì)比

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

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

模型并行數(shù)據(jù)對(duì)比

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

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

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

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

小結(jié)

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

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