開源賽道太擠了！月之暗面開源新版Muon最佳化器
机器之心發表於2025-02-24
省一半算力跑出2倍效果，月之暗面開源最佳化器Muon，同預算下全面領先。
上次是論文，兩家幾乎前後腳放出改進版的注意力機制，可參考《撞車 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 最佳化器為 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 縮放至這一範圍：月之暗麵糰隊還提出了一種基於 ZeRO-1 的分散式解決方案，稱為分散式 Muon（Distributed Muon）。分散式 Muon 遵循 ZeRO-1 在資料並行（DP）上對最佳化器狀態進行劃分，並與普通的 ZeRO-1 AdamW 最佳化器相比引入了兩個額外的操作，演算法 1 描述了分散式 Muon 的實現。為了使所有矩陣引數更新的 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 的計算成本更低，因此被選用於後續實驗。為了與 AdamW 公平比較，該團隊在一系列基於 Llama 架構的模型上對 Muon 進行了擴充。對於 Muon，由於其 RMS 與 AdamW 匹配，團隊直接複用了 AdamW 的超引數。實驗結果顯示，擬合的 Scaling Law 曲線表明，在計算最優設定下，Muon 僅需約 52% 的訓練 FLOPs 即可達到與 AdamW 相當的效能。這進一步說明了 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 的擴充套件優勢。在完整訓練後，Moonlight 與類似規模的開源模型（如 LLAMA3-3B、Qwen2.5-3B 和 Deepseek-v2-Lite）進行了比較。結果顯示，Moonlight 在效能上優於使用相同數量 tokens 訓練的模型，與更大引數規模模型相比，也較有競爭力。此外，研究團隊還發現，Muon 可以讓模型的權重更新更「多樣化」，尤其在 MoE 模型中表現突出。在微調階段，在預訓練和微調階段都使用 Muon，模型表現會比用 AdamW 的組合更好，但如果微調和預訓練的最佳化器不一致，優勢就不明顯了。https://github.com/MoonshotAI/Moonlight?tab=readme-ov-filehttps://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdfhttps://x.com/Kimi_Moonshot/status/1893379158472044623
