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

OneFlow 用一致性視角輕松填平分布式訓(xùn)練難的鴻溝

對于上面 PyTorch 遇到的分布式難題， OneFlow 如何輕松解決？

OneFlow 的兩大獨(dú)特設(shè)計是：

1. 運(yùn)行時 Actor 機(jī)制

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

對于分布式集群環(huán)境（多機(jī)多卡訓(xùn)練場景），OneFlow 會把整個分布式集群抽象成為一個超級設(shè)備，用戶只需要在這個超級設(shè)備上搭建深度學(xué)習(xí)模型即可。這個虛擬出來的超級設(shè)備我們就稱之為邏輯視角（Logical View），而實際上的分布式集群的多機(jī)多卡就是物理視角（Physical View），OneFlow維護(hù)邏輯視角和物理視角之間的數(shù)學(xué)上的正確性就稱之為一致性視角（Consistent View）。

一致性視角（Consistent View）抽象

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

1.流水并行 ，只需要配置 Placement 就夠了

Placement 描述了邏輯上的 Op 和 物理上的 Op 的映射關(guān)系， 一個 Op 分布在哪些機(jī)器、哪些設(shè)備上，只需要在邏輯圖里指定其 Placement 即可。在流水并行里，用戶不需要像 PyTorch 一樣關(guān)心什么時機(jī) send/recv 、 到底是執(zhí)行 forward 還是 forward + backward，用戶只需要配置 Placement （現(xiàn)階段我們還需要用戶幫忙指定 stage id）就完成了流水并行，以下就是 OneFlow 來做 GPT 的流水并行的代碼（其實是包含了所有的并行方式）：

class Transformer(object):
    def __call__(self, hidden_states):
        for i in range(self.num_layers):
            with distribute.layer_placement_scope(i):
                h = self.layers[i](h)

其中 layer_placement_scope 就是在 scope 中配置了 placement 和 stage id：

def layer_placement_scope(layer_idx, device="gpu"):
    dist_util = get_dist_util()
    with flow.scope.placement(
        device, dist_util.get_layer_placement(layer_idx), dist_util.parallel_hierarchy,
    ):
        if dist_util.is_pipeline_parallel():
            with flow.experimental.scope.config(
                pipeline_stage_id_hint=dist_util.get_layer_stage(layer_idx)
            ):

其余的所有流水并行的細(xì)節(jié)：如各個 stage 之間怎么通信， 做前向還是后向， 怎么流水起來，怎么保證正確性，這些全都不用用戶操心。

在下一章節(jié)我們還會介紹 OneFlow 內(nèi)部是怎么實現(xiàn)流水并行的，相較于 Megatron 的復(fù)雜實現(xiàn)， OneFlow 系統(tǒng)層面的工作對框架開發(fā)者而言也更加友好。

上圖展示了一個可能的 Placement 例子，用于 GPU0 和 GPU1 之間的流水并行。圖中負(fù)責(zé)在 CPU 和 GPU、GPU 與 GPU 之間進(jìn)行數(shù)據(jù)搬運(yùn)的Op（CopyH2D、CopyD2D）是 OneFlow 系統(tǒng)自動添加的。

2.數(shù)據(jù) + 模型的混合并行，只需要配置 Variable 的 SBP 就夠了

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

Split：表示物理上的多個 Tensor 是由邏輯上的 Tensor 進(jìn)行切分后得到的。Split 會包含一個參數(shù) Axis，表示被切分的維度。如果把所有物理上的 Tensor 按照 Split 的維度進(jìn)行拼接，就能還原出邏輯上的 Tensor。

Broadcast：表示物理上的多個 Tensor 是邏輯上 Tensor 的復(fù)制，兩者數(shù)據(jù)完全相同。

PartialSum：表示物理上的多個 Tensor 跟邏輯上的 Tenso r的形狀相同，但每個對應(yīng)位置上元素的值 是邏輯Tensor對應(yīng)位置元素的值的一部分。如果把所有物理上的 Tensor 按照對應(yīng)位置相加（element-wise），即可還出原邏輯上的 Tensor。

下圖展示了SBP的幾種簡單情形。

SBP 邏輯與物理 Tensor 的對應(yīng)關(guān)系

