一、背景

之前的文章中我們?cè)敿?xì)介紹了大規(guī)模分布式訓(xùn)練中的數(shù)據(jù)并行（Data Parallelism，DP）、張量并行（Tensor Parallelism，TP）和流水線并行（Pipeline Parallelism，PP）。這篇文章中我們繼續(xù)介紹 MoE 中經(jīng)常使用的專家并行（Expert Parallelism，EP），以及 EP 中涉及的相關(guān) All2All 操作和優(yōu)化手段等。

二、引言

2.1 標(biāo)準(zhǔn) All2All

AlltoAll 是集合通信庫(kù)（比如 NCCL）中另一種常見(jiàn)的通信原語(yǔ)，用于多個(gè)設(shè)備之間進(jìn)行數(shù)據(jù)交換。AlltoAlll 操作允許每個(gè)參與的設(shè)備將其本地?cái)?shù)據(jù)分發(fā)到其他設(shè)備，同時(shí)從其他設(shè)備接收數(shù)據(jù)。

如下圖所示是一種標(biāo)準(zhǔn)的 AlltoAll 操作，有 4 個(gè) GPU，每個(gè) GPU 包含 4 個(gè)數(shù)據(jù)。通過(guò) AlltoAll 操作之后每個(gè)設(shè)備都將 4 個(gè)數(shù)據(jù)分發(fā)給 4 個(gè) GPU，同時(shí)從 4 個(gè) GPU 接收數(shù)據(jù)。可以看出，AlltoAll 很像一個(gè)矩陣的轉(zhuǎn)置操作：

如下圖所示為 Pytorch 實(shí)現(xiàn)一個(gè)上述標(biāo)準(zhǔn) AlltoAll 的示例：

實(shí)際上 NCCL 中并沒(méi)有 AlltoAll 通信原語(yǔ)，需要通過(guò) ncclSend 和 ncclRecv 實(shí)現(xiàn)，其中 ncclSend 和 ncclRecv 是一個(gè) P2P 通信。如下圖所示，每個(gè) Rank 都發(fā)送 nranks 塊數(shù)據(jù)，同時(shí)接收 nranks 塊數(shù)據(jù)就實(shí)現(xiàn)了 AlltoAll 的功能。（可以參考 Point-to-point communication — NCCL 2.22.3 documentation）

類似的方式就可以實(shí)現(xiàn) one2all（Scatter）操作：

類似的方式也可以實(shí)現(xiàn) all2one（Gather）操作：

2.2 非標(biāo)準(zhǔn) All2All

實(shí)際上有些場(chǎng)景并非均勻發(fā)送和接收，有可能發(fā)送到不同設(shè)備的數(shù)據(jù)量不同，從不同設(shè)備接收的數(shù)據(jù)量也可能不同。Pytorch 的 “torch.distributed.all_to_all_single” 提供了 input_split_sizes 和 output_split_sizes 參數(shù)來(lái)支持：

• input_split_sizes 表示向每個(gè)設(shè)備發(fā)送的數(shù)據(jù)量。
• output_split_sizes 表示從每個(gè)設(shè)備接收的數(shù)據(jù)量。

如下圖所示，4 個(gè) GPU，每個(gè) GPU 都包含 10 個(gè)數(shù)據(jù)：

• 4 個(gè) GPU 都向 GPU k 發(fā)送 k+1 個(gè)數(shù)據(jù)（比如，都向 GPU 3 發(fā)送 4 個(gè)數(shù)據(jù)）。
• GPU k 從所有 GPU 都接收 k+1 個(gè)數(shù)據(jù)（比如，GPU 2 從所有 GPU 都接收 3 個(gè)數(shù)據(jù)）。?

如下圖所示為 Pytorch 實(shí)現(xiàn)一個(gè)上述非標(biāo)準(zhǔn) all2all 的示例：

PS：需要指出的是，上述接口中 output_split_sizes 和 input_split_sizes 的個(gè)別值也可以是 0，表示不從某個(gè)設(shè)備接收，或者不向某個(gè)設(shè)備發(fā)送數(shù)據(jù)。如上所示，all2all 底層是用 ncclSend 和 ncclRecv 實(shí)現(xiàn)，很容易可以做到這一點(diǎn)。

2.3 兩次 All2All

