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