需要注意的是，對于同一個邏輯上的 Tensor，其物理上的 Tensor 的映射關(guān)系可能會有多種，這取決于生產(chǎn)這個 Tensor 的 Op 和消費(fèi)這個 Tensor 的 Ops 是如何看待這個邏輯上的 Tensor 的。

那么用 OneFlow 做數(shù)據(jù)并行、模型并行，需要用戶做什么呢？其實只需要配置 Variable 的 SBP 即可。我們簡單介紹一下數(shù)據(jù)并行和模型并行在 OneFlow 里的配置方式：

數(shù)據(jù)并行下，每個設(shè)備上都有整份的模型，所以 Variable 的 SbpParallel 是 Broadcast， 表示物理上的每個設(shè)備上的模型都是邏輯上的完整模型的復(fù)制。其余的用戶就不用再做任何操作了（其實數(shù)據(jù)并行下，反向梯度更新的同步操作 AllReduce 是 OneFlow 系統(tǒng)內(nèi)部根據(jù) SBP 的推導(dǎo)規(guī)則自動插入的。）

模型并行下，每個設(shè)備都把模型切分并只保留一部分， 所以 Variable 的 SbpParallel 是 Split(0)， 表示物理上的每個設(shè)備上的模型都是邏輯上的完整模型經(jīng)過第0維切分后的。其余的用戶就不用再做任何操作了。前后向的數(shù)據(jù)同步操作也是 OneFlow 系統(tǒng)內(nèi)部根據(jù) SBP 推導(dǎo)規(guī)則自動插入的。

其實對于 Linear Layer 里的 Variable （假設(shè)是 row major 存儲），Split(0) 和 Split(1) 分別表示了兩種不同的模型并行方式。如果是 Split(0) 前后向就需要插入 AllReduce， 如果是 Split(1) ，前后向就需要插入 AllGather 操作了。至于為什么要插入 AllReduce 或者 AllGather，我會在下一章節(jié)介紹 SBP 推導(dǎo)的時候詳細(xì)解釋。另外，其實數(shù)據(jù)并行梯度更新要插入 AllReduce 做梯度同步，在 OneFlow 里也是自動推導(dǎo)出來的，并不是一個像 PyTorch DDP 一樣的模塊做特判。

而且 OneFlow 的 Consistent View 還保證了：任何配置 SBP 得到的并行結(jié)果， OneFlow 都保證了 其計算在數(shù)學(xué)上是完全一致的， 我們從機(jī)制上保證了分布式訓(xùn)練的正確性難題， 這一點是現(xiàn)在的 PyTorch 無法做到的。

3. 2D SBP

那么如何同時讓一個 Op 既做數(shù)據(jù)并行，又做模型并行（分組）？這里就用到了 2-D SBP。

在 2-D SBP 下，（其實 OneFlow 還支持?jǐn)U展到任意維度 N-D）集群和設(shè)備呈現(xiàn)為一個 2-D 的拓?fù)浣Y(jié)構(gòu)。比如我們一共有 2 機(jī) 8 卡（每臺機(jī)器 4 張 GPU），我們可以將 8 個設(shè)備表示成一個 (2 x 4) 的矩陣， 那么 如何在機(jī)器間數(shù)據(jù)并行、機(jī)器內(nèi)模型并行呢？用戶只需要將 Variable 的 2-D SBP 配置成： [ Broadcast, Split(0) ] 即可，那么實際上 各個設(shè)備上的 物理 Tensor 跟 邏輯 Tensor 的映射關(guān)系如下圖所示。在第一維上的 Broadcast， 表示 GPU0 和 GPU4、GPU1 和 GPu 5、 GPU2 和 GPU6、 GPU3 和 GPU7 在機(jī)器間做數(shù)據(jù)并行，在第二維上的 Split(0)， 表示 GPU0，1，2，3 、 GPU4，5，6，7 在機(jī)器內(nèi)做模型并行。

2D-SBP 下 邏輯 Tensor 和 物理 Tensor 的對應(yīng)關(guān)系

