為 MoE 模型優化提供新視角。

原標題：阿里云通義大模型新技術：MoE模型訓練專家平衡的關鍵細節
文章來源：機器之心
內容字數：7053字

機器之心AIxiv專欄：阿里云通義千問團隊解決MoE模型訓練關鍵問題

機器之心AIxiv專欄長期關注并報道全球AI領域的學術前沿，近日報道了阿里云通義千問Qwen團隊的一篇論文，該論文揭示并解決了MoE（混合專家模型）訓練中的一個關鍵問題，并提出了一種顯著提升MoE模型性能和專家特異性的新方法。

1. MoE模型訓練中的關鍵問題：局部負載均衡的局限性

MoE模型通過路由機制動態激活部分模型參數，從而高效地擴展模型規模。然而，基于TopK機制的稀疏激活會導致專家激活不均衡：部分專家被頻繁選擇，而其他專家被閑置，造成資源浪費。為了解決這個問題，通常引入負載均衡損失（LBL）來鼓勵專家激活均衡。然而，現有MoE訓練框架大多實現的是局部（micro-batch）負載均衡，即在每個小批量數據內進行均衡分配。這會導致一個問題：如果小批量數據來自同一領域，則模型會將這些數據均勻分配給所有專家，阻礙專家在特定領域的分化，從而限制模型性能。

2. 從局部均衡到全局均衡：輕量級通信的解決方案

為了克服局部均衡的局限性，阿里云通義千問團隊提出了一種全新的方法：將局部均衡放松為全局均衡。通過輕量級的通信機制，該方法能夠同步所有micro-batch的專家選擇頻率，并在全局范圍內計算負載均衡損失。這種方法的關鍵在于，只需要在各個節點之間同步一個專家數大小的向量，通信開銷極小。即使對于需要梯度積累的情況，通過緩存機制也能有效降低通信成本。

3. 實驗結果：顯著提升模型性能和專家特異性

研究人員在不同規模的模型（3.4B、15B、43B參數）上進行了實驗，結果表明，將均衡范圍從局部擴展到全局后，模型的性能（Benchmark指標和PPL）均得到了顯著提升。實驗還表明，增加全局均衡范圍能持續提升模型效果，并在一定范圍后達到飽和。此外，通過對比實驗，研究人員驗證了全局均衡的提升主要源于在更通用、多樣化的token集合上計算損失，而不是單純的token數量增加。

4. 效率提升：少量局部均衡損失的補充

完全依賴全局均衡可能會降低局部均衡，影響計算效率。因此，研究人員進一步實驗了在主要使用全局均衡的同時，添加少量局部均衡損失。結果表明，這種方法既能提高模型訓練速度，又能保持模型性能。

5. 與現有工作的比較

與現有的一些工作相比，該研究更系統地驗證了全局負載均衡的有效性，詳細分析了均衡范圍對性能的影響，并通過消融實驗證明了全局均衡通過納入更多樣化的領域信息來提升性能。

6. 結論

阿里云通義千問團隊的研究解決了MoE模型訓練中的一個關鍵問題，通過輕量級通信實現全局負載均衡，顯著提升了MoE模型的性能和專家特異性。這項工作為MoE模型的優化提供了新的視角，并有助于構建更大規模、更有效的MoE模型。

聯系作者

文章來源：機器之心
作者微信：
作者簡介：專業的人工智能媒體和產業服務平臺