算力需求比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)化算法一同開源。 Muon技術(shù)博客發(fā)布當(dāng)時(shí)主要適用于較小的模型和數(shù)據(jù)集,作者留下了三個(gè)懸而未決的問題: Muon能否用于更大規(guī)模的訓(xùn)練? Muon能否在更大規(guī)模的GPU集群上使用? Muon是否同樣適用于微調(diào)和強(qiáng)化學(xué)習(xí)? 消息一出,當(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ì)的工作之前,先來了解一下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的核心思想是通過正交化梯度更新矩陣,避免參數(shù)更新陷入局部極小,使模型能夠?qū)W習(xí)到更加多樣化的特征表示。 在94%的精度下,Muon把CIFAR-10在A100上的訓(xùn)練時(shí)間從3.3秒縮短至2.6秒。 不過當(dāng)時(shí)Muon團(tuán)隊(duì)只證明了其在小型模型和數(shù)據(jù)集上的可行性,對(duì)于較大的模型能否適用則是個(gè)未知數(shù)。 現(xiàn)在經(jīng)過月暗團(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))的過程中,團(tuán)隊(duì)對(duì)比了AdamW、無權(quán)重衰減的Muon和帶權(quán)重衰減的Muon。 結(jié)果顯示,帶權(quán)重衰減的Muon在過擬合階段取得了最佳效果,驗(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ì)過小,而將每個(gè)head看作獨(dú)立矩陣時(shí),更新又會(huì)過大。 此外,這個(gè)幅度也與AdamW不一致,給超參數(shù)的設(shè)置帶來困難。 為了讓不同矩陣參數(shù)的更新幅度匹配,并與AdamW保持一致,作者嘗試了幾種改進(jìn)方案,最終選擇直接基于形狀調(diào)整每個(gè)參數(shù)的學(xué)習(xí)率。 其中0.2是通過實(shí)驗(yàn)確定的一個(gè)常數(shù),用于將Muon的更新尺度與AdamW對(duì)齊。 除了對(duì)Muon本身的改進(jìn),要想將Muon用于更大規(guī)模的訓(xùn)練,還需要將其擴(kuò)展到分布式訓(xùn)練環(huán)境中。 由于Muon需要完整的梯度矩陣來計(jì)算正交化的更新量,而現(xiàn)有的分布式訓(xùn)練框架(如ZeRO-1、Megatron-LM等)都假設(shè)優(yōu)化器狀態(tài)可以獨(dú)立地按元素切分到不同設(shè)備上,所以它們無法直接支持Muon。 為了解決這個(gè)問題,論文作者提出了分布式Muon的并行化策略。 它在ZeRO-1的基礎(chǔ)上引入了兩個(gè)額外的操作: 一是在每個(gè)數(shù)據(jù)并行組內(nèi)做梯度聚合通信,將分散的梯度切片合并成完整的矩陣; 二是基于聚合后的梯度矩陣并行計(jì)算正交化的更新量,然后只保留與本地參數(shù)對(duì)應(yīng)的那一部分。 這種實(shí)現(xiàn)方式在最小化內(nè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在英語理解與推理(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è)知名語言模型的對(duì)比表明,Moonlight在性能-訓(xùn)練預(yù)算平面上推進(jìn)了帕累托前沿(Pareto Frontier)。 (注:帕累托前沿是一個(gè)經(jīng)濟(jì)學(xué)和管理學(xué)中的概念,描述的是在多目標(biāo)決策問題中所有可能的最優(yōu)解的集合,這些解在多個(gè)目標(biāo)之間取得了最佳平衡。在帕累托前沿上的每一個(gè)點(diǎn),都意味著一個(gè)目標(biāo)的改善必然以犧牲另一個(gè)目標(biāo)為代價(jià),因此它代表了在多個(gè)目標(biāo)之間實(shí)現(xiàn)的最佳權(quán)衡。) 為了進(jìn)一步分析Muon更新矩陣參數(shù)的內(nèi)在機(jī)制,作者對(duì)比了Muon和AdamW訓(xùn)練得到的模型在不同訓(xùn)練階段的參數(shù)矩陣奇異值譜。 結(jié)果發(fā)現(xiàn),Muon優(yōu)化的矩陣在各層各類參數(shù)上,總是比AdamW有更高的奇異值熵。這從經(jīng)驗(yàn)上驗(yàn)證了Muon通過正交化來學(xué)習(xí)更多樣化表示的直覺。 最后,在Moonlight模型的基礎(chǔ)上,作者還探索了Muon在指導(dǎo)微調(diào)階段的效果,結(jié)果表明,在預(yù)訓(xùn)練和微調(diào)階段均使用Muon的效果是最佳的。 技術(shù)報(bào)告: 本文來源:量子位 |
原創(chuàng)欄目
IT百科
網(wǎng)友評(píng)論
聚超值•精選