【元來如此】第四章——MOE專家并行再升級
文 | 軟件生態中心-模型應用部-思成
正文2000字 閱讀5分鐘
“
在之前的文章中，我們介紹了MoE模型推理相關的探索內容，還沒有了解的同學可以點擊下方圖片鏈接跳轉閱讀。接下來為大家帶來MoE大模型并行訓練的內容。
MoE架構
▲圖0 Switch Transformers[1]論文中的MoE Layer
如圖0所示，MoE網絡主要的改進是將Transformer模型中的每個FFN層替換為MoE結構，所謂的MoE結構主要由一個Router和N個Expert組成，其中Router負責將Token路由到對應的Expert。之后在對應的專家完成對應的計算操作，最后將多個專家的結果通過Router路由概率計算加權和。
后續幾年，陸續有很多關于MoE的研究。比如對于上述提到的Router，演進出了很多Token- Choice、Expert-Choice、Token Merge、Expert Merge算法。針對Expert網絡結構，有在FNN和Attention兩個主要方面的工作，比較有代表性的比如GShard和MoA。針對Share Expert算法衍生出了比如DeepseekMoE、Qwen1.5-MoE等方面的工作。
本篇內容會從基礎的MoE結構入手，更多的集中在分布式方案設計上，并不會過多關注MoE相關算法相關內容（如果有感興趣的同學，可以額外開章節進行介紹）。
下文內容我們做如下約定和假設：
各層MoE網絡結構是同構的，圖示中只展示其中一層網絡
為方便分析，我們圖中只展示Top1-Expert的情況
各個Router是負載均衡的，且不考慮溢出或者丟棄等
同時定義以下符號：
DP(Data Parallelism)
TP(Tensor Parallelism)
EP(Expert Parallelism)
方案0
從中易得，前向計算過程中Attention部分不存在通信，在MoE部分需要每一個DP維度內的EP4之間的All2All操作，經過MoE的計算之后，再進行一次All2All操作還原數據。 在模型反向計算過程中GPU0-GPU4，GPU1-GPU5，GPU2-GPU6，GPU3-GPU7的Expert做Allreduce操作，在Attention部分做8卡的Allreduce操作。
▲ 圖1 Attention-DP8-MoE-EP4-DP2
在這樣的方案設計中，我們不難得到這樣的結論：非MoE部分，即Attention部分在DP維度完成Allreduce操作。在MoE部分，專家之間在EP維度完成All2All操作，在DP維度完成Allreduce操作。
方案1
接下來讓我們稍微進階一下。假設方案0中GPU卡的顯存并不夠。其中一種方案是將Expert維度改變成EP4-TP2，如圖2所示。單機8卡，4個專家，Attention部分（藍色）采用DP8，MoE部分采用EP4-TP2。
可以非常直觀的看到，前向過程中Attention部分保持不變。MoE部分（為展示清晰，圖中只畫出了GPU3中的數據在MoE結構中的流程，其余7張卡同理）首先進行了一次變種的All2All通信，經過MoE計算之后，通過另一個正常的All2All操作還原數據，注意，經過還原的數據需要在各自的rank內再額外完成一次reduce操作，當然這個reduce操作并不產生額外的通信。
反向計算過程中，Attention分布保持8卡做Allreduce操作。MoE部分沒有額外的通信。
▲ 圖2 Attention-DP8-MoE-EP4-TP2
在這個方案中。不管是哪個All2All操作，對比方案1，通信量都增加了1倍，同時這個All2All操作會的范圍擴大了1倍，這種通信模式在非兩兩互聯的link中對性能的挑戰都會進一步增大。另外，在第一次All2All的時候存在一個變種的All2All操作（雖然通過一些trick方案可以繞過，感興趣的同學可以想想看）對于代碼可讀性和維護性上也存在一定的破壞。
方案2
在前向計算的時候， Attention部分和之前2個方案不同，GPU0-GPU1，GPU2-GPU3，GPU4-GPU5，GPU6-GPU7之間需要進行Allreduce通信。之后看起來在MoE部分完成2次All2All操作。細心的同學可能已經發現，如果按照常規的All2All操作，第一次會多發送1倍的重復數據。第二次進行的是一個類似與上一個方案的All2All操作。 在反向計算的時候，GPU0-GPU4，GPU1-GPU5，GPU2-GPU6，GPU3-GPU7的Expert做Allreduce操作，在Attention部分做GPU0-GPU2-GPU4-GPU6和GPU1-GPU3-GPU5-GPU7的Allreduce操作。
▲ 圖3 Attention-TP2-DP4-MoE-EP4-DP2
▲ 圖4 Attention-TP2-DP4-MoE-EP4-DP2
▲ 圖5 Attention-TP2-DP4-MoE-EP4-DP2
方案3
前向過程中，Attention部分和之前介紹的一致。在MoE部分，首先需要一次All2All（2組，分別是GPU0-GPU2-GPU4-GPU6和GPU1-GPU3-GPU5-GPU7）操作。之后在GPU0-GPU1、GPU2-GPU3等各自完成Allreduce計算，最后再通過一次All2All操作（同樣也是2組）。 這里不進一步贅述有關反向的操作流程了，相信聰明的你一定可以根據上面幾種情況的學習推理出這里發生了什么。
▲ 圖6 Attention-TP2-DP4-MoE-EP4-TP2
從上述分析來看，為了高效訓練一個確定的MoE模型，針對不同的硬件特點，設計貼合硬件特點的分布式訓練方案是一種可以有效提升性能的手段。我們以上一篇提到的Mixtral-8x7B模型為例，可以得到單層Attention參數量42million，而8個Expert對應的參數量為1409million，兩者參數差異在33倍以上。從顯存容量角度看，盡可能的分割專家就已經可以很好的釋放顯存壓力，非必要的時候并不需要額外切分Attention部分的參數。
免費試用申請
如果您對我們的產品感興趣，
可點擊底部“閱讀原文”，進行試用申請。
參考文獻
[1] https://arxiv.org/abs/2101.03961