上述非標(biāo)準(zhǔn) All2All 中有個(gè)問(wèn)題：有些時(shí)候當(dāng)前設(shè)備只知道要向其他設(shè)備發(fā)送多少數(shù)據(jù)，而并不知道需要從其他設(shè)備接收多少數(shù)據(jù)。這個(gè)問(wèn)題可以通過(guò) 2 次 all2all 來(lái)解決：

• 第一次 all2all 交換要傳輸?shù)臄?shù)據(jù)量信息，這是一個(gè)標(biāo)準(zhǔn)的 all2all 操作，如下圖紅框所示。
• 第二次 all2all 根據(jù)上述獲取的數(shù)據(jù)量信息來(lái)執(zhí)行真正的數(shù)據(jù)傳輸，此時(shí)是一個(gè)非標(biāo)準(zhǔn) all2all 操作。?

三、Google Sparsely-Gated MoE

3.1 摘要

在 [1701.06538] Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer 中，作者（也包括大名鼎鼎的 Geoffrey Hinton 和 Jeff Dean）將 MoE 引入到 LSTM 模型中，并提出了稀疏 MoE（Sparse MoE）的概念。在 Sparse MoE 中可以包含數(shù)千個(gè)前饋?zhàn)泳W(wǎng)絡(luò)，并由可訓(xùn)練的門控網(wǎng)絡(luò)（Gating Network）確定這些專家的稀疏組合。

作者將 MoE 應(yīng)用于語(yǔ)言建模和機(jī)器翻譯任務(wù)，在相應(yīng)基準(zhǔn)測(cè)試中，這些模型可以以較低的計(jì)算成本獲得優(yōu)于 SOTA 的效果。

3.2 方案

如下圖 Figure 1 所示，作者引入了 Gating Network 機(jī)制，該機(jī)制可以選出 Topk 的 Expert（Expert 2 和 Expert n-1）進(jìn)行計(jì)算。這種稀疏性意味著只有部分專家被激活處理特定的輸入，從而可以大大降低計(jì)算量：

作者也進(jìn)一步證明可以通過(guò)靈活控制專家數(shù)，來(lái)獲得不同容量的模型。如下圖 Table 8 所示，作者分別構(gòu)建了 32/256/1024/4096/16384/65535/131072 個(gè)專家的模型，其最大為 137B 的 LSTM 模型。由于稀疏性的存在，雖然 137B 參數(shù)量很大，但可以比當(dāng)時(shí) SOTA 模型更低的計(jì)算成本下獲得更好的效果：

3.3 Expert Parallelism

然而，想要高效訓(xùn)練上述的模型卻非常有挑戰(zhàn)，假設(shè)有 d 個(gè)設(shè)備，采用 Data Parallelism 的 Micro Batch Size 為 b，Total 的 Mini Batch Size 為 d*b。同樣假設(shè)模型中有 n 個(gè) Expert，每個(gè) Sample 都選擇 k 個(gè)激活。

基于以上條件，對(duì)于每個(gè) DP Worker，每個(gè) Expert 平均有 k*b/n 個(gè) Sample。由于往往 k 遠(yuǎn)小于 n，比如 n 可以是 256/1024/4096/16K/64K/128K，而 k 只是 2 或 4，也就導(dǎo)致 k*b/n << b。也就是說(shuō)，平均到每個(gè) Expert 上的樣本會(huì)遠(yuǎn)小于 b，這種方式很不利于模型的擴(kuò)展。

作者采用了特殊的模型并行方案，不同的 DP Worker 會(huì)共享同一份 Expert 參數(shù)，非 Expert 不共享。比如說(shuō)模型有 256 個(gè) Expert，32  個(gè) DP 的 Worker，則每個(gè) Worker 上各有 8 個(gè) Expert，不同 DP 的 Sample 經(jīng) Gate 后重新組合分發(fā)到對(duì)應(yīng)的 Expert，此時(shí)平均每個(gè) Expert 的 Sample 數(shù)為 d*k*b/n，只要 Mini Batch Size 大小 d*b 隨著 Expert 數(shù)的增加（相當(dāng)于參數(shù)量的增加）而增加，就可以保證每個(gè) Expert 的 Sample 數(shù)為常數(shù)，對(duì)顯存和帶寬的需求也就基本是常數(shù)。

如下圖所示，DP 為 3，Expert 個(gè)數(shù)為 9，每個(gè) Worker 上分 3 個(gè) Expert，而其他部分參數(shù)不共享：

