ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

大模型训练优化器避坑指南:从SGD、AdamW到Muon的演进与实操

大模型训练优化器避坑指南:从SGD、AdamW到Muon的演进与实操 训练大模型和训练小模型一个最容易被低估的差距在哪里小模型的loss哪怕飘一会儿重跑也就几个小时大模型一旦在优化器层面出问题损失的是整个集群几天的算力。很多人遇到loss不降、崩塌或者震荡第一反应是“学习率调大点/调小点”但真正决定训练能否稳定走完的往往是优化器optimizer本身——它怎么更新参数、状态怎么存、跨卡怎么同步、续训怎么恢复。这一篇是大模型训练系列的第二节我专门把优化器拆开讲从最朴素的SGD到Adam、AdamW再到最近社区里讨论度很高的Muon每个名字背后到底改了什么以及我在实操中因为优化器吃过的亏。不管你是刚准备拉起第一个大模型训练还是已经在大规模集群上跑了很久这篇文章都值得你对照自己的配置过一遍。1. 优化器在训练中的职责梯度之后的那一步决定一切1.1 反向传播结束参数并不会自动更新很多人看完深度学习入门教程会以为“反向传播”和“梯度下降”是一回事。其实不是。反向传播只负责计算梯度——也就是每个参数相对loss的导数真正拿着梯度去改动参数的值是优化器干的事。这条界限在写简单模型时很容易被忽略因为PyTorch里一行optimizer.zero_grad(); optimizer.step()就把更新动作隐藏了。但到了大模型训练两者必须严格分清梯度只是一份“建议”优化器才是那个决定“走多快、往哪走、要不要刹车”的机构。最朴素的更新公式是SGDθ θ - lr * g。这里的lr学习率是全局步长g是本轮梯度。“全局”两个字是关键——不管哪个参数只要梯度大就走得多。听起来很合理但真实训练里不同参数的梯度尺度差异是数量级的word embedding层、attention输出层、FFN中间层的梯度分布完全不同。这就像让同一套规则同时管市区小电驴和高速公路货车结果必然是有的地方堵死、有的地方漂移。1.2 大模型场景下优化器要解决三个核心问题第一个是方向问题。mini-batch的梯度噪声在大模型里一点不比小模型小随机梯度天然抖动剧烈没有平滑机制的训练曲线会非常难看甚至根本不收敛。第二个是尺度问题。不同参数、不同层的学习节奏应该不一样自适应类优化器就是为了回答“每个参数该走多大步”而生的。第三个是鲁棒性问题。训练算力这么贵优化器不能三天两头因为数值问题炸掉——NaN、loss spike、梯度爆炸这些在LLM训练里不是偶发事件而是常态。所以我一直把优化器比作训练里的方向盘和减震器SGD是没减震的纯手动档Adam是自适应巡航Muon这一代新方法则试图在“方向”层面做文章。方向盘没配好引擎再好也开不成直线。1.3 优化器不是独立组件它和学习率、batch策略深度耦合大模型训练圈常说“优化器 lr schedule batch size 梯度裁剪”是一套组合拳。同一个Adam配置warmup步数不同、梯度累积步数不同表现可能天差地别。我见过太多案例有人把别人跑通的超参数原样搬过来只是把batch size翻倍却没同步调整学习率结果两三千步后loss直接起飞。优化器本身没问题是它接收到的“输入环境”变了。所以这篇文章讲优化器我不会只讲公式而是把优化器放到它实际运行的整套系统里看它如何被学习率调度影响如何在分布式状态下保持一致如何决定显存开销以及续训时怎么恢复。这些才是大模型训练时真正决定成败的细节。2. 从SGD到Adam到AdamW主流优化器的演进逻辑2.1 Momentum给梯度加“惯性”SGD最让人头疼的是震荡。山谷形状的损失面里梯度在陡峭方向来回摆在平缓方向又推不动。这时候给更新加一项“历史动量”就像给小球一个惯性当前方向的梯度与上一个时刻的移动方向做加权平均。公式很简单m_t β1 * m_{t-1} (1 - β1) * g_tθ_{t1} θ_t - lr * m_t注意这里的m是梯度的指数滑动平均β1一般取0.9左右。作用是在震荡方向相互抵消、在一致方向持续叠加。SGDMomentum在大模型里不流行但它的物理直觉是后面所有优化器的地基优化器本质上是在对梯度做滤波和塑形。2.2 RMSProp给每个参数单独配一个“步长尺”SGDMomentum仍然只有一个全局学习率但现实是稀疏参数比如embedding里出现次数很低的token对应向量梯度长期很小高频参数梯度又可能爆大。于是有了RMSProp这类方法对每个参数的梯度平方做指数滑动平均再用lr / sqrt(v eps)作为该参数的学习率。梯度平方大的参数步长自动变小梯度一直很小的参数步长自动变大。这个思想直接沿用到了Adam里。2.3 Adam一阶动量与二阶动量的结合Adam把Momentum和RMSProp组合起来m_t β1 * m_{t-1} (1 - β1) * g_tv_t β2 * v_{t-1} (1 - β2) * g_t²m_hat m_t / (1 - β1^t)v_hat v_t / (1 - β2^t)θ_{t1} θ_t - lr * m_hat / (sqrt(v_hat) eps)前两项分别是梯度的一阶矩和二阶矩估计。后半部分 m_hat 和 v_hat 是偏差校正训练最初几步滑动平均还没有充分观测数据不校正的话前几步更新会偏小或偏大。有人说Adam是“默认先猜一个方向再用统计量修正”其实更准确说它是用二阶矩动态调整每个参数的步长同时用一阶矩平滑更新方向。大模型训练普遍选择Adam核心原因有两个第一它对梯度尺度不敏感embedding层、attention层、FFN层的更新幅度都能维持在合理范围第二它对稀疏梯度友好大规模词表下大部分位置梯度为0Adam会在这些维度上“见机行事”而SGD遇到稀疏梯度时训练极其痛苦。2.4 Adam的显存代价这是大模型最硬的约束Adam不是免费的。标准实现里每个参数除了自身数值还要额外保存两段状态一阶动量m和二阶动量v而且大类上都是FP32。算一下7B参数模型模型权重FP16约14GBAdam状态FP32的m FP32的v2 * 4 * 7 56GB混合精度训练还需要一份FP32主权重副本28GB也就是说7B模型用标准AdamW训练光优化器相关显存就占84GB左右。很多同学做小模型从来没想过这回事但到了70B级别优化器状态直接决定单卡能不能塞下、是否能上ZeRO也决定了显存换通信的权衡。这也是为什么业界主流不是换掉Adam而是“压缩”Adam状态8bit优化器、ZeRO分片、混合精度状态管理都是为了在这个角落抠显存。优化器每参数额外状态7B模型下优化器显存FP16主权重不计典型行为SGD无0需要精细调lr对大模型基本失控SGDMomentum1个FP32动量28GB缓解震荡仍无自适应RMSProp1个FP32平方梯度28GB自适应步长但不稳定Adam/AdamW2个FP32动量56GB稳定、适配度好显存代价大AdamW(8bit)2个8bit动量14GB显存大量节省精度损失有限3. AdamW与超参数解剖大模型训练真正要理解的东西3.1 Adam vs AdamWweight decay的位置不是小事很多人只知道“大模型要用AdamW不要用Adam”却说不清为什么。关键在weight decay的写法。经典Adam里正则项是加在loss里的L2正则梯度更新时等价于在原始梯度上额外加一项wd * θ。但Adam会对每个参数的梯度做二阶矩归一化权重衰减项也会被“自适应地缩放”——这等于不同参数施加了不同强度的正则结果完全扭曲了L2正则的初衷。AdamWDecoupled Weight Decay把权重衰减从loss里拆出来直接在参数更新后独立执行θ_{t1} θ_t - lr * m_hat / (sqrt(v_hat) eps) - lr * wd * θ_t这时正则强度与Adam的自适应缩放无关每个参数都以相同比率衰减。实际表现上AdamW在大部分LLM任务里比Adam稳定泛化也更好。所以我现在看到还在用传统Adam训大模型的项目第一反应都是先让它们改成AdamW。3.2 β1、β2、epsilon分别管什么很多人的调参经验是“Adam不用调”这话在small model时代勉强成立在大模型训练里很快会被打脸。β1是一阶动量系数管梯度方向的平滑程度。默认0.9表示当前时刻只保留10%的新鲜信息历史占90%。练大模型时我见过有人把β1调到0.95效果是loss曲线更平滑但代价是更新反应变慢一旦出现异常梯度刹车也会变慢。β2是二阶动量系数管步长估计的历史窗口。默认0.999意味着v统计的是过去大约1000步的梯度平方平均。大模型训练里一个常见的做法是把β2调小到0.95甚至0.9因为训练步数很长时早期梯度平方的统计会对当前产生滞后导致步长估计失真。β2越小对最近梯度变化越敏感越能快速响应loss spike但也会让自适应步长更“神经质”。epsilon就更微妙了。很多人只记得它是防止除零的但在FP16/FP32混合精度下epsilon还充当了“数值地板”。sqrt(v) eps里如果eps太小比如1e-8在极端情况下二阶矩v计算到极小值时依然可能产生异常大的更新如果把eps调大到1e-6或1e-4更新步长会整体变保守有时反而能抑制loss spike。大模型混合精度训练里我自己会先用1e-8跑短实验如果数值抖动明显会试1e-6。3.3 混合精度下的优化器状态FP32 master weight是必须的吗标准做法是模型参数以FP16存到显存里做前向反向但另存一份FP32主权重优化器更新在FP32主权重上完成后再cast回FP16。原因是FP16只有大约3位十进制的数值精度梯度累积几百步后误差会像滚雪球一样不可接受。这带来的显存负担就是前面那张表里多出的主权重副本。所以很多训练框架会在优化器层面做“状态压缩”Adam的二阶动量v对精度不太敏感可以安全降成FP16或8bit一阶动量m则更“娇贵”一些降精度后某些任务里会掉点需要实测。比较稳的路线是先跑标准AdamW做基线确认收敛趋势后再把v换8bit最后再考虑m降精度。4. Muon另一种思路的隐藏层优化器4.1 Muon的出发点隐藏层可能不需要“逐参数缩放”2024年末以来“Muon优化器”在大模型训练社区里讨论度一直很高。它的出发点很反直觉Adam的逐参数缩放真的很厉害但我们真的处处需要它吗在μPMaximal Update Parameterization最大化更新参数化视角下隐藏层的核心问题是更新的方向而不是每个参数被放大的幅度。隐藏层的权重矩阵其奇异值分布的规范程度、谱范数大小才真正决定深层网络能否稳定训练。而Adam对每个元素独立缩放本质上是把梯度矩阵的每个元素当作互不相关的标量处理这种“精细”有时候反而打乱了矩阵层面的结构化信息。所以Muon选择了另一个操作对梯度矩阵做近似正交化处理用一个固定的、可预测的尺度去更新隐藏层而不是像Adam那样逐元素缩放。这个思路在多层Transformer的线性投影层上很有效。4.2 核心操作Newton-Schulz迭代Muon里的关键机制是Newton-Schulz迭代。它的目标是把梯度矩阵往正交矩阵方向“掰”但不去做昂贵的SVD。对矩阵X迭代下面这个式子数次典型取5次X ← (3X - X Xᵀ X) / 2这个操作会让X的奇异值分布向1收敛太小的奇异值被放大太大的奇异值被压缩。迭代之后再用类似SGD带Momentum的方式更新参数。这样每个层的更新幅度天然有界不需要像Adam那样配一整套二阶矩状态。注意Newton-Schulz迭代是针对二维矩阵的所以Muon天然适合把层的权重整理成矩阵来用Transformer里绝大多数线性层都满足这个条件。要注意的是迭代次数不是越多越好迭代次数过高矩阵会被“过度正交化”反而丢掉梯度里的有效信息太低又起不到校准作用。论文和社区实现里5次是常见选择我实测时也先用5次再根据训练稳定性微调。4.3 实操中的MuonAdam组合目前Muon在大模型里比较靠谱的用法是“分工制”隐藏层比如FFN的线性层、attention里的Q/K/V/O投影用Muon而embedding层、输出head层、残差分支这类需要逐元素精细调节的地方仍然用Adam。道理也很直白残差连接和输出层的更新不需要太“结构感”反而希望Adam的自适应能力去处理梯度稀疏和尺度不一致。我跑过一次实验对比相同的数据和模型规模下纯AdamW的超参数需要做不少warmup调试而MuonAdam的组合在较长序列、较深模型上loss曲线更稳对学习率初始值的敏感度也小一些。起步学习率上Muon部分用比Adam大得多的lr一个数量级左右是常有的事因为正交化后更新幅度天然被限制住了。这里我不建议直接照抄任何固定数值你的模型宽度、深度、数据批次都会有影响。Muon也不是没有麻烦。它的二阶“状态”少了一半显存和通信开销显著降低但当梯度矩阵本身很稀疏时正交化操作反而会放大稀疏结构里的噪声。另外Muon对库版本比较敏感——用的是PyTorch的SDDMM、矩阵乘法路径误差累积和GPU上的kernel选择都有影响。想从AdamW跳到Muon别只看“新优化器降显存”就盲目切先用小规模实验对比一下数值曲线。5. 实操中的优化器状态、显存与分布式踩坑记5.1 warmup、学习率曲线和“塑形”聊优化器绕不开学习率调度。我听到一个有意思的类比游戏本工具箱里有个“曲线塑型优化器”本质上是通过改功耗、风扇曲线去影响性能曲线的形状。我们做训练其实也在做“曲线塑型”——不过对象是loss下降曲线工具是warmup和lr schedule。大模型训练中warmup几乎是必修课。原因不完全是“学习率从小到大更稳”更重要的是Adam这类优化器的二阶动量统计在最初几步是失真的。如果一开始就给大学习率v统计还没起来更新步长会被严重高估直接冲进损失面的陡峭区域。常见做法是前面2000-4000步线性warmup之后进入稳定期。到了训练后期用余弦退火或固定lr后接一小段cooldown都比一直保持恒定lr更有利于收敛到平坦区域。我见过不少同学把warmup调到几百步就以为够用了。对小模型也许没错但大模型动辄几万步训练前期几百步的数据不足让Adam的二阶动量形成稳定估计所以我把warmup步数通常设成总步数的1%到2%太短会明显感受到前期loss下降不光滑太长又浪费算力。5.2 梯度累积与分布式更新的一致性大模型受限于显存经常要用gradient accumulation多batch的梯度先在本地累加再统一更新一次参数。这里最常见的坑是梯度累积后忘了给loss做归一化。比如你想等效batch size 32而每卡只能塞batch size 4、累积8步那么每步的loss应该统一除以8或等效的累积次数否则梯度会被放大8倍前面的lr schedule全部失效。分布式训练里优化器状态的同步性也很容易出问题。DDP会把梯度all-reduce之后再step所有rank上的优化器状态理论上是一致的只要初始化方式相同。但如果你在代码里对优化器做了非确定性的操作比如手改m/v、条件分支里只让部分rank更新马上就会看到各卡状态漂移。FSDP和DeepSpeed ZeRO下优化器状态被切片分到不同设备参数更新后状态在rank间并不天然一致需要保存完整状态时尤其要小心否则恢复训练后会出现“检查点里状态尺寸对不上”的怪问题。5.3 断点续训不保存optimizer state等于白续这是我踩过最狠的一个坑。跑大模型到第40000步卡了续训时只从model checkpoint恢复参数没加载optimizer state。结果前几步loss陡增看起来像模型损坏其实是Adam的m和v都清零了更新步长全部重新“猜”。恢复训练后不得不再花上千步让二阶动量重新积累——浪费的时间比重新训练还多。所以续训检查点一定同时存三样东西模型权重、优化器状态、学习率调度器状态。DeepSpeed场景下不只是model_dict还有optimizer.state_dict()和lr_scheduler.state_dict()并且要确认保存时所有rank都参与了存储否则在ZeRO并行下恢复时会读到一个残缺的m/v。经验法则每次保存前跑一次torch.distributed.barrier()并且把checkpoint写到共享文件系统上避免恢复时跨节点复制损坏。5.4 混合精度下的优化器数值故障与排查思路FP16混合精度常见故障就是overflow与NaN。很多人一看到NaN就去查数据其实一大半是优化器更新那一步把数值推爆了。我的排查顺序是先看m和v的状态数值有没有出现Inf再看clip_grad_norm后的梯度是否仍然异常最后才去看数据和label。因为优化器状态里一旦混入Inf后续每一个step都会把NaN传染给所有权重这也是为什么很多框架单独给optimizer state做一层“Inf守卫”检测到异常时跳过当前step。另一个高频问题是梯度裁剪与Adam的兼容。Adam的自适应步长可以让大梯度得到抑制很多人就觉得不需要clip。但loss spike时二阶动量来不及响应仍然会出现单步更新过大。我的习惯是保留梯度裁剪作为安全网不过阈值要放大一些比如max_grad_norm1.0在大模型上会比小模型保守得多有时2.0或5.0更合适。具体值取决于你的模型规模和loss尺度需要自己测一轮。写在最后的实操体会优化器这个东西表面上是公式选择实际上是一整套系统工程。我在跑大模型之后的体会是不要频繁切换优化器先把一个优化器的状态管理、显存布局、调度配合吃透再考虑换新方案。AdamW在今天仍然是性价比最高的起点Muon这类新优化器值得在小规模实验里积极尝试但生产环境切换前一定要把分布式保存、混合精度数值表现都验证干净。最后分享一个每次都能帮到我的小习惯在训练脚本里定期打印优化器状态的内存占用和数值统计m的均值、方差并记录到log。很多问题在loss曲线异常之前优化器状态就已经给出预告了。比如二阶动量v突然下降好几个数量级往往是梯度消失的前兆m如果出现极大值说明某一步梯度异常。盯住优化器比只盯loss能早好几步发现问题。
返回列表