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

作者：量子位 2025-02-24 09:15:00

人工智能開(kāi)源

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

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

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

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

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

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

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

消息一出，當(dāng)時(shí)Muon的作者也都很激動(dòng)，主要作者Keller Jordan表示這是Muon規(guī)模化的首個(gè)成功報(bào)告，為團(tuán)隊(duì)送上了祝賀。

另一名貢獻(xiàn)者，當(dāng)時(shí)負(fù)責(zé)Muon規(guī)模化實(shí)驗(yàn)的Hyperbolic Labs聯(lián)創(chuàng)兼CTO Yuchen Jin也表示，月暗團(tuán)隊(duì)的這項(xiàng)成果，是Muon的一次勝利。

將AdamW特點(diǎn)引入Muon

在介紹月暗團(tuán)隊(duì)的工作之前，先來(lái)了解一下Muon是個(gè)什么樣的技術(shù)。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

其中0.2是通過(guò)實(shí)驗(yàn)確定的一個(gè)常數(shù)，用于將Muon的更新尺度與AdamW對(duì)齊。

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

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

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

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

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

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

證明Muon擴(kuò)展可行性

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

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

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

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

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

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

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

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

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