LLM 學習筆記-Deepspeed-MoE 論文

論文 DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale

1. Introduction

現有的 MoE 方法在正式使用場景中存在的挑戰：

1. 場景侷限：大都是 encoder-decoder 模型或者 sequence-to-sequence 任務；
2. 訓練時的內存需求巨大：
3. 推理性能還不太行：通常單個 GPU 放不下 MoE 做推理。另一方面多 GPU 的 MoE 推理方法還欠缺研究。MoE 推理還收到內存帶寬的影響。

Deepspeed-MoE針對上述挑戰做了下面的改進：

1. 把 MoE 的任務擴展到了各種自迴歸的 NLG 任務
2. 提出 PR-MoE 來減少 MoE 參數
3. 設計了 Deepspeed-MoE 推理系統，減少 7.3 倍推理延時和開銷。

2. Related work

• 先講了 large-scale 的 dense model
• 然後引出 moe
• 最後介紹了一下現有的 MoE 訓練和推理優化系統
  • 各種並行策略：dp, pp, zero, ep
  • fastmoe, tutel

3. 將 MoE 擴展到其他下游任務

MoE 模型是基於 dense 模型設計的，即在 dense 模型的基礎上把單個 MLP 替換成 多個 MLP，再加上 gate 等模塊。例如 1.3B+MoE-128 就是基於參數量爲 1.3B 的dense模型，每一層 MLP 替換成 128 個 expert 組成的模塊。

擴展的任務包括：

• completion prediction task (LAMBADA [16])
• common sense reasoning task (PIQA [35])
• 2 個 reading comprehension tasks (BoolQ [18], RACEh [36])
• 2 個 question answering tasks (TriviaQA [21], WebQs [20])

實驗結果不贅述了，反正MoE贏過dense，而且開銷更小就完事了。

4. PR-MoE & MoS：減少模型大小，提高parameter efficiency

PR-MoE （Pyramid Residual MoE）結構如下

PR-MoE 的是基於兩個現象設計的：

1. 現有的 MoE 每一層專家模塊的 expert 數量都是一樣，所以做了個對比試驗：實驗 A 指在模型前半部分添加專家模塊，實驗 B 則只在後半部分添加專家模塊。實驗結果顯示 B 表現更高，所以說明模型深層（即後半部分）更需要專家模塊；所以最終 PR-MoE 後半部分的專家數量會設置的更多一些。

2. 要提高 MoE 的泛化能力，有兩種方法：

   1. 增加 expert 數量，固定 expert capacity：這個的問題在於會增加闡述了和內存開銷
   2. 固定 expert 數量，增加 expert capacity：這個的問題在於會增加通信量

   上面提到的 expert capacity 可以理解成每次選 top-k 個 expert 參與計算，K 越大，expert capacity 就越大，通常認爲這樣模型泛化能力越強。這是因爲我們會認爲不同 expert 彼此之間會做一個糾正，不至於預測結果偏差太大。所以一種思路是固定一個 expert，然後再動態選擇一個 expert，這樣就可以既減少了通信量（因爲只有一個 expert 需要選擇），又保證了 capacity。

MoE 訓練加速的方法是專家並行，最簡單和最高效的辦法是每個 GPU 上均等平分 expert。例如總共有 64 個 expert，然後這個時候使用 64 個 GPU 的話，那就只需要 給每個 GPU 分配一個 expert就可以了。如果只有 32 個 GPU，那就每個 GPU 分配 2 個 expert。

但是現在PR-MoE 的訓練的難點在於不同層 expert 的數量不一樣了。其實，如果 GPU 數量少於最少的專家數量的話，還挺好劃分的，簡單平均分就好了。但是這種情況呢？假設有三層，每層專家數量分別是 {32, 64, 128}，GPU 數量是 128。按照原論文的說法，deepspeed-moe 支持這樣的並行方式：

“a PR-MoE model running on 128 GPUs, with 32, 64, and 128 experts at different MoE layers, can be trained with 128-way data parallelism for the non-expert parallelism, and {32, 64, 128} expert parallelism plus {4, 2, 1} data parallelism for MoE parameters.”

簡單來說，就是非 expert 模塊採用 128 路的數據並行；但是對於專家模塊，這才用專家並行和數據並行相結合的策略。例如第一層只有 32 個專家，那麼我們就把第一層複製 4 遍，即對第一層做 4 路數據並行，每一路內部會做 32 路的專家並行。其他同理。

爲進一步降低參數量，論文還用知識蒸餾得到更小的 MoE 模型，稱作 Mixture of Students (MoS)。

5. 推理優化

我們首先對於 1.3B+MoE-128 模型（總參數量是接近 52B）考慮兩種推理的極端情況：

1. 只有一個輸入的 token 數據做推理，那這個時候其實只會激活一條路徑，所以只需要使用 1.3B 對應的內存開銷
2. 當輸入一個大 batch 的數據時，很可能每個 expert 都被激活了，那這個時候需要的是 52B 對應的內存開銷

換言之，模型推理的開銷介於 1.3～52B 之間。

Deepspeed-MoE 的優化思路是下面 3 條：

1. 支持各種並行策略混合，而且把路徑相同的 token 打包一起做推理，這樣每次就只需要激活一個路徑了，即單個 dense 模型的開銷。
2. 優化通信scheduling
3. 優化 transformer和 MoE 相關的 kernel

上面 3 條優化思路詳細介紹如下：

5.1 靈活組合多種並行策略：TP，EP，DP

簡單理解就多種並行方式混合使用。下圖是以單個 MoE 層爲例解釋瞭如何使用混合並行策略，假設總共有16 個 GPU，專家數量爲 8 個，可以看到並行模式是針對專家和非專家模塊分別設計的，具體如下：

