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

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

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

OneFlow 的兩大獨特設計是：

1. 運行時 Actor 機制

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

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

一致性視角（Consistent View）抽象

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

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

Placement 描述了邏輯上的 Op 和 物理上的 Op 的映射關系， 一個 Op 分布在哪些機器、哪些設備上，只需要在邏輯圖里指定其 Placement 即可。在流水并行里，用戶不需要像 PyTorch 一樣關心什么時機 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)
            ):

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

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

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

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

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

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

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

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

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

SBP 邏輯與物理 Tensor 的對應關系

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

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

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

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

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

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

3. 2D SBP

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

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

2D-SBP 下 邏輯 Tensor 和 物理 Tensor 的對應關系

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

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

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

PyTorch 分布式的困境：

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

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

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

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

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

OneFlow 分布式的易用性：

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

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

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

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

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