四、Google Gshard

4.1 摘要

在 [2006.16668] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding 中，作者首次將 MoE 引入到 Transformer 模型中。此外，作者發(fā)布了 GShard，它由一組輕量級(jí) API 和 XLA 編譯器組成，提供了一種優(yōu)雅的方式來(lái)表達(dá)各種并行計(jì)算模式，與此同時(shí)只需對(duì)模型代碼進(jìn)行少量的更改。使用 GShard 可以將 Transformer 模型擴(kuò)展到 600B 參數(shù)，并且可以在 2048 個(gè) TPU v3 上用 4 天訓(xùn)完。

4.2 方案

具體的方案如下圖 Figure 3 所示：

• 增加Position-wise Sparsely Gated MoE層，將FFN 層替換為 MoE 結(jié)構(gòu)，MoE 中的每個(gè)專家都是一個(gè) FFN（每個(gè)專家大小相同）

Gating 模塊：通過(guò)Gating 模塊將輸入路由到不同的專家（Transformer 模型輸入的是 Token 序列，因此每個(gè) Token都會(huì)通過(guò) Gating 選擇不同的專家，而不是整個(gè)序列使用相同的專家，默認(rèn)為top2）。

Random routing：有些時(shí)候 Gating 模塊得到的排名第二的專家的分?jǐn)?shù)會(huì)很低，此時(shí)可以簡(jiǎn)單的忽略第二個(gè)專家。

• 并非是每一層的 FFN 都替換為 MoE，而是間隔一層替換，如果有 12 層，則只有 6 層有 MoE（通常是可配置的）。
• 采用專家并行（Expert Parallel，EP）策略，每個(gè)設(shè)備一個(gè)專家，除 MoE 之外的模型其它部分在所有設(shè)備存儲(chǔ)一份相同的副本。（如果有 128 個(gè)專家，則使用 128 個(gè) TPU Core；2048 個(gè)專家，則使用 2048 個(gè) TPU Core）?

PS：從上可以看出，其實(shí)與上述的 Sparsely-Gated MoE 非常類似，只是擴(kuò)展到了 Transformer 模型。在 Gating 后通過(guò) All-to-All 來(lái)分布 Token，在 Expert 后通過(guò) All-to-All 再次重組 Token，后文詳細(xì)介紹。

五、Switch Transformer

5.1 摘要

在 [2101.03961] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity 中，作者相比 Gshard 等方案主要做了三點(diǎn)改進(jìn)：簡(jiǎn)化稀疏路由、高效稀疏路由，以及增強(qiáng)的訓(xùn)練和微調(diào)技巧。通過(guò)這些手段可以降低通信和計(jì)算成本，并減輕訓(xùn)練的不穩(wěn)定性。在相同的計(jì)算資源下，基于 T5-Base 和 T5-Large 設(shè)計(jì)的模型的預(yù)訓(xùn)練可以加速 7 倍。

PS：這里我們主要關(guān)注前兩個(gè)改進(jìn)。

5.2 簡(jiǎn)化稀疏路由

在 [1701.06538] Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer  中，稀疏專家的數(shù)目需要 > 1，在 Gshard 中作者也是使用的 top2 專家。而 Switch Transformer 中，作者發(fā)現(xiàn)僅使用一個(gè)專家也能保證模型的質(zhì)量。這樣有 3 個(gè)好處：

• Router 計(jì)算更簡(jiǎn)單，通信量也更少。
• 一個(gè) Token 僅對(duì)應(yīng)一個(gè)專家，計(jì)算量也更少。
• 平均每個(gè)專家對(duì)應(yīng)的batch size至少可以減半。

如下圖 Figure 2 所示，其模型結(jié)構(gòu)和 Gshard 中類似，圖中的紅框和綠框是同樣的 MoE，只是對(duì)應(yīng)不同的輸入，經(jīng) Router 后也只連接一個(gè)專家：

5.3 高效稀疏路由