• 對於非 expert 參數模塊：
  • 使用4 路數據並行，也就是說該部分參數複製了 4 遍
  • 每一路數據並行的內部採用 4 路的tensor 並行
• 對於 expert 參數模塊：
  • 使用 8 路專家並行
  • 每個專家通過 2 路的 tensor 並行進行參數拆分

5.2 優化通信

我們知道專家並行需要把專家分配到不同的 GPU 上，因此計算是時需要用到 all-to-all 通信來互相交換 token。論文作者發現使用基於 NCCL 的 torch.distributed 接口會有很大的 overhead，尤其是當規模上去的時候。於是他們改用了微軟的 SCCL，實驗結果顯示比 NCCL 更好。作者還設計了兩個新的通信優化策略：

5.2.1 Hierachical all-to-all

簡單來說，分層 AlltoAll 設計的工作原理如下：

1. 首先，每個 GPU 將其數據進行本地轉換。這可以是任何類型的轉換，例如將數據排序或壓縮。
2. 然後，每個 GPU 將其轉換後的數據與同一節點上的其他 GPU 進行交換。這可以理解成是局部 AlltoAll。
3. 最後，每個節點將其轉換後的數據與其他節點進行交換。這可以理解成全局 AlltoAll。

假設總共有 \(p\) 個 GPU，每個節點包含 \(G\) 個GPU，所以總共有 \(p/G\) 個節點。複雜度對比：

• 原本直接所有 GPU 之間做 all-to-all 的複雜度是 \(O(p)\)
• 現在分層 all-to-all 的複雜度爲 \(O(G+p/G)\)包括兩部分：
  • 局部通信複雜度：因爲不同節點之間可以同時做通信，所以複雜度是\(O(G)\)
  • 全局通信複雜度：其實就是節點之間做通信，所以複雜度是 \(O(p/G)\)

5.2.2 Tensor 和 expert 協同並行的通信優化

在插入 MoE 模塊後的 transformer layer 通常長這樣：inputs → MLP1 → MoE layer

爲了加速我們可以對 MLP 層分別使用 tensor 並行，對 MoE layer做專家並行，這是前提。

• 我們先考慮在 MLP1 上執行 tensor 並行：我們知道 tensor 並行的特點是參與計算的 GPU 的輸入數據是相同的，輸出數據也是相同的（column parallel 需要用到 all-gather, row parallel 需要用到 all-reduce）。
• 接着在 MoE 上做專家並行：正如前面提到的，MLP1 最後在不同設備上的輸出結果會保持一致，那麼也就是說對於後面的 MoE layer 而言，同一個 tensor parallel group 的進程之間因爲他們的輸入數據是一樣的了，所以就不需要再做all-to-all 通信了，也就是說 GPU0 和 GPU1 沒必要在浪費時間做 all-to-all 了，同理 GPU2 和 GPU3 也不需要了。那換言之，爲了節省通信開銷，我們可以先在 GPU0 和 GPU2 之間做 all-to-all 通信，結束之後， 同一個tensor parallel group內部的 GPU 再做 all-gather即可，即 GPU0 把數據同步給 GPU1。

小結：這樣一來，原本 all-to-all 的通信開銷就從 \(O(p)\)降到了 \(O(L)+O(p/L)\)，其中 p 表示總共的 GPU 數量， L 表示tensor parallelism degree。對應到下面圖中 p=4，L=2。當總的 GPU 數量增大的時候，這種方式的通行效率才能體現出來

5.3 MoE-related Kernel優化

基於前面的介紹我們知道 MoE 的核心主要是

• gate 模塊
• 各種 data layout transformation，例如：
  • 稀疏張量 one-hot vector，用來指示每個 token 分配給了哪個 expert
  • Sparse-Dense Einsum Operations，比如是根據上面onehot vector 和 inputs tokens 進行矩陣乘法，實現 token 的 masked 計算
  • 數據的重新排序和重排（Reordering and Rearranging of Data）：MoE 模型中的最終 einsum 運算會對 tokens 進行縮放和重新排序，以使其回到原始的順序。這個過程可能會涉及對數據進行重新排列和重組。

gate 模塊包含了很多算子，例如生成 token 的 mask，top-k 操作，scatter token 到對應 expert等等，這些算子都需要啓動 kernel call,而每一次都需要開銷。所以論文 的優化思路是兩點：

• 使用 dense representation代替稀疏張量表示，簡單理解就是弄了一個 mapping table，記錄每個 token和 expert 的對應關係
• 把 gate 模塊內的所有算子融合成了單個 kernel 算子

實驗結果顯示優化後的 MoE kernel 能夠降低 6 被的延遲。

疑問

整篇文章主要是第 5 章讀起來費勁，讀完還是有些問題沒有搞明白，希望和大家討論一下：

• 圖 8 和圖 9 中的小方格分別表示什麼含義，是模型參數還是 input token？
• all-to-all 具體是指下面的哪一種，因爲看到了不同的說法：
  1. 參與通信的 GPU 把自身所有數據都同步給其他 GPU，然後每個 GPU 通過 mask 執行計算
  2. 參與通信的 GPU 根據 gate 的結果只傳輸需要的 token 給對應的 GPU
• 爲什麼圖 8 和圖 9 的 local transform 都是把 GPU 0 的 B 和 C 做交換，可以把 A 和 C 交換嗎？
• 每個 GPU 都只傳輸了 3/4 的數據給其他 GPU，這是爲什麼？另外我理解的是不同的 expert 處理的數據量應該是不一樣的

參考：

• https://zhuanlan.zhihu.com/p/466363675
• https://www.cnblogs.com/marsggbo/p/16871789.html

微信公衆號：AutoML機器學習
MARSGGBO♥原創
如有意合作或學術討論歡迎私戳聯繫~
郵箱:[email protected]