一、背景

之前的文章中我們詳細介紹了大規(guī)模分布式訓練中的數(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 標準 All2All

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

如下圖所示是一種標準的 AlltoAll 操作，有 4 個 GPU，每個 GPU 包含 4 個數(shù)據(jù)。通過 AlltoAll 操作之后每個設備都將 4 個數(shù)據(jù)分發(fā)給 4 個 GPU，同時從 4 個 GPU 接收數(shù)據(jù)?？梢钥闯?，AlltoAll 很像一個矩陣的轉(zhuǎn)置操作：

如下圖所示為 Pytorch 實現(xiàn)一個上述標準 AlltoAll 的示例：

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

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

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

2.2 非標準 All2All

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

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

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

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

如下圖所示為 Pytorch 實現(xiàn)一個上述非標準 all2all 的示例：

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

2.3 兩次 All2All

上述非標準 All2All 中有個問題：有些時候當前設備只知道要向其他設備發(fā)送多少數(shù)據(jù)，而并不知道需要從其他設備接收多少數(shù)據(jù)。這個問題可以通過 2 次 all2all 來解決：

• 第一次 all2all 交換要傳輸?shù)臄?shù)據(jù)量信息，這是一個標準的 all2all 操作，如下圖紅框所示。
• 第二次 all2all 根據(jù)上述獲取的數(shù)據(jù)量信息來執(zhí)行真正的數(shù)據(jù)傳輸，此時是一個非標準 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ù)千個前饋子網(wǎng)絡，并由可訓練的門控網(wǎng)絡（Gating Network）確定這些專家的稀疏組合。

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

3.2 方案

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

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

3.3 Expert Parallelism

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

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

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

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

四、Google Gshard

4.1 摘要

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

4.2 方案

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

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

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

Random routing：有些時候 Gating 模塊得到的排名第二的專家的分數(shù)會很低，此時可以簡單的忽略第二個專家。

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

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

五、Switch Transformer

5.1 摘要

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

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

5.2 簡化稀疏路由

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

• Router 計算更簡單，通信量也更少。
• 一個 Token 僅對應一個專家，計算量也更少。
• 平均每個專家對應的batch size至少可以減半。

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

5.3 高效稀疏路由

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

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

• 容量因子為 1.0：如下圖中所示，則對應的專家容量為 2：

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

Expert 2 有 2 個 Token，正好。

Expert 3只有 1 個 Token，需要Padding 1 個空的 Token。

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

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

六、FastMoE

6.1 摘要

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

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

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

6.2 系統(tǒng)設計

6.2.1 靈活性

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

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

6.2.2 擴展模型容量

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

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

異構(gòu)同步模塊。模型的不同部分可能在不同的 Worker 組間重復，這非常有挑戰(zhàn)，因為分布式模塊不得不識別是否需要對參數(shù)的梯度進行同步，以及與誰同步。因此，F(xiàn)astMoE 引入了數(shù)據(jù)并行通信組標簽：

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

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

6.3 優(yōu)化激活

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

6.4 多 CUDA Stream 調(diào)度

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

6.5 MoE 代碼實現(xiàn)

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

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

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

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

七、Tutel

7.1 摘要

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

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

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

7.2 自適應 MoE

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

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

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

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

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

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

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

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

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

7.3 優(yōu)化

7.3.1 Flexible All-to-All

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

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

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)存空間對齊到連續(xù)內(nèi)存空間，并將多個小的通信合并成大的通信：

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

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

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

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

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

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

7.3.4 Adaptive Pipelining

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

八、MegaBlocks

8.1 摘要

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

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

PS：需要說明的是，MegaBlocks 主要針對的還是單個 GPU 上包含多個專家的場景。

8.2 方法

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

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

如下圖 Figure 5 所示為對應的稀疏分塊矩陣表示方式：

九、參考鏈接

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閑談 ????