作者采用 Mesh-TensorFlow（PS：之前的文章中介紹過(guò)），其提供和 TensorFlow 相似的 API，提供了更簡(jiǎn)單的分布式數(shù)據(jù)并行和模型并行。作者的模型主要針對(duì) TPU 設(shè)計(jì)，其在模型訓(xùn)練中不支持動(dòng)態(tài) Tensor shape，也就是要求每個(gè)專家輸入的 Tensor shape 是固定的。然而，路由是動(dòng)態(tài)的，相應(yīng)路由到每個(gè)專家的 Tensor 的 shape 也是動(dòng)態(tài)的，為了解決這一問(wèn)題，作者使用了專家容量（Expert Capacity），如下所示，專家容量為每個(gè) Batch 中總的 Token 數(shù)除以專家數(shù)，然后再乘以容量因子（Capacity Factor），即可得到專家容量（每個(gè)專家對(duì)應(yīng)的 Token 數(shù)）。

如下圖 Figure 3 所示，有 6 個(gè) Token，3 個(gè)專家，平均每個(gè)專家 2 個(gè) Token：

• 容量因子為 1.0：如下圖中所示，則對(duì)應(yīng)的專家容量為 2：

Expert 1有 3 個(gè) Token，則需要丟棄一個(gè)通過(guò)殘差連接直接傳到下一層。

Expert 2 有 2 個(gè) Token，正好。

Expert 3只有 1 個(gè) Token，需要Padding 1 個(gè)空的 Token。

• 容量因子為 1.5：如下圖右所示，則對(duì)應(yīng)的專家容量為 3：
• Expert 1 有 3 個(gè) Token，正好。
• Expert 2 只有 2 個(gè) Token，需要 Padding 1 個(gè)空的 Token。
• Expert 3 只有 1 個(gè) Token，需要 Padding 2 個(gè)空的 Token。?

從上也可以看出，容量因子越大，需要 Padding 的 Token 也就越多，無(wú)效計(jì)算越多；負(fù)載越不均衡，需要 Padding 的 Token 也就越多，無(wú)效計(jì)算越多。為了更好的負(fù)載均衡，作者同樣添加了 Load Balancing Loss。

六、FastMoE

6.1 摘要

之前的高性能分布式 MoE 訓(xùn)練系統(tǒng)主要是針對(duì) Google 的硬件（TPU）和軟件（Mesh TensorFlow），并且不向公眾開(kāi)放，針對(duì) NVIDIA GPU 和 Pytorch 還沒(méi)有相應(yīng)方案。

在 [2103.13262] FastMoE: A Fast Mixture-of-Expert Training System 中，作者提出 FastMoE，其是一個(gè)基于 Pytorch 的分布式 MoE 訓(xùn)練系統(tǒng)，并提供高度優(yōu)化的高性能加速方案。該系統(tǒng)支持將不同的專家放置在多個(gè)節(jié)點(diǎn)上的多個(gè) GPU 中，從而實(shí)現(xiàn)專家數(shù)量和 GPU 數(shù)量線性增加。

PS：如下圖所示（來(lái)自 fastmoe/doc/readme-cn.md at master），F(xiàn)astMoE 主要針對(duì)的是 Expert 比較多的場(chǎng)景，也就是一個(gè) GPU 上有 1 個(gè)或多個(gè) Expert。在 2021 年底的 v0.3.0 版本中集成了 Megatron-LM，通過(guò) Megatron-LM 的 Tensor Parallel 來(lái)實(shí)現(xiàn)一個(gè) Expert 分布在不同的 GPU 上。

6.2 系統(tǒng)設(shè)計(jì)

6.2.1 靈活性

FastMoE 的靈活性主要體現(xiàn)在以下幾個(gè)方面：

• 支持任意的網(wǎng)絡(luò)作為專家。作者對(duì)專家模塊做了抽象，用戶可以專注設(shè)計(jì)專家模塊；此外，F(xiàn)astMoE 也支持將多個(gè)專家放在同一個(gè) Worker 上。
• 針對(duì) Transformer 模型高度優(yōu)化的 FFN。尤其是當(dāng)多個(gè)專家放在一個(gè) Worker時(shí)，常見(jiàn)的方式是通過(guò) for 循環(huán)串行的執(zhí)行 Worker 上的多個(gè)專家，而作者實(shí)現(xiàn)了并行執(zhí)行不同專家的方案。（Batched Gemm）
• 插件式支持 Pytorch 和 Megatron-LM。作者對(duì) FastMoE 進(jìn)行了必要的抽象，使其很容易與其他框架集成，如下圖所示為與 Megatron-LM 集成的示例：

6.2.2 擴(kuò)展模型容量

