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

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

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

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

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

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

一致性視角（Consistent View）抽象

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

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

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

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é)：如各個(gè) stage 之間怎么通信， 做前向還是后向， 怎么流水起來(lái)，怎么保證正確性，這些全都不用用戶(hù)操心。

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

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

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

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

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

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

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

下圖展示了SBP的幾種簡(jiǎn)單情形。

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

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

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

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

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

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

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

3. 2D SBP

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

在 2-D SBP 下，（其實(shí) OneFlow 還支持?jǐn)U展到任意維度 N-D）集群和設(shè)備呈現(xiàn)為一個(gè) 2-D 的拓?fù)浣Y(jié)構(gòu)。比如我們一共有 2 機(jī) 8 卡（每臺(tái)機(jī)器 4 張 GPU），我們可以將 8 個(gè)設(shè)備表示成一個(gè) (2 x 4) 的矩陣， 那么 如何在機(jī)器間數(shù)據(jù)并行、機(jī)器內(nèi)模型并行呢？用戶(hù)只需要將 Variable 的 2-D SBP 配置成： [ Broadcast, Split(0) ] 即可，那么實(shí)際上 各個(gè)設(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 的對(duì)應(yīng)關(guān)系

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

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

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

PyTorch 分布式的困境：

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

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

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

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

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

OneFlow 分布式的易用性：

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

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

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

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

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