開源賽道太擠了！月之暗面開源新版Muon優化器

作者：機器之心 2025-02-24 13:50:00

人工智能開源

省一半算力跑出2倍效果，月之暗面開源優化器Muon，同預算下全面領先。

月之暗面和 DeepSeek 這次又「撞車」了。

上次是論文，兩家幾乎前后腳放出改進版的注意力機制，可參考《撞車 DeepSeek NSA，Kimi 楊植麟署名的新注意力架構 MoBA 發布，代碼也公開》、《剛剛！DeepSeek 梁文鋒親自掛名，公開新注意力架構 NSA》。

這次是開源。

上周五，DeepSeek 剛剛官宣這周要連續開源 5 個代碼庫，卻被月之暗面深夜截胡了。

昨天，月之暗面搶先一步開源了改進版 Muon 優化器，比 AdamW 優化器計算效率提升了 2 倍。

團隊人員表示，原始 Muon 優化器在訓練小型語言模型方面表現出色，但其在擴展到更大模型方面的可行性尚未得到證明。因此，團隊人員確定了兩種對擴展 Muon 至關重要的技術：

添加權重衰減：對擴展到更大模型至關重要。
一致的 RMS 更新：在模型更新上執行一致的均方根。

這些技術使得 Muon 能夠在大規模訓練中直接使用，而無需調整超參數。Scaling law 實驗表明，與計算最優訓練的 AdamW 相比，Muon 的計算效率提升了 2 倍。

基于這些改進，月之暗面推出了 Moonlight，這是一個 3B/16B 參數的 Mixture-of-Expert（MoE）模型，使用 Muon 進行了 5.7 萬億 tokens 的訓練。該模型刷新了當前的「帕累托前沿」，換句話說，在相同的訓練預算下，沒有其他模型能在所有性能指標上同時超越它。

與之前的模型相比，Moonlight 也以更少的訓練 FLOPs 獲得了更好的性能。

如下圖所示，該研究進行了 Scaling law 研究，將 Muon 與強大的 AdamW 基線進行了比較，結果展示了 Muon 的卓越性能。Muon 實現了與 AdamW 訓練相當的性能，同時僅需要大約 52% 的訓練 FLOP。

月之暗面不但開源了內存優化且通信高效的 Muon 實現代碼，并且還發布了預訓練、指令調優以及中間檢查點，以支持未來的研究。

論文《 MUON IS SCALABLE FOR LLM TRAINING 》。

論文地址：https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
代碼地址：https://github.com/MoonshotAI/Moonlight
模型地址：https://huggingface.co/moonshotai/Moonlight-16B-A3B

研究介紹

擴展 Muon

Muon 優化器為 OpenAI 研究者 Keller Jordan 等人在 2024 年提出的，他們的研究表明在小規模訓練中 Muon 的表現顯著優于 AdamW。

但月之暗面發現，當將其擴展到訓練更大模型并使用更多 token 時，模型性能提升逐漸減弱。他們觀察到，權重和層輸出的 RMS 值持續增長，最終超出了 bf16 的高精度范圍，這可能會損害模型的性能。

為了解決這個問題，月之暗面在 Muon 中引入了標準的 AdamW（Loshchilov 等人，2019）權重衰減機制。

為了探究這一機制，研究者對 Muon 進行了有無權重衰減的實驗，他們訓練了一個包含 800M 參數和 100B token（約為最優訓練 token 量的 5 倍）的模型。圖 2 展示了使用 AdamW、無權重衰減的原始 Muon 以及帶權重衰減的 Muon 訓練的模型的驗證損失曲線。

結果顯示，雖然原始 Muon 在初期收斂速度更快，但一些模型權重隨時間的推移增長過大，可能會影響模型的長期性能。

加入權重衰減后解決了這一問題 —— 結果表明，帶權重衰減的 Muon 優于原始 Muon 和 AdamW，獲得了更低的驗證損失。公式 3 為表達式，其中 λ 為權重衰減比率。

一致的 RMS 更新。研究者發現 Adam 和 AdamW 的一個重要特性是，它們將更新的 RMS 維持在 1 左右。然而，月之暗面發現 Muon 更新 RMS 會根據參數矩陣形狀的變化而變化，具體如下引理 1 所示：

為了在不同形狀矩陣之間保持一致的 RMS 更新，該研究通過來擴展 Muon 矩陣更新，從而抵消引理 1 中提到的影響。

在實際應用中，研究者通常將 AdamW 與 Muon 結合使用，以處理非矩陣參數。本文希望優化器超參數（學習率 η、權重衰減 λ）能夠在矩陣參數和非矩陣參數之間共享。

因此他們提出將 Muon 更新的 RMS 調整到與 AdamW 相似的范圍。他們通過以下調整將 Muon 更新 RMS 縮放至這一范圍：

分布式 Muon

月之暗面團隊還提出了一種基于 ZeRO-1 的分布式解決方案，稱為分布式 Muon（Distributed Muon）。分布式 Muon 遵循 ZeRO-1 在數據并行（DP）上對優化器狀態進行劃分，并與普通的 ZeRO-1 AdamW 優化器相比引入了兩個額外的操作，算法 1 描述了分布式 Muon 的實現。

實驗

RMS 的一致性

為了使所有矩陣參數更新的 RMS 值與 AdamW 的 RMS 保持一致，研究團隊嘗試了兩種方法來控制參數更新的 RMS，并將其與只用了 AdamW 的基線的 RMS 進行了對比。

由于大規模訓練模型時，會出現各種意料之外的情況，因此，研究團隊測試了 Muon 對訓練早期階段的影響。當矩陣維度差異增大時，更新 RMS 不一致的問題會更加明顯。該團隊對模型架構進行了微調，用標準的 2 層 MLP 替換了 Swiglu MLP，并將其矩陣參數的形狀從 [H, 2.6H] 改為 [H, 4H]。

團隊評估了模型的損失，并監控了關鍵參數的 RMS，尤其是形狀為 [H, H] 的注意力查詢權重和形狀為 [H, 4H] 的 MLP 權重。

實驗結果表明（見表 1），Update Norm 和 Adjusted LR 均優于基線方法，且 Adjusted LR 的計算成本更低，因此被選用于后續實驗。

Muon 的 Scaling Law

為了與 AdamW 公平比較，該團隊在一系列基于 Llama 架構的模型上對 Muon 進行了拓展。

對于 Muon，由于其 RMS 與 AdamW 匹配，團隊直接復用了 AdamW 的超參數。

實驗結果顯示，擬合的 Scaling Law 曲線表明，在計算最優設置下，Muon 僅需約 52% 的訓練 FLOPs 即可達到與 AdamW 相當的性能。這進一步說明了 Muon 在大規模語言模型訓練中的高效性。

使用 Muon 進行預訓練

為了評估 Muon 在模型架構中的表現，該團隊使用 DeepSeek-V3-Small 架構從頭開始預訓練了 Moonlight 模型。

Moonlight 模型總共進行了 5.7 萬億 tokens 的訓練，但在訓練到 1.2 萬億 tokens 的階段，團隊將其與 DeepSeek-V3-Small（使用 1.33T tokens 訓練的 2.4B/16B 參數 MoE 模型）和 Moonlight-A（與 Moonlight 設置相同，但使用 AdamW 優化器）進行了比較。如表 4 所示，Moonlight 在語言、數學和編碼等任務上都顯著優于 Moonlight-A，證明了 Muon 的擴展優勢。