為 MoE 模型優(yōu)化提供新視角。

原標(biāo)題：阿里云通義大模型新技術(shù)：MoE模型訓(xùn)練專家平衡的關(guān)鍵細(xì)節(jié)
文章來源：機(jī)器之心
內(nèi)容字?jǐn)?shù)：7053字

機(jī)器之心AIxiv專欄：阿里云通義千問團(tuán)隊(duì)解決MoE模型訓(xùn)練關(guān)鍵問題

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

1. MoE模型訓(xùn)練中的關(guān)鍵問題：局部負(fù)載均衡的局限性

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

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

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

3. 實(shí)驗(yàn)結(jié)果：顯著提升模型性能和專家特異性

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

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

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

5. 與現(xiàn)有工作的比較

與現(xiàn)有的一些工作相比，該研究更系統(tǒng)地驗(yàn)證了全局負(fù)載均衡的有效性，詳細(xì)分析了均衡范圍對性能的影響，并通過消融實(shí)驗(yàn)證明了全局均衡通過納入更多樣化的領(lǐng)域信息來提升性能。

6. 結(jié)論

阿里云通義千問團(tuán)隊(duì)的研究解決了MoE模型訓(xùn)練中的一個關(guān)鍵問題，通過輕量級通信實(shí)現(xiàn)全局負(fù)載均衡，顯著提升了MoE模型的性能和專家特異性。這項(xiàng)工作為MoE模型的優(yōu)化提供了新的視角，并有助于構(gòu)建更大規(guī)模、更有效的MoE模型。

聯(lián)系作者

文章來源：機(jī)器之心
作者微信：
作者簡介：專業(yè)的人工智能媒體和產(chǎn)業(yè)服務(wù)平臺