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