FastMoE 的模型并行方案。FastMoE 支持將專家分布在多個(gè)節(jié)點(diǎn)的多個(gè) Worker 上，并且將不同 Worker 之間的數(shù)據(jù)通信隱藏起來(lái)，模型開(kāi)發(fā)人員不用考慮。

此外，在分布式 MoE 系統(tǒng)中的一個(gè)主要挑戰(zhàn)為：動(dòng)態(tài)路由導(dǎo)致分配給不同專家的輸入樣本數(shù)可能存在很大的差異。作者的方案為：在 Worker 之間交換實(shí)際的數(shù)據(jù)之前，先在 Worker 之間交換大小信息，Worker 根據(jù)相應(yīng)信息分配 Buffer，然后傳輸真實(shí)的數(shù)據(jù)。（PS：這就是 2.3 的兩次 All2All 操作）

異構(gòu)同步模塊。模型的不同部分可能在不同的 Worker 組間重復(fù)，這非常有挑戰(zhàn)，因?yàn)榉植际侥K不得不識(shí)別是否需要對(duì)參數(shù)的梯度進(jìn)行同步，以及與誰(shuí)同步。因此，F(xiàn)astMoE 引入了數(shù)據(jù)并行通信組標(biāo)簽：

• world：需要與所有 Worker 同步。
• data parallel：需要與模型并行組正交的數(shù)據(jù)并行組中的 Worker 同步。
• none：不需同步。

例如，無(wú)論模型并行設(shè)置如何，Gating Network 需要在所有 Worker 之間復(fù)制，因此標(biāo)簽為 world。注意力層可以劃分為模型并行子層，因此其標(biāo)簽為 data parallel。每個(gè) Worker 都包含幾個(gè)特定的專家網(wǎng)絡(luò)，其標(biāo)簽為 none。

6.3 優(yōu)化激活

FastMoE 將所有輸入樣本一起 Batching 后發(fā)給同一個(gè)專家。由于數(shù)據(jù)表示的限制，F(xiàn)astMoE 使用專門開(kāi)發(fā)的 CUDA Kernel 進(jìn)行內(nèi)存移動(dòng)，以減少開(kāi)銷。如下圖 Figure 4 所示，給定每個(gè)樣本要進(jìn)入的索引（Gating 輸出），通過(guò) Scatter 操作（PS：其實(shí)也是 All2All，詳情看 6.5）將所有樣本按照對(duì)應(yīng)順序進(jìn)行排布，執(zhí)行完專家計(jì)算之后，再按照相反的 Gather 操作（PS：其實(shí)也是 All2All，詳情看 6.5）進(jìn)行復(fù)原。（gate output 應(yīng)該為 0, 1, 2, 1, 1, 0 ?）

6.4 多 CUDA Stream 調(diào)度

如下圖 Figure 8 所示，S 表示 Send，R 表示 Receive，C 表示 Compute，通過(guò)利用 CUDA 的 Multi Stream 機(jī)制，可以最大限度實(shí)現(xiàn)通信和計(jì)算的 overlap，實(shí)現(xiàn)加速的目的：

6.5 MoE 代碼實(shí)現(xiàn)

需要說(shuō)明的是，F(xiàn)astMoE 中并不是直接調(diào)用 Pytorch 的 all2all 來(lái)實(shí)現(xiàn) Token 的分發(fā)。如下圖所示：

• prepare_forward()：通過(guò) all2all 獲得相應(yīng)的索引、大小信息。（PS：有些實(shí)現(xiàn)也會(huì)用 AllGather）
• scatter()：稀疏 all2all，分發(fā) Token。
• expert_fn()：執(zhí)行 Expert 函數(shù)。
• gater()：稀疏 all2all，Token 重新返回 Rank，此時(shí)不需要再額外獲取索引、大小信息。?

如下圖所示為 prepare_forward() 的具體實(shí)現(xiàn)，可以看出其為非稀疏 all2all() 操作：

如下圖所示，作者實(shí)現(xiàn)了 Global_Scatter 函數(shù)，可以理解為一個(gè)稀疏的 all2all 操作，因?yàn)椴⒉皇敲總€(gè) Rank 都會(huì)向其他 Rank 分發(fā)，也不是每個(gè) Rank 都從所有 Rank 接收。（參考：fastmoe//cuda/global_exchange.h）

七、Tutel

7.1 摘要

