2025年最让我意外的,不是哪个模型又刷了新榜单,而是翻开几份旗舰模型的技术报告,发现它们不约而同地在优化器里动了刀。尤其是GLM-5那份报告里出现的"Muon Split",让我专门停下来,把另外几家同期报告翻出来对比了一遍。结论挺有意思:大家都在同一个地方,把Muon拆开来用。
这篇笔记不打算停留在"它们用了什么"这种层面,重点想聊聊三份报告为什么会在同一时间指向同一个改动,Muon Split到底在拆什么,以及这背后对预训练工程有什么实际影响。如果你在做大模型训练、正在折腾优化器,或者只是好奇这一波"换优化器"的风到底怎么刮起来的,这篇应该能给你一个比较完整的视角。
1. 为什么2025年的旗舰模型都在换优化器
1.1 AdamW的瓶颈:显存与更新的结构性缺陷
先说个背景。过去几年,大模型预训练基本被AdamW统治,它好用、稳、调参经验丰富。但到了千亿甚至万亿参数这个规模,AdamW的问题越来越藏不住了。
最直接的问题是显存。AdamW要维护一阶动量(momentum)和二阶动量(variance)两个状态,每个参数都要存两份额外的浮点数。参数本身如果是BF16,那还好;但优化器状态如果是FP32,光两块动量就把显存吃掉了8字节/参数,比模型权重还大。一个700B模型,光优化器状态就是5.6TB的量级,分布式训练里光搬运这些状态就够折腾的。
但显存只是表面问题。更深的隐患在于,AdamW把每个参数当作独立的标量来处理,二阶动量只对每个坐标单独缩放。当参数组织成矩阵、张量时,这种"逐元素自适应"其实没有利用梯度矩阵内部的结构信息——同一行的参数之间、同一列的参数之间存在相关性,AdamW完全不看这些。结果就是:收敛稳定,但每一步的方向不一定是最优的,在极大规模训练时,需要的有效步数越来越多。
于是大家开始找替代方案。不是所有场景都需要AdamW这种级别的自适应,矩阵结构本身就可以用来做更好的更新方向。
1.2 Muon是什么:把梯度矩阵"拧"成正交方向
Muon这个名词2025年突然在各大技术报告里冒出来,但它并不是什么横空出世的新东西。它的核心思路其实很朴素:把某个参数矩阵的动量矩阵,做一次近似正交化,然后用一个标量学习率去更新。
拆开看是这样:
- 第一步和Adam一样,维护一个动量矩阵,指数移动平均累积梯度。
- 第二步不再是像Adam那样算二阶矩做逐元素缩放,而是对这个动量矩阵做一次"白化"——目的是让更新方向在不同维度上尽量不相关,相当于把矩阵"拧"成一个接近正交的方向。
- 第三步直接用标量学习率去更新参数,不再需要逐参数的自适应缩放。
这个操作的数学直觉可以这么理解:如果你把一个矩阵看作一个线性映射,那么好的更新方向应该是让这个映射变化得均匀、稳定。正交矩阵恰好满足这个性质——它的列向量彼此独立、长度一致。把动量矩阵尽力投影到正交矩阵附近,相当于在做"矩阵层面的归一化"。
实现上不需要做昂贵的SVD分解,用Newton-Schulz迭代就能近似逼近正交化。这个迭代每次只做矩阵乘法,成本可控,迭代三五次就已经能得到一个不错的近似正交方向。
Muon真正吸引人的地方在于显存。它只需要维护一个动量,省掉了二阶动量那一份状态,优化器状态直接砍半。而且它对矩阵结构参数的效果,理论上比AdamW更接近二阶方法的水平。代价是:第一次接触它的人很容易调崩——它对学习率、对参数形状、对哪些参数能用哪些不能用,都有讲究。
1.3 "同一处改动"指的是什么
GLM-5报告里提到的Muon Split,拆开是两个词:Muon优化器,加上Split拆分。三家旗舰在同一年、同一个地方改了同一个东西,这个"地方"指的就是优化器的落地点:它们都没有直接照搬原版Muon,而是针对大规模预训练的实际约束,对Muon的使用方式做了拆分。
为什么必须拆?因为原版Muon有一个比较"娇气"的设定。它最初的理论推导基于Maximal Update Parametrization(μP)框架,对每层的学习率缩放有严格要求。实际模型中,并不是所有参数都是矩阵——偏置是向量、Embedding和LM Head虽然是大矩阵但形状特殊、LayerNorm里的scale和shift也是向量。这些参数如果用Muon的矩阵正交化逻辑,要么形状不对没法算,要么算出来的结果不稳定。
我记得第一次在一份报告里看到"split"这个词时,第一反应是它可能指某种把更新矩阵切成几块的新算法。翻完三家的描述才发现,大家改的其实是同一件事:把参数按矩阵结构分好组,矩阵类参数走Muon路径,非矩阵参数保留AdamW路径,同时对超大矩阵再做块级拆分,控制正交化的计算和显存成本。这个改动看似简单,却是Muon从"论文里的数学技巧"变成"可以上生产线的训练组件"的关键一步。
2. Muon Split到底在拆什么
2.1 按参数维度拆分:向量参数不能硬套Muon
先说最容易理解、也是三份报告里最一致的一层拆分:按参数维度分组。
Muon的正交化操作,针对的是二维及以上矩阵。一个形状为(out_features, in_features)的权重矩阵,可以做矩阵乘法、可以做正交化迭代。但一个形状为(hidden,)的向量参数,本质上只是一串标量,没有任何"列与行"的结构,硬塞进Muon的框架里,连Newton-Schulz迭代都跑不了。
于是最基本的Split逻辑就是:
- 二维及以上的参数矩阵,例如Attention里的QKV投影、MLP里的上下投影、以及绝大部分weight矩阵,走Muon路径。
- 一维参数,例如所有bias、LayerNorm的gamma和beta,继续走AdamW或者简单的SGD路径。
- Embedding矩阵和LM Head这种形状特殊的大矩阵,各家处理方式略有差别,下一节细说。
这一步拆分的直接收益是:优化器整体还是Muon主导,但规避了向量参数带来的数值稳定性问题。你不需要手动给几百个vector参数单独调学习率,AdamW会帮你处理这些"非结构化参数"。
这里有个细节值得注意:某些实现里,一维参数不光保留AdamW路径,还会单独设置weight decay策略。因为Muon路径下的权重衰减是直接做在正交化更新方向上的,而向量参数不能这么做,Split之后两边各管各的,反而清楚。
2.2 按张量块拆分:正交化的性价比权衡
第二层拆分更工程化,也是"Split"这个命名里最关键的意图。
Muon的正交化虽然避免了SVD那种高不可攀的复杂度,但Newton-Schulz迭代还是要做多次矩阵乘法。对于一个小矩阵,比如(4096, 4096),花不了多少算力;但对于超大矩阵,比如专家模块里的(16384, 16384)甚至更大,一次正交化迭代的计算量就上来了,而且迭代过程中产生的中间矩阵显存开销同样不能忽略。
三家不约而同做了同一件事:把大矩阵沿着行或列方向切成若干个块,对每个块分别跑动量累积和正交化。比如一个(8192, 4096)的矩阵,可以按输出维度切成4个(2048, 4096)的块,每个块独立走Muon,更新时再拼接回原形状。
块级拆分的收益非常实际:
- 单块的最大维度变小,Newton-Schulz迭代的矩阵乘法成本按维度三次方下降,省算力。
- 块与块之间可以并行处理,在张量并行或流水线并行的框架里更容易做负载均衡。
- 显存峰值被控制住了,正交化迭代产生的中间张量不会因为矩阵太大而撑爆显存。
但块切分不是切得越小越好。正交化本来是为了捕捉矩阵整体的相关性,如果切成只有(64, 4096)的细条,每个块内部的结构信息太少,正交化效果大打折扣,最终表现可能还不如直接用AdamW。从三份报告体现的经验看,块的大小通常在百到千级别维度比较合适。
2.3 三份报告的微差异
既然说的是"三家旗舰",具体名字我不展开,但可以聊聊它们在细节取舍上的差异。GLM-5报告里,我注意到它们对LM Head和Embedding的处理比较谨慎——这两个层是模型输入和输出的必经之路,对数值尺度敏感度极高,它们在Muon路径之外给这些层保留了独立的缩放机制。另一家偏重推理效率的旗舰,更强调块拆分与张量并行的配合,做法是把拆分边界和GPU卡数对齐,让每卡只处理属于自己的块,省去额外通信。还有一家则在momentum的计算上做了改动,把原始Muon里的固定beta改成了按训练阶段变化的值,前期高beta追求稳定,后期降beta加快响应。
所以"同一个地方"并不是说三家代码完全一致,而是它们都意识到了同一件事:Muon这个新优化器,必须经过拆分和混合,才能真正服务于旗舰级预训练。原版Muon是一个数学上漂亮的概念,工程落地版本则都是"Muon + AdamW + 块切分 + 缩放策略"的缝合体。
3. 站在工程侧落地Muon Split
3.1 最小实现:一个支持分组的优化器
文本层面理解了split逻辑之后,落到代码其实没有那么玄乎。你需要的不是一个全新的优化器,而是一个能按参数组走不同路径的优化器外壳。下面这段伪代码是我在模拟项目里用的结构,整理出来供参考。
import torch import torch.nn as nn def orthogonalize_(x, iters=5): """Newton-Schulz迭代,把矩阵近似投影到正交方向""" for _ in range(iters): xtx = x @ x.transpose(-1, -2) x.data = x.data * (1.5 - 0.5 * xtx) def split_groups(model, use_muon_names=("weight", "w_g")): muon_params, adam_params = [], [] for name, p in model.named_parameters(): if p.ndim >= 2 and any(k in name for k in use_muon_names): muon_params.append(p) else: adam_params.append(p) return muon_params, adam_params这段代码做了一件事:按维度把参数分成两组。偏置、LayerNorm参数、还有那些名字里不带weight的二维参数,统一走AdamW;二维以上的矩阵权重走Muon。
注意orthogonalize_里的迭代次数iters,这是一个非常关键的超参。迭代次数越多,正交化越彻底,但算力开销也越大。我的经验是:3次是底线,5次是常用值,8次以上基本看不出额外收益,纯属浪费。
3.2 块拆分与正交化迭代
当参数矩阵特别大时,直接对全矩阵做正交化会顶到显存和算力上限。这时就需要做块级别的Split。
def muon_update(param, momentum, lr, beta=0.95, chunk_dim=0, n_chunks=4): grad = param.grad # 动量累积 momentum.lerp_(grad, 1 - beta) # 按chunk_dim拆分 chunks = momentum.chunk(n_chunks, dim=chunk_dim) updated = [] for c in chunks: # 对每个块独立做正交化 orthogonalize_(c, iters=5) updated.append(c) # 拼回原形状 m_orth = torch.cat(updated, dim=chunk_dim) # 标量学习率更新 param.data.add_(-lr, m_orth)核心思路是:动量矩阵先按某个维度切成n_chunks块,每块独立做正交化,再拼接回来。切块维度一般选输出维度(对于参数形状为(out, in)的情况,选dim=0)。因为输出维度的相关性通常更依赖列方向上的投影,按行切能保留更完整的列结构信息。
但这里有个隐藏坑:直接对momentum的chunk做orthogonalize_,如果chunk不是平方矩阵,Newton-Schulz迭代本身不会出错,但收敛行为会变。我的建议是:当块形状不是方阵且长宽比超过3:1时,增加迭代次数,或者先做一次行/列归一化预处理,帮助迭代稳定。
3.3 推荐初始配置与缩放技巧
Muon Split落地时,最让人抓狂的不是算法本身,而是学习率怎么设。Muon在μP框架下对学习率缩放有严格要求,具体到实践,记住一个原则:宽矩阵用更小的学习率,窄矩阵可以用更大的学习率。很多实测结果表明,通常缩放系数取1 / sqrt(max(fan_in, fan_out))或类似形式是合理的起点。
我常用的初始配置如下:
| 参数 | 初始值 | 说明 |
|---|---|---|
| muon_beta | 0.95 | 动量系数,比Adam的0.9更平滑 |
| ortho_iters | 5 | 正交化迭代次数 |
| chunk_size_max | 8192 | 超过这个维度的矩阵才做块拆分 |
| adam_lr_ratio | 0.1 | AdamW组的lr = muon_lr * 0.1 |
| weight_decay | 0.01 | 只加在Muon组的更新方向上 |
AdamW组的学习率给到Muon组的十分之一,是因为向量参数和特殊层一般不需要太激进的更新。这个比例在很多规模实验里都成立,但不绝对,如果你的模型对embedding层特别敏感,可以把比例压到0.03。
另外一个技巧:Muon路径下不要单独给momentum设weight decay。动量累积的是梯度,正则化应该作用在参数更新上,而不是污染动量。这也是Split的价值——两组参数各走各的路径,才不会出现"每个参数都被同一个优化器逻辑硬套"的情况。
3.4 显存、吞吐与稳定性的实测经验
我自己在模拟项目X上跑过Muon Split对比,规模谈不上旗舰,但够说明趋势。直接说观察到的数据点:
- 显存:相比AdamW,优化器状态显著压缩。700B这个量级下,AdamW光二阶动量省掉的部分就是大头。实测显存峰值大概能低15%到20%,具体取决于模型里二维矩阵参数的占比。
- 吞吐:正交化迭代带来的额外计算量不可忽视。如果每个step都做5次迭代,小batch下训练吞吐会掉3%到5%。块拆分之后,这个开销可以被并行度摊薄,实际损失更小。
- 稳定性:前几千step的loss曲线明显比AdamW平滑,尤其是梯度噪声大的早期阶段,Muon的方向约束让更新不容易走偏。
需要提醒的是:不要一上来就把所有层都切进Muon。我建议先用一个小规模脱模实验验证分组逻辑,跑通后再逐步放大。Muon Split本身是个很有弹性的框架,改动的空间比换一个完全不同的优化器大得多。
4. 训练过程中最常踩的坑
4.1 影子状态与NaN
最经典的问题是训练到几千步后突然loss变NaN,而且往往不是梯度爆炸引起的,而是优化器状态本身出问题了。我排查过几次这类问题,最后定位到的原因基本集中在两个地方:一个是正交化迭代前没有对动量矩阵做任何数值保护,一旦某个块出现极端大值,Newton-Schulz迭代会让异常越放越大;另一个是momentum里泄露了权重衰减的贡献,导致累积量失真。
解决办法也直接:在正交化之前对动量矩阵做一次整体缩放,让它的谱范数落在安全区间。实操中可以加一个简单的scale:
scale = m.norm() / (m.numel() ** 0.5) m.data.mul_(1.0 / (scale + 1e-8))这不算理论最优,但能显著降低NaN概率。另外多提一句:使用混合精度训练时,动量建议保留在更高精度里,尤其是BF16环境下,直接用BF16动量矩阵做正交化,迭代误差会被快速放大。
4.2 收敛变慢与lr缩放
另一个常见问题是:训练稳定是稳定了,但曲线掉得比AdamW还慢,loss下不去。这时候八成不是算法问题,而是学习率没跟模型宽度联动。
Muon和AdamW的lr敏感度完全不同。很多人从AdamW切到Muon时,保留了一样的基础学习率,结果就是模型学得太慢。按μP框架的思路,Muon组的lr应该跟具体参数矩阵的fan_in、fan_out挂钩,不同宽度的层要用不同的lr。如果整个模型共用一个lr,那窄层基本学不动,宽层又可能震荡。
我见过一份报告里的做法是给每个参数矩阵单独存一份"形状缩放因子",在step里乘到lr上。这个方案不复杂,但效果立竿见影,比全局搜lr高效很多。
4.3 显存、精度与并行
还有一个容易被忽略的点:Muon Split不等于完全没有自适应优化器开销。块拆分之后,虽然每个块的中间矩阵变小了,但块数量变多,如果显存碎片化严重,实际峰值可能比你预想的高。建议在显存估算时按"最大块"而不是"总模型"来做预算。
张量并行场景下,split边界怎么切也很有讲究。我的经验是:尽量让每个块落在单卡内部,避免一个块的orthogonalize迭代需要跨卡通信才能完成。否则你会在通信日志里看到一个诡异的现象:优化器计算本身没占多少时间,通信却成了瓶颈。
4.4 性能验证方法
最后聊一个方法论层面的问题:怎么确定你的Split逻辑真的生效了?
不要只看下游任务指标,那是滞后的信号。我会在训练早期就盯两个指标:动量矩阵的正交性残差,以及不同参数块之间的更新方向余弦相似度。正交性残差用||X^T X - I||来算,数值越小说明正交化越彻底。更新方向余弦相似度则用来判断块拆分是否破坏了整体方向的一致性——如果块之间方向互相打架,说明拆分维度选错了。
这两个指标在loss曲线上不一定能立刻看出来,但它们能帮你提前发现优化器层面的问题,而不是等到训练了几天之后才追悔莫及。
5. 读完三份报告后的一些想法
5.1 为什么是"拆"而不是"换"
说实话,看完三份报告的对比之后,我最大的感受不是"Muon好强",而是"工程化落地的关键往往不是引入新东西,而是知道怎么把新东西塞进旧体系里"。
Muon的原版设计很漂亮,但直接端到端替换AdamW,在真实预训练里会撞上无数细节:向量参数怎么办、超大矩阵怎么控制算力、不同层该不该用同一套超参、embedding区域太敏感怎么办。三家旗舰做的是同一个动作——拆。把参数按形状拆开,把更新逻辑按块拆开,把学习率按宽度拆开。所有拆法的目标不是让算法更复杂,而是让新优化器能适应旧训练管线里已经成熟的分布式、混合精度、并行策略。
这其实是预训练工程里反复出现的一个模式:论文给的是数学上美的东西,生产线上要的是能跟已有约束共存的东西。Muon Split是"从论文到生产"的一个缩影。
5.2 下一步想验证的扩展
我自己下一步想试的方向,是把Muon Split和层间优化器调度结合。具体来说:训练前期用AdamW做稳定探索,中期切换到Muon加速收敛,后期再用AdamW收尾。这个想法不成熟,但从三份报告的实验曲线看,Muon的优势集中在中期收益,后期边际收益不大。如果调度层能把"前期稳、中期快、后期细"变成现实,也许能在同等算力下再挤出一截性能。
不过这些都是后话。当前这个阶段,如果你正在做预训练,我建议你先不要急着全量上线Muon Split。先在几个block上做demo验证,用前面提的正交性残差和方向一致性指标评估效果,再逐步扩大范围。优化器这个东西,用户自适应很重要,别人的最佳实践只能当参考,不能当标准答案。
最后分享一个小技巧,也是我在实际使用中发现的:无论Split怎么设置,都建议在checkpoint里额外存一份优化器配置记录。因为你每次改lr缩放比例或块大小,模型行为都会明显变化,如果没有存档,等训练到中后期发现效果不理想,很难判断是哪一次改动造成的。优化器的调试也是一种实验管理,做好版本记录,比多跑几轮实验更值钱。