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

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

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

1. 使用新一代SparseMixer來精確估計專家路由的梯度，解決傳統方案中利用門控梯度代替路由梯度的問題。

2. 專家并行不要了，訓練中改用數據、pipeline和張量并行，避免了傳統方法丟棄token的問題。

論文地址：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的數學推理能力：

作者注：我們對新發布的GAOKAO（即全國普通大學和學院入學統一考試）的數學問題進行案例研究，這是中國一年一度的全國本科入學考試。

該考試以其嚴格的安全協議而聞名，是評估AI模型回答數學問題的能力的理想測試平臺。請注意，GRIN MoE的訓練于太平洋標準時間6月3日結束，2024年GAOKAO于中國標準時間6月7日開始。