之前的 MoE 分布式訓(xùn)練系統(tǒng)往往采用靜態(tài)執(zhí)行方式（Tensor 的 Shape 在執(zhí)行中不能改變），導(dǎo)致經(jīng) Token 路由之后可能存在 Token 丟棄或者 Padding 無(wú)效計(jì)算的問(wèn)題，導(dǎo)致計(jì)算效率比較低。

在 [2206.03382] Tutel: Adaptive Mixture-of-Experts at Scale 中，作者提出了 Tutel，其具備動(dòng)態(tài)自適應(yīng)并行和流水并行（PS：非流水線并行）機(jī)制。Tutel 中作者設(shè)計(jì)了一個(gè)統(tǒng)一布局來(lái)分發(fā) MoE 模型參數(shù)和輸入數(shù)據(jù)，并利用其實(shí)現(xiàn)可切換并行性和動(dòng)態(tài)流水并行，而無(wú)需引入數(shù)學(xué)不等價(jià)操作或者 Tensor 遷移開(kāi)銷，可以在運(yùn)行時(shí)以零成本實(shí)現(xiàn)自適應(yīng)并行/流水并行優(yōu)化。基于這一關(guān)鍵設(shè)計(jì)，Tutel 實(shí)現(xiàn)了各種 MoE 加速技術(shù)，包括 Flexible All-to-All、二維分層（2DH）All-to-All，以及快速編碼、解碼等。綜合所有技術(shù)，Tutel 相比之前的方案，在 16 個(gè)和 2048 個(gè) A100 GPU 上，單個(gè) MoE 層的速度提升 4.96x 和 5.75x。

作者評(píng)估表明，Tutel 可以高效地運(yùn)行 SwinV2-MoE，其基于 Swin Transformer V2 構(gòu)建。使用 Tutel 訓(xùn)練和推理 SwinV2-MoE 比 Fairseq 加速 1.55x 和 2.11x。同時(shí)，SwinV2-MoE 在預(yù)訓(xùn)練及下游視覺(jué)任務(wù)中比對(duì)應(yīng)的密集模型實(shí)現(xiàn)了更高的準(zhǔn)確性。

7.2 自適應(yīng) MoE

鑒于 EP、DP 和 MP 派生了 7 種不同的并行方法組合，一種方案是為每種方法設(shè)計(jì)一個(gè)執(zhí)行流程，并使其可與其他方法切換。然而，實(shí)際上沒(méi)有必要設(shè)計(jì) 7 個(gè)執(zhí)行流程，因?yàn)槠淇梢院?jiǎn)化為更小但效率相當(dāng)?shù)膯?wèn)題。作者的方法是分析所有并行方法的復(fù)雜性，以將它們縮小到最小子集（這里作者只考慮最重要的通信復(fù)雜性，所有 GPU 都執(zhí)行相同的計(jì)算，計(jì)算復(fù)雜度相同，通信復(fù)雜性直接決定了一種并行方法相比其他方法的效率）。如果它們滿足以下條件則將其刪除：

1. 在任何情況下都不是最佳的。
2. 是另一種方法的特例。

如下圖 Table 3 所示為一些常見(jiàn)的參數(shù)：

作者在參數(shù)表里沒(méi)有具體介紹 r 參數(shù)，只在后文介紹，表示每個(gè)專家的 TP 數(shù)，也就是每個(gè)專家分布在幾個(gè) GPU 上：

• 如果 r=1，則表示 EP+DP+MP 變?yōu)?EP+DP
• 如果 r= W/E，則表示 EP+DP+MP 變?yōu)?EP+MP

如下圖 Table 4 所示，經(jīng)過(guò)一系列比較，作者得出結(jié)論，該子集只包含 DP（1） 和 EP+DP+MP（7）：

• 對(duì)于DP（1）：僅數(shù)據(jù)并行，不過(guò)采用的是ZeRO-DP Stage-3，可以將模型參數(shù)分布在多個(gè) GPU 設(shè)備，在前向計(jì)算的時(shí)候通過(guò) All-Gather 操作獲取所有模型參數(shù)進(jìn)行計(jì)算。在反向時(shí)，執(zhí)行一次 Reduce-Scatter。
• 對(duì)于MP（2）：僅模型并行，每個(gè) GPU 上都只有模型的 1/W，所有 GPU 加起來(lái)有一份完整模型。只要能使用 EP，則總會(huì)差于EP+MP（6）。
• 對(duì)于EP（3）：只有專家數(shù)量 >= GPU 數(shù)量才有意義，因此作者假設(shè)專家數(shù)量 < GPU 數(shù)量，這也是當(dāng)前 LLM-MoE 的現(xiàn)狀，不用考慮純 EP 的方案。?

