專家模型不要專家並行！微軟開源MoE新路徑

新智元報道

編輯：alan

【新智元導讀】近日，來自微軟的研究人員開源了使用全新方法訓練的MoE大模型，不走尋常路，且編碼和數學表現出色。

繼Phi家族之後，微軟又開源了新的混合專家大模型——GRIN MoE。

與Phi-3.5同樣的個頭（16 * 3.8B），卻採用了截然不同的訓練方法。

這個「不走尋常路」如果寫個太長不看版，那就是兩句話：

論文地址：https://arxiv.org/abs/2409.12136

當然了，上面兩句話是小編說的，多少有點糙，文中細節，還請諸君繼續閱讀~

這年頭，新來一個LLM，當然要先刷分了——

參數要少，效果要好，所以要在左上角：

GRIN作爲MoE架構，總參數量約42B，推理時激活的參數爲6.6B，打同級別（7B）的非MoE模型是手拿把攥，甚至比14B的Phi-3還要略勝一籌。

在上面的這份成績單中，GRIN MoE表現優異，尤其是在編碼和數學測試中。

比如，在衡量數學問題解決能力的GSM-8K中，GRIN MoE得分爲90.4，而在編碼任務基準HumanEval上拿到了74.4分。

在MMLU（大規模多任務語言理解）基準測試中GRIN得分爲79.4，超過了同爲MoE架構的Mixtral（70.5分），以及自家的Phi-3.5（78.9分）。

如果對比流行的商用模型，GPT-3.5表示感受到時代的力量，默默退出羣聊。

開放權重：https://huggingface.co/microsoft/GRIN-MoE

demo：https://github.com/microsoft/GRIN-MoE

MoE全新訓練路徑

GRIN MoE由常規的Transformer塊構成，採用分組查詢注意力（GQA）和滑動窗口注意力來提高計算效率。

採用RoPE進行位置編碼，以便在預訓練後實現長上下文能力。

在MoE架構中，模型通過路由網絡爲每個輸入token挑選適合的專家模塊。對於有n個專家的網絡，一個用於推理的MoE模塊的輸出爲：

其中z = Router（x，r），本文中Router採用線性網絡，Gating是門控函數（通常爲softmax），Expert是FNN層。

MoE通過TopK函數進行專家分配，這個專家路由的過程是不可微的，所以反向傳播的時候沒法求導。

對此，傳統的MoE訓練將TopK視爲常數，僅通過Gating來反向傳播計算路由權重梯度，相當於用門控的梯度代替了路由的梯度。

這多少有點糙。

不可導怎麼辦

恰好，本文一作之前有一篇工作（SparseMixer）：

論文地址：https://arxiv.org/pdf/2310.00811

受到直通梯度估計器的啓發，作者擴展了前作，提出了SparseMixer-v2。

作者首先將TopK函數替換爲模型訓練中離散變量的隨機採樣，然後應用heun’s third order method來近似專家路由梯度，並構建一個改進的反向傳播，爲專家路由給出數學上合理的梯度估計。

前作中，SparseMixer的有效性在神經機器翻譯任務和ELECTRA語言模型訓練中得到了證明。

而在GRIN MoE的開發過程中，SparseMixer-v2終於有機會大規模應用於自迴歸語言模型訓練。

作者用2.5T token訓練了兩個16×0.9B MoE。其中一個遵循GRIN MoE中使用的相同方案，另一個用傳統的GShard方法替換 SparseMixer-v2。

如上圖所示，將SparseMixer-v2的性能提升推廣到16×0.9B尺度的自迴歸語言模型訓練。

在前0.5T token上GShard表現更好，但SparseMixer-v2在訓練後期取得了更強的性能。

專家模型不要專家並行

傳統的MoE訓練採用專家並行，簡單理解就是把不同的專家分配到不同的顯卡上。

一個明顯的問題是負載不均衡，有的專家會分到更多的token，有的專家卻很閒。

之前的做法是設定一個閾值，比如1000個token分給4個專家，每人應該是250，這時候每張卡就最多隻算250個token，超過後直接丟棄（送到下一層）。

而在本文中，作者利用數據並行、pipeline並行和張量並行來訓練GRIN MoE。

此外，對於沒有專家並行性的MoE計算，作者發現Megablocks包非常有用，它的grouped_GEMM內核和包裝器的性能更好。

應用這些新的工程化方法避免了專家並行，也就不用丟棄token了。

最終，與具有相同激活參數的密集模型相比，本文的方法實現了超過80%的訓練效率提升。

上表中，作者將兩種不同大小的MoE模型與具有相同激活參數量的密集模型進行了比較，使用相同的硬件測量了它們的訓練吞吐量。

儘管MoE總的參數量是密集模型的六倍多，但在實驗中達到了超過80%的相對吞吐量，證實了使用GRIN MoE方法的模型具有顯著的計算擴展潛力。

（PS：密集模型的吞吐量是在與MoE模型相同的並行度設置下測量的，這裡的比較是爲了研究密集激活網絡（非MoE）和稀疏激活網絡（MoE）的GPU內核效率）

此外，在擴大模型大小時，密集模型和MoE模型顯示出相似的減速模式，比如6.6B密集模型的訓練吞吐量大約比1.6B密集模型的訓練吞吐量慢4.19倍（後者的參數少4倍）。同樣，42B MoE模型的訓練吞吐量比10B MoE 模型的訓練吞吐量慢約3.96倍（對應參數少4.2倍）。

並行實驗

在只使用pipeline並行的情況下，通過在GPU之間進一步劃分不同層，可以將最大專家數量從16個擴展到32個。但是，如果再增加專家數量，則會導致單個層的參數過多，一個GPU就放不下了。

所以下一個維度採用張量並行。

專家並行在前向和後向計算中有兩個all-to-all通信開銷，而張量並行在前向和後向計算中有兩個all-reduce通信開銷。

相比之下all-reduce操作的延遲更高一點，但可以通過精心排布前向和反向的計算來overlap掉一部分開銷。

如上圖所示，通過結合pipeline並行和張量並行，系統支持的最大專家數量擴展到52個（總共132B參數）。

這個數量是因爲實驗只用了64個GPU，最多能將模型劃分爲64個階段，如果有更多的GPU，那麼還能繼續向上擴展。

不過作者也表示，使用更復雜的並行通常會導致計算吞吐量降低。

負載均衡

如前所述，本文沒有采用專家並行，但是負載不均衡的事實依然存在。

作者在這裡通過調整負載均衡損失來調節全局的負載均衡。常見的負載均衡損失定義爲：

其中α是超參數，n是專家數量，fi是調度給專家的token比例。

傳統方法在本地不同的GPU上計算fi，因此負載均衡損失將調節本地專家負載均衡並緩解token丟棄。

在本文中，作者通過計算全局的fi（比如數據並行過程中組內的all-reduce）來修改負載均衡損失，調節專家負載以達到全局平衡。

儘管這種調整會產生額外的通信開銷，但類似於張量並行，這些通信也可以與計算overlap，從而在很大程度上減少額外的延遲。

最後，放一個測試結果來show一下GRIN MoE的數學推理能力：

參考資料：

https://venturebeat.com/ai/microsofts-grin-moe-ai-model-takes-on-coding-and-math-beating-competitors-in-key-benchmarks/