在 2D SBP 下， 設(shè)備拓?fù)浔环Q之為 ParallelConf::hierarchy ，表示如何將一維的 world_size = 8 映射成為二維的層次結(jié)構(gòu)：(2, 4)。而 2-D SBP 被稱之為ParallelDistribution， 是由兩個 SbpParallel 組成的 list: [Broadcast, Split(0)]。只需要簡單配置一下 SBP， 用戶就可以輕松設(shè)定網(wǎng)絡(luò)是數(shù)據(jù)并行、模型并行還是混合并行。

3. OneFlow：讓每一位算法工程師都有能力訓(xùn)練 GPT

綜合比較 PyTorch 和 OneFlow 在分布式并行上的設(shè)計以及開放給用戶的接口，有如下總結(jié)：

PyTorch 分布式的困境：

PyTorch 只有物理視角，沒有邏輯視角，所以 PyTorch 的用戶想要做 分布式并行，任何時候都需要自己推導(dǎo)深度學(xué)習(xí)模型中哪處需要跟其他的物理設(shè)備進(jìn)行通信和數(shù)據(jù)同步操作，既要推導(dǎo)通信所在的位置，又要推導(dǎo)通信的操作類型，還要推導(dǎo)跟其他哪些設(shè)備通信。這個在簡單的數(shù)據(jù)并行下可以使用 DDP 或者 Horovod 來實現(xiàn)，但是在復(fù)雜的模型并行、混合并行下，做并行的門檻會非常高。

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

PyTorch 在非對稱的并行方式里（如流水并行），各個設(shè)備的調(diào)度邏輯需要用戶自己手寫，用戶需要自己精細(xì)的控制每個設(shè)備上的啟動以及執(zhí)行邏輯，且執(zhí)行邏輯把前后向執(zhí)行和send/recv通信操作糅合在了一起，即使在最規(guī)整的 Transformer Layer 的流水并行下也很復(fù)雜，想要擴(kuò)展到其他模型上的工作量也很大。

PyTorch 沒有機(jī)制保證分布式并行訓(xùn)練中的正確性 和 數(shù)學(xué)一致性。即使用戶寫錯了通信操作、插錯了位置、 跟錯誤的設(shè)備進(jìn)行了通信，PyTorch也檢查不出來。

這些都使得用戶使用 PyTorch 開發(fā)復(fù)雜分布式訓(xùn)練的腳本極為困難，以至于只有 NVIDIA 、 微軟 等大企業(yè)的分布式訓(xùn)練專家才能開發(fā)出 GPT 的 PyTorch 版本。

OneFlow 分布式的易用性：

OneFlow 有一致性視角 Consistent View， 將分布式訓(xùn)練下的多機(jī)通信和 算法邏輯解耦，使得用戶可以不用關(guān)心分布式訓(xùn)練的細(xì)節(jié)，降低了分布式訓(xùn)練的使用門檻。

OneFlow 通過 Placement + SBP 機(jī)制解決了分布式訓(xùn)練中任意并行場景的需求，用戶只需要配置 op 的 Placement 就可以完成流水并行，只需要配置 Tensor 的 SBP 就可以實現(xiàn)數(shù)據(jù)并行、模型并行和混合并行。而且任何并行方式都是 Placement + SBP 的一種特例， OneFlow 從系統(tǒng)層面不需要做任何的特判，SBP 才是各種分布式并行的本質(zhì)。

OneFlow 的通信邏輯可以復(fù)用，不需要為任何特定網(wǎng)絡(luò)和特定算子實現(xiàn)相應(yīng)的通信邏輯。通信邏輯由 OneFlow 的 Boxing 機(jī)制完成，與具體的算子和模型無關(guān)。

OneFlow 的 SBP 還保證了數(shù)學(xué)上的一致性， 相同的邏輯上的模型腳本，使用任意的并行方式（數(shù)據(jù)并行、模型并行、流水并行）、使用任意的集群拓?fù)洌琌neFlow 都從數(shù)學(xué)上保證了模型分布式訓(xùn)練的正確性。

因此，我們才說 OneFlow 用一套簡潔設(shè)計解決了分布式并行的各種難題。對于用戶而言， OneFlow 讓算法工程師不需要成為一位分布式訓(xùn)練的專家也有能力做復(fù)雜的分布式訓(xùn)練， 只要有硬件資源，任何一位算法工程師都可以訓(xùn)練 GPT， 任何一位算法工程師都可以開發(fā)一個新的大規(guī)模分布式訓(xùn)練的模型。