如下圖 Figure 6 所示為相應(yīng)的 Zero-DP，假設(shè)有 4 個(gè) GPU，模型有 2 個(gè)專家，則每個(gè) GPU 都只存儲(chǔ)某個(gè)專家的 1/2。在前向計(jì)算時(shí)需要一次 All-Gather 獲取到 2 個(gè)完整的專家參數(shù)。

經(jīng)過(guò)如上的分析后，作者得出了不同的分布式方案，如下圖 Figure 8 所示，假設(shè) ZeRO-DP 為 r=0，根據(jù) r 的不同值可以選擇不同的策略，特殊情況為上述介紹的 r=1 和 r=W/E：

7.3 優(yōu)化

7.3.1 Flexible All-to-All

常規(guī)的 FFN 層計(jì)算時(shí)，All-to-All 的 data layout 會(huì)和 Word-Size 有關(guān)，當(dāng) Word-Size（GPU）數(shù)目比較大時(shí)，性能可能會(huì)下降比較多：

PS：出現(xiàn)這一問(wèn)題的主要原因是：FFN layer 主要為矩陣乘法，GPU 處理大矩陣乘法非常高效，而如果矩陣中的某一維度比較小時(shí)，會(huì)導(dǎo)致矩陣乘法處于 Roofline-Model 的 Memory-Bound 區(qū)域，導(dǎo)致無(wú)法充分發(fā)揮 GPU 算力，并且維度越小此瓶頸越明顯。當(dāng) World-Size 為 256 時(shí)，對(duì)應(yīng)的矩陣短邊為 16384/256=64，可能正好在 Roofline-Model 的轉(zhuǎn)折點(diǎn)，這也是為什么當(dāng) Worhd-Size 進(jìn)一步增大時(shí)性能會(huì)進(jìn)一步降低。

Flexible All-to-All 的目的是去除和 World-Size 的相關(guān)性，如下圖為優(yōu)化后的效果：

7.3.2 2DH All-to-All

如下圖 Figure 15 所示，2DH All-to-All 的主要思路是充分考慮數(shù)據(jù)的局部性（GPU 內(nèi)，同 node GPU、多 node GPU），將非連續(xù)內(nèi)存空間對(duì)齊到連續(xù)內(nèi)存空間，并將多個(gè)小的通信合并成大的通信：

• 第一列 -> 第二列：GPU 內(nèi)部交換數(shù)據(jù)（無(wú)通信）
• 第二列 -> 第三列：同 node 的 GPU 間交換數(shù)據(jù)（NVLink）
• 第三列 -> 第四列：GPU 內(nèi)部交換數(shù)據(jù)（無(wú)通信）
• 第四列 -> 第五列：跨 node 的 GPU 間交換數(shù)據(jù)（網(wǎng)絡(luò)）?

如下圖 Figure 20 和 Figure 21 所示，提出的 2DH All-to-All 比基線提升明顯：

7.3.3 Fast Encode 和 Decode Kernel 優(yōu)化

如下圖 Figure 3 所示，在專家并行模式下，專家層的前后會(huì)分別引入 All-to-All 通信操作。前一個(gè) All-to-All 用于將每個(gè) Worker 上的 Token 按照 Router 后對(duì)應(yīng)的專家發(fā)送到專家所在的 GPU，也叫 All-to-All（Dispatch）；而后一個(gè) All-to-All 用于將專家計(jì)算后的 Token 重新按照原來(lái)的方式排列，也叫 All-to-All（Combine）。

在 All-to-All（Dispatch）操作之前需要準(zhǔn)備好 All-to-All 的輸入，也叫 Encode；在 All-to-All（Combine）操作之后需要解包 All-to-All 的輸出，組織為原始的順序，也叫 Decode。而很多框架中 Encode 和 Decode 的實(shí)現(xiàn)都不夠高效，有很多無(wú)效計(jì)算，因此作者定制了高性能 CUDA Kernel 來(lái)優(yōu)化，如下圖（a）為未優(yōu)化的 Encode，（b）為優(yōu)化后的 Encode。

