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

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

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

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

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

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

同時團隊還基于DeepSeek架構(gòu)訓(xùn)練出了一個16B的MoE模型，與改進后的優(yōu)化算法一同開源。

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

• Muon能否用于更大規(guī)模的訓(xùn)練？
• Muon能否在更大規(guī)模的GPU集群上使用？
• Muon是否同樣適用于微調(diào)和強化學(xué)習(xí)？
  現(xiàn)在月暗團隊用實驗給出了回答——全部都是Yes。

消息一出，當(dāng)時Muon的作者也都很激動，主要作者Keller Jordan表示這是Muon規(guī)?；氖讉€成功報告，為團隊送上了祝賀。

另一名貢獻者，當(dāng)時負責(zé)Muon規(guī)?；瘜嶒灥腍yperbolic Labs聯(lián)創(chuàng)兼CTO?Yuchen Jin也表示，月暗團隊的這項成果，是Muon的一次勝利。

將AdamW特點引入Muon

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

這是一種神經(jīng)網(wǎng)絡(luò)隱藏層的2D參數(shù)優(yōu)化器，主要作者是OpenAI深度學(xué)習(xí)團隊的Keller Jordan。

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

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

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

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

現(xiàn)在經(jīng)過月暗團隊的改進之后，Muon被證明對于更大的模型和數(shù)據(jù)集同樣適用。

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

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

這樣做的原因是作者發(fā)現(xiàn)直接將Muon應(yīng)用到大規(guī)模訓(xùn)練時，模型權(quán)重和層輸出的幅度會持續(xù)增長，最終超出bf16的高精度表示范圍，損害模型性能。

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

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

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

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

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

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

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

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

除了對Muon本身的改進，要想將Muon用于更大規(guī)模的訓(xùn)練，還需要將其擴展到分布式訓(xùn)練環(huán)境中。

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

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

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

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

這種實現(xiàn)方式在最小化內(nèi)存占用和通信開銷的同時，最大限度地保留了原始Muon算法的數(shù)學(xué)性質(zhì)。

證明Muon擴展可行性

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

結(jié)果表明,在計算預(yù)算最優(yōu)的情況下，Muon的樣本效率是AdamW的1.92倍，即訓(xùn)練FLOPS只需AdamW的52%，就能達到相當(dāng)?shù)男阅堋?/p>

這一發(fā)現(xiàn)證實了Muon在大規(guī)模訓(xùn)練中的效率優(yōu)勢。

在此基礎(chǔ)之上，作者以DeepSeek-V3-Small架構(gòu)作為基礎(chǔ)，用改進的Muon訓(xùn)練了Moonlight模型。

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

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

即使與使用更大數(shù)據(jù)集訓(xùn)練的稠密模型相比，Moonlight也展現(xiàn)了極強的競爭力。

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

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

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

結(jié)果發(fā)現(xiàn)，Muon優(yōu)化的矩陣在各層各類參數(shù)上，總是比AdamW有更高的奇異值熵。這從經(jīng)驗上驗證了Muon通過正交化來學(xué)習(xí)更多樣化表示的直覺。

最后，在Moonlight模型的基礎(chǔ)上，作者還探索了Muon在指導(dǎo)微調(diào)階段的效果，結(jié)果表明，在預(yù)訓(xùn)練和微調(diào)階段均使用Muon的效果是最佳的。

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