月之暗面開源改進版Muon優化器，算力需求比AdamW銳減48%，DeepSeek也適用

克雷西 2025-02-23 17:31:33 來源：量子位

克雷西 發自 凹非寺
量子位 | 公眾號 QbitAI

算力需求比AdamW直降48%，OpenAI技術人員提出的訓練優化算法Muon，被月之暗面團隊又推進了一步！

團隊發現了Muon方法的Scaling Law，做出改進并證明了Muon對更大的模型同樣適用。

在參數量最高1.5B的不同Llama架構模型上，改進后的Muon算力需求僅為AdamW的52%。

同時團隊還基于DeepSeek架構訓練出了一個16B的MoE模型，與改進后的優化算法一同開源。

Muon技術博客發布當時主要適用于較小的模型和數據集，作者留下了三個懸而未決的問題：

• Muon能否用于更大規模的訓練？
• Muon能否在更大規模的GPU集群上使用？
• Muon是否同樣適用于微調和強化學習？
  現在月暗團隊用實驗給出了回答——全部都是Yes。

消息一出，當時Muon的作者也都很激動，主要作者Keller Jordan表示這是Muon規模化的首個成功報告，為團隊送上了祝賀。

另一名貢獻者，當時負責Muon規模化實驗的Hyperbolic Labs聯創兼CTO?Yuchen Jin也表示，月暗團隊的這項成果，是Muon的一次勝利。

將AdamW特點引入Muon

在介紹月暗團隊的工作之前，先來了解一下Muon是個什么樣的技術。

這是一種神經網絡隱藏層的2D參數優化器，主要作者是OpenAI深度學習團隊的Keller Jordan。

這項成果發表于去年的12月8日，而Keller也是去年12月加入的OpenAI。

Muon的核心思想是通過正交化梯度更新矩陣，避免參數更新陷入局部極小，使模型能夠學習到更加多樣化的特征表示。

在94%的精度下，Muon把CIFAR-10在A100上的訓練時間從3.3秒縮短至2.6秒。

不過當時Muon團隊只證明了其在小型模型和數據集上的可行性，對于較大的模型能否適用則是個未知數。

現在經過月暗團隊的改進之后，Muon被證明對于更大的模型和數據集同樣適用。

針對模型本身，團隊吸收了AdamW中的一些特點，移植到了Muon當中，具體包括兩個方面。

一是引入了權重衰減機制，在權重更新公式中添加了一個帶有衰減系數的項。

這樣做的原因是作者發現直接將Muon應用到大規模訓練時，模型權重和層輸出的幅度會持續增長，最終超出bf16的高精度表示范圍，損害模型性能。

在訓練一個8億參數模型至100B tokens（約5倍計算預算最優）的過程中，團隊對比了AdamW、無權重衰減的Muon和帶權重衰減的Muon。

結果顯示，帶權重衰減的Muon在過擬合階段取得了最佳效果，驗證了權重衰減的必要性。

第二項改進，是調整了Muon的參數更新尺度，使不同形狀矩陣參數的更新幅度保持一致，并與AdamW的更新幅度匹配。

Muon的一個特點是，對于形狀為[A,B]的矩陣參數，其理論更新幅度為sqrt(1/max(A,B))。

這導致不同形狀矩陣參數的更新幅度差異很大，比如對于MLP這種寬矩陣，更新會過小，而將每個head看作獨立矩陣時，更新又會過大。

此外，這個幅度也與AdamW不一致，給超參數的設置帶來困難。

為了讓不同矩陣參數的更新幅度匹配，并與AdamW保持一致，作者嘗試了幾種改進方案，最終選擇直接基于形狀調整每個參數的學習率。

其中0.2是通過實驗確定的一個常數，用于將Muon的更新尺度與AdamW對齊。

除了對Muon本身的改進，要想將Muon用于更大規模的訓練，還需要將其擴展到分布式訓練環境中。

由于Muon需要完整的梯度矩陣來計算正交化的更新量，而現有的分布式訓練框架（如ZeRO-1、Megatron-LM等）都假設優化器狀態可以獨立地按元素切分到不同設備上，所以它們無法直接支持Muon。

為了解決這個問題，論文作者提出了分布式Muon的并行化策略。

它在ZeRO-1的基礎上引入了兩個額外的操作:

• 一是在每個數據并行組內做梯度聚合通信，將分散的梯度切片合并成完整的矩陣；
• 二是基于聚合后的梯度矩陣并行計算正交化的更新量，然后只保留與本地參數對應的那一部分。

這種實現方式在最小化內存占用和通信開銷的同時，最大限度地保留了原始Muon算法的數學性質。

證明Muon擴展可行性

基于上述Muon改進,作者取得了以下成果，作者在Llama架構的一系列稠密模型上，進行了Muon和AdamW的模型縮放對比實驗。

結果表明,在計算預算最優的情況下，Muon的樣本效率是AdamW的1.92倍，即訓練FLOPS只需AdamW的52%，就能達到相當的性能。

這一發現證實了Muon在大規模訓練中的效率優勢。

在此基礎之上，作者以DeepSeek-V3-Small架構作為基礎，用改進的Muon訓練了Moonlight模型。

Moonlight是一個MoE模型，具有15.29B的總參數和2.24B激活參數，訓練token量為5.7T。

與相同規模和數據量的模型相比，Moonlight在英語理解與推理（MMLU、TriviaQA、BBH）、代碼生成（HumanEval、MBPP）、數學推理（GSM8K、MATH、CMATH）、中文理解（C-Eval、CMMLU）等各類任務上都取得了明顯更好的性能。

即使與使用更大數據集訓練的稠密模型相比，Moonlight也展現了極強的競爭力。

與多個知名語言模型的對比表明，Moonlight在性能-訓練預算平面上推進了帕累托前沿（Pareto Frontier）。

（注：帕累托前沿是一個經濟學和管理學中的概念，描述的是在多目標決策問題中所有可能的最優解的集合，這些解在多個目標之間取得了最佳平衡。在帕累托前沿上的每一個點，都意味著一個目標的改善必然以犧牲另一個目標為代價，因此它代表了在多個目標之間實現的最佳權衡。）

為了進一步分析Muon更新矩陣參數的內在機制，作者對比了Muon和AdamW訓練得到的模型在不同訓練階段的參數矩陣奇異值譜。

結果發現，Muon優化的矩陣在各層各類參數上，總是比AdamW有更高的奇異值熵。這從經驗上驗證了Muon通過正交化來學習更多樣化表示的直覺。

最后，在Moonlight模型的基礎上，作者還探索了Muon在指導微調階段的效果，結果表明，在預訓練和微調階段均使用Muon的效果是最佳的。

技術報告：
https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
Code：
https://github.com/MoonshotAI/Moonlight
Moonlight模型：
https://huggingface.co/moonshotai/Moonlight-16B-A3B