如下圖 Figure 15 所示，優(yōu)化后 Encode、Decode 相關(guān)的時(shí)間大幅降低（此外也可以有效節(jié)約顯存）：

7.3.4 Adaptive Pipelining

此外，在 Tutel 中，作者也采用了 Multi-Stream 機(jī)制來(lái)實(shí)現(xiàn)計(jì)算和通信的重疊，以提升效率，這里不再展開(kāi)。

八、MegaBlocks

8.1 摘要

MegaBlocks（[2211.15841] MegaBlocks: Efficient Sparse Training with Mixture-of-Experts） 是斯坦福大學(xué)、微軟及谷歌聯(lián)合發(fā)布的在 GPU 上高效訓(xùn)練 MoE 的系統(tǒng)。之前我們提到過(guò)，MoE 的 Router 負(fù)載不均衡會(huì)導(dǎo)致需要?jiǎng)h除 Token 或者 Padding 填充，本文中作者采用塊稀疏操作對(duì) MoE 計(jì)算進(jìn)行了重新調(diào)整，并開(kāi)發(fā)了新的塊稀疏 GPU Kernel，以高效處理 MoE 中存在的動(dòng)態(tài)性。作者提出的方法中從不丟棄 Token，并能與現(xiàn)有硬件很好的結(jié)合。

與最先進(jìn)的 Tutel 庫(kù)相比，端到端訓(xùn)練速度提高 40%；與使用高度優(yōu)化的 Megatron-LM 框架訓(xùn)練的 DNN 相比，端到端訓(xùn)練速度提高 2.4x。

PS：需要說(shuō)明的是，MegaBlocks 主要針對(duì)的還是單個(gè) GPU 上包含多個(gè)專家的場(chǎng)景。

8.2 方法

MegaBlocks 主要解決的是 1 個(gè) GPU 上有多個(gè)專家時(shí)，由于負(fù)載不均衡導(dǎo)致的 Token 丟棄或者 Padding 無(wú)效計(jì)算問(wèn)題。如下圖 Figure 3 所示，假設(shè)有 3 個(gè)專家，每個(gè)專家的 Capability 為 2 個(gè) Token，Router 后分配給 3 個(gè)專家的 Token 分別為 3,1,2，因此 Expert-0 需要丟棄一個(gè) Token，Expert-1 需要 Padding 一個(gè) Token。假設(shè) Token Embedding 維度為 1024，F(xiàn)FN 第一個(gè) MLP 升維后為 4096：

• （A）：對(duì)應(yīng) 3 個(gè) (2, 1024) 和 (1024, 4096) 的矩陣乘法，每個(gè)輸出都是 (2, 4096)
• （B）：可以表示為 Batch Gemm 來(lái)計(jì)算，輸出為 (6, 12288)，但只有對(duì)角線上有 3 個(gè) (2, 4096) 的子矩陣，其他位置為 0。采用稀疏計(jì)算不會(huì)增加額外的計(jì)算量。
• （C）：同樣可以表示為 Batch Gemm（可變 Shape），但是不丟棄 Token，也不 Padding，相當(dāng)于 (3, 1024)，(1, 1024) 和 (2, 1024) 的 3 個(gè)矩陣分別不同的 (1024, 4096) 的矩陣相乘，稀疏表示后生成的還是 (6, 12288) 矩陣。PS：這個(gè)圖很容易讓人迷惑，圖中的列分塊是作者想要支持可變大小的專家，但并沒(méi)有實(shí)現(xiàn)。實(shí)際上當(dāng)前用的專家大小都相同，所以各個(gè)專家列分塊的大小也應(yīng)該相同。?

如下圖 Figure 5 所示為對(duì)應(yīng)的稀疏分塊矩陣表示方式：

九、參考鏈接

1. ??https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/p2p.html#??
2. ??https://arxiv.org/abs/1701.06538??
3. ??https://arxiv.org/abs/2006.16668??
4. ??https://arxiv.org/abs/2101.03961??
5. ??https://arxiv.org/abs/2103.13262??
6. ??https://github.com/laekov/fastmoe/blob/master/doc/readme-cn.md??
7. ??https://github.com/laekov/fastmoe/blob/master/cuda/global_exchange.h??
8. ??https://arxiv.org/abs/2206.03382??
9. ??https://arxiv.org/abs/2211.15841??

本文轉(zhuǎn)載自??AI閑談??，作者： AI閑談 ????