ARTICLE DETAIL

资讯详情

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

大模型训练优化器实战:从SGD到AdamW、EMA与Muon

大模型训练优化器实战:从SGD到AdamW、EMA与Muon 1. 大模型训练里优化器到底在干什么很多人第一次接触大模型训练注意力全在模型结构、数据规模、显卡数量上优化器往往被当成一个“调包就行”的配角。但真到了训练跑不起来、loss 不收敛、显存爆掉、梯度爆炸的时候你会发现十有八九问题就出在优化器这一环。我这些年参与过从几亿参数到百亿参数级别的训练任务踩过的坑里优化器相关的至少占三成。先把话说直白一点优化器就是决定“模型参数往哪个方向走、走多大一步”的那个东西。模型前向传播算出 loss反向传播算出每个参数的梯度但梯度只告诉你“往哪个方向调能降低 loss”它没告诉你“这一步该迈多大”。优化器干的就是这件事——把梯度翻译成实际的参数更新量。你可以把它理解成开车时的油门和方向盘梯度是路况信息优化器是真正踩油门、打方向的那个动作。大模型训练和传统小模型训练在优化器选择上有本质区别。小模型时代SGD 加个动量就能跑得不错因为参数量小、loss 曲面相对简单、训练轮次多。但大模型参数量动辄几十亿上百亿loss 曲面极其复杂存在大量鞍点和平坦区域而且训练成本极高你不可能像小模型那样跑几百个 epoch 慢慢磨。这就要求优化器在有限步数内尽可能高效地找到好的解。这就是为什么 Adam 系列几乎成了大模型训练的默认选择而 SGD 在大模型场景下基本退居二线。这篇文章我会把大模型训练中优化器的核心问题拆开讲从最基础的 SGD 到 Adam、AdamW再到 EMA 权重平均和最近讨论度很高的 Muon 优化器。每个优化器我都会说清楚它解决什么问题、为什么这么设计、实际训练中怎么配参数、以及我踩过的具体坑。适合正在做或准备做大模型训练的同学也适合想搞清楚“为什么大家都用 AdamW”这个问题的朋友。2. 从 SGD 到 AdamW优化器的演进逻辑2.1 朴素 SGD 为什么在大模型上不够用SGD 的更新公式简单到不能再简单参数等于参数减去学习率乘以梯度。就这么一个公式在小模型上能 work但放到大模型上问题立刻暴露。第一个问题是更新方向噪声大。大模型每个 batch 的梯度方差很高因为参数量大、数据分布复杂单个 batch 的梯度方向可能和整体最优方向偏差很大。SGD 完全信任当前 batch 的梯度导致更新方向来回震荡收敛极慢。你可能会想那我把 batch size 调大不就行了理论上可以降低梯度方差但大模型的 batch size 受显存限制不可能无限增大而且 batch size 太大又会降低泛化性能这是个两难。第二个问题是所有参数共用一个学习率。大模型里不同层的参数梯度尺度差异巨大。Embedding 层的梯度和最后几层 attention 的梯度可能差好几个数量级。用同一个学习率要么大的那边爆炸要么小的那边几乎不动。你手动给每层设不同学习率几百层你设得过来吗第三个问题是鞍点。高维空间里鞍点远比局部最小值多SGD 在鞍点附近梯度接近零更新几乎停滞但鞍点并不是好解。SGD 没有机制帮助它逃离鞍点。我早期做过一个对比实验同样的 1.3B 模型同样的数据SGD 加动量跑了两天 loss 还在 3.5 附近震荡换成 Adam 半天就降到 2.8。这个差距不是调参能弥补的是优化器机制本身的差异。2.2 动量与自适应学习率的核心思想SGD 的改进方向有两个一是加动量二是做自适应学习率。这两个思路后来被 Adam 合并到了一起。动量的思路很直观不要只看当前这一步的梯度把历史梯度也考虑进来。具体做法是维护一个梯度的一阶矩估计也就是梯度的指数移动平均。更新时用这个平滑后的梯度而不是原始梯度。这样做的好处是如果某个方向 consistently 有梯度动量会累积更新步长变大如果梯度方向来回变动量会相互抵消减少震荡。你可以把它理解成给优化过程加了惯性像球从山坡滚下来不会因为一个小坑就停下来。自适应学习率的思路是每个参数有自己的学习率根据这个参数历史梯度的大小动态调整。梯度一直很大的参数说明它已经很活跃了学习率应该调小避免过冲梯度一直很小的参数说明它还没怎么被训练到学习率应该调大让它快点跟上。具体实现上Adam 维护梯度的二阶矩估计也就是梯度平方的指数移动平均然后用一阶矩除以二阶矩的平方根得到每个参数的实际更新量。这两个机制合在一起就是 Adam 的核心。它既平滑了更新方向又给每个参数分配了不同的有效学习率。这就是为什么 Adam 在大模型上比 SGD 好用得多——它自动处理了不同参数梯度尺度差异的问题你只需要设一个全局学习率就行。2.3 Adam 的偏差修正与 AdamW 的解耦权重衰减Adam 刚出来的时候有个细节容易被忽略偏差修正。因为一阶矩和二阶矩都初始化为零训练初期这些估计值会偏向零导致更新量偏小。Adam 通过除以一个修正项来补偿让训练初期的更新量也保持在合理范围。这个修正项在训练后期会趋近于 1影响消失。如果你自己实现 Adam 而忘了偏差修正训练初期会明显偏慢这个坑我踩过。但 Adam 有个更隐蔽的问题权重衰减的实现方式。传统做法是把权重衰减加到梯度里也就是 L2 正则化。但在 Adam 里这样做会导致权重衰减的效果和自适应学习率耦合在一起——梯度大的参数权重衰减的实际效果也被放大了。这显然不合理权重衰减应该对所有参数一视同仁。AdamW 的贡献就是把权重衰减从梯度更新里解耦出来直接作用在参数上。具体来说参数更新时先做 Adam 的自适应更新然后单独减去学习率乘以权重衰减系数乘以参数值。这样权重衰减的效果就独立于梯度大小了。这个改动看起来小但在大模型训练里影响很大。我实测过同样的模型和数据AdamW 比 Adam 加 L2 正则的最终 loss 低 0.05 到 0.1而且训练更稳定不容易出现后期 loss 反弹。现在你去看主流大模型训练代码几乎清一色用 AdamW。不是 Adam 不能用而是 AdamW 在权重衰减这件事上更正确长期训练下来优势明显。3. AdamW 在大模型训练中的参数配置实战3.1 学习率、betas 和 eps 怎么设AdamW 有几个关键参数学习率、betas、eps、weight_decay。每个参数都有讲究我一个个说。学习率是最重要的。大模型训练通常用 1e-4 到 3e-4 这个量级具体取决于模型大小和 batch size。模型越大学习率通常要越小因为大模型对更新更敏感。batch size 越大学习率可以适当调大有个经验公式是学习率和 batch size 的平方根成正比。但这不是铁律我见过 7B 模型用 3e-4 跑得很好的也见过 13B 模型用 1e-4 还嫌大的。关键是要做学习率 warmup前几百到几千步从零线性增加到设定值避免训练初期梯度不稳定导致发散。betas 控制一阶矩和二阶矩的指数移动平均衰减率默认是 (0.9, 0.999)。第一个值 0.9 控制动量越大越平滑但响应越慢。第二个值 0.999 控制自适应学习率的记忆长度。大模型训练里有些工作会把第二个值调到 0.95 或 0.98让自适应学习率对近期梯度更敏感。我试过 0.95在训练后期 loss 下降确实更快但前期震荡也更大。如果你训练步数不多比如只跑几万步用 0.95 可能更合适如果跑几十万步0.999 的长期记忆更稳。eps 是个小常数防止除零默认 1e-8。大模型训练里如果梯度很小eps 会影响更新量。有些实现会用 1e-6 甚至 1e-5让数值更稳定。但 eps 太大会导致小梯度参数的更新被过度抑制。我的经验是混合精度训练时 eps 用 1e-6 比较稳纯 fp32 训练 1e-8 就行。weight_decay 通常设 0.01 到 0.1。大模型训练里 0.1 很常见因为模型大、参数多需要更强的正则化防止过拟合。但也不是越大越好太大模型欠拟合loss 降不下去。我一般从 0.01 开始试如果验证集 loss 比训练集高很多再往上调。3.2 学习率调度与 warmup 的配合学习率调度和优化器是绑在一起用的。大模型训练很少用固定学习率通常用 cosine 衰减或 linear 衰减配合 warmup。warmup 的作用是让训练初期稳定。大模型随机初始化后前几步梯度可能很大直接上大学习率容易发散。warmup 让学习率从零慢慢升上去给模型一个适应过程。warmup 步数一般是总步数的 1% 到 5%比如总步数 10 万步warmup 1000 到 5000 步。太少起不到稳定作用太多浪费训练预算。cosine 衰减是训练后期学习率逐渐降到接近零让模型精细收敛。linear 衰减也类似但下降更均匀。我一般用 cosine因为它在训练中期保持较高学习率后期下降平滑实测最终 loss 略好于 linear。但 cosine 有个问题如果你中途要停学习率还没降到位模型可能没完全收敛。所以如果训练预算不确定linear 更保险。这里有个实操细节warmup 期间优化器的二阶矩估计还在积累偏差修正也在起作用所以 warmup 结束时的学习率不要设得太激进。我见过有人 warmup 结束直接上 1e-3结果 loss 直接飞了。稳妥做法是 warmup 结束的学习率就是你设定的最大学习率不要额外放大。3.3 梯度裁剪与优化器的协同梯度裁剪不是优化器的一部分但和优化器配合使用。大模型训练里梯度爆炸是常见问题尤其是训练初期或数据里有异常样本时。梯度裁剪把梯度的范数限制在一个阈值内防止单步更新过大。PyTorch 里用torch.nn.utils.clip_grad_norm_阈值通常设 1.0。这个值不是随便定的太小会抑制正常梯度太大起不到保护作用。我一般从 1.0 开始如果训练稳定就保持如果经常出现 loss spike 就降到 0.5 试试。梯度裁剪和 AdamW 的配合有个细节裁剪是在优化器 step 之前做的裁剪后的梯度再喂给 AdamW。所以 AdamW 看到的是裁剪后的梯度它的一阶矩和二阶矩估计也是基于裁剪后的梯度。这意味着如果裁剪频繁触发AdamW 的自适应学习率会被扭曲。我的经验是如果训练稳定裁剪很少触发影响可以忽略如果裁剪经常触发说明学习率可能太大了应该先调学习率而不是依赖裁剪。4. EMA 权重平均让模型更稳的隐藏技巧4.1 EMA 是什么以及为什么有效EMA 全称是 Exponential Moving Average指数移动平均。在训练里EMA 不是优化器而是一种权重后处理技术。具体做法是维护一份模型参数的影子副本每步训练后影子参数按一定比例向当前参数靠近。公式是影子参数等于衰减率乘以影子参数加上一减衰减率乘以当前参数。衰减率通常设 0.999 或 0.9999。为什么 EMA 有效因为训练过程中模型参数一直在波动尤其是用 Adam 这种自适应优化器时参数更新噪声较大。最终时刻的参数不一定是最好的可能刚好在一个局部波动的高点。EMA 把历史参数平均了一下相当于平滑了训练轨迹得到的参数更接近损失曲面的平坦区域。平坦区域的解通常泛化更好对输入扰动更鲁棒。我在多个项目里对比过用 EMA 的模型在验证集上通常比不用 EMA 的最终 checkpoint 好 0.5 到 1 个点。这个提升不需要额外训练成本只是多存一份参数性价比很高。但 EMA 也有代价它需要额外显存存影子参数对大模型来说这是实打实的开销。7B 模型 fp16 参数大概 14GBEMA 就要多 14GB。如果显存紧张可以用 CPU offload 或者降低 EMA 更新频率。4.2 EMA 衰减率的计算与调优EMA 衰减率不是随便设的它和训练步数有关。衰减率越接近 1平均窗口越长平滑效果越强但响应越慢。如果训练步数少衰减率设太高影子参数还没跟上当前参数训练就结束了EMA 反而拖后腿。有个经验公式衰减率等于 1 减去 1 除以训练步数乘以一个系数。比如训练 10 万步系数取 0.1衰减率就是 1 减 1 除以 1 万约等于 0.9999。如果训练只有 1 万步衰减率应该设 0.999 左右。我一般会设一个 warmup 阶段前几千步衰减率从 0.9 慢慢升到目标值避免训练初期影子参数被随机初始化的参数带偏。还有个细节EMA 应该在优化器 step 之后更新用的是更新后的参数。如果顺序搞反了EMA 平均的是旧参数效果会打折扣。这个坑我在早期实现时踩过当时纳闷为什么 EMA 没效果后来发现是更新顺序错了。4.3 EMA 与模型保存、推理的配合EMA 参数怎么用训练结束后你可以选择保存 EMA 参数作为最终模型也可以保存原始参数。我的做法是两个都存然后在验证集上对比哪个好用哪个。有时候 EMA 好有时候原始参数好取决于训练稳定性和数据分布。推理时用 EMA 参数有个好处模型输出更稳定对输入的小扰动不敏感。这在生成任务里尤其明显用 EMA 参数的模型生成的文本更连贯重复和胡言乱语更少。但 EMA 参数也可能过于平滑导致模型失去一些锐利度在需要精确匹配的任务上反而不如原始参数。所以不要盲目用 EMA要验证。另外如果你用 EMA学习率调度要相应调整。因为 EMA 平滑了参数最终学习率可以比不用 EMA 时稍高一点让模型在训练后期保持一定的探索能力。但这个调整幅度不大我一般不动学习率只在 EMA 衰减率上做文章。5. Muon 优化器值得关注的新选择5.1 Muon 的核心思路与适用场景Muon 是最近在社区讨论比较多的一个优化器全称是 Momentum Orthogonalized by Newton-Schulz。它的核心思路和 Adam 不同Adam 是逐参数自适应学习率Muon 是对梯度矩阵做正交化处理让更新方向更“干净”。具体来说Muon 维护一个动量缓冲区然后对动量矩阵做 Newton-Schulz 迭代近似计算矩阵的正交化。正交化的作用是让更新矩阵的奇异值都接近 1这样每个方向的更新量更均衡不会出现某些方向更新过大、某些方向更新过小的情况。你可以把它理解成给梯度做了一次“白化”去掉了方向上的尺度差异。Muon 适合什么场景目前看它在中小规模模型和特定架构上表现不错尤其是 transformer 的隐藏层矩阵。有实验显示 Muon 在同等计算预算下比 AdamW 收敛更快最终 loss 更低。但 Muon 对超参数更敏感学习率、动量系数、Newton-Schulz 迭代次数都需要仔细调。而且 Muon 目前主要针对二维参数矩阵对 embedding 层和 bias 这些一维参数还是用 AdamW 处理。所以实际训练里往往是混合优化器矩阵参数用 Muon其他用 AdamW。5.2 Muon 与 AdamW 的对比实验观察我最近在一个 1B 级别的模型上做了 Muon 和 AdamW 的对比。同样的数据、同样的训练步数Muon 在前 20% 步数里 loss 下降明显更快但后期优势缩小最终 loss 只比 AdamW 低 0.02 左右。训练稳定性方面Muon 对学习率更敏感我用 AdamW 时学习率 3e-4 很稳Muon 用同样学习率就出现了几次 loss spike降到 1e-4 才稳。另一个观察是 Muon 的显存开销比 AdamW 小。AdamW 要存一阶矩和二阶矩两份状态Muon 只存一份动量显存省了将近一半。这对大模型训练是个不小的优势意味着你可以用更大的 batch size 或者更长的序列长度。但 Muon 的工程实现比 AdamW 复杂Newton-Schulz 迭代有计算开销虽然不大但累积起来也不能忽略。而且 Muon 的分布式训练支持还不如 AdamW 成熟如果你用 FSDP 或 DeepSpeed可能需要自己适配。我的建议是如果你在做探索性训练、追求极致效率可以试试 Muon如果是生产环境、要求稳定可靠AdamW 还是首选。5.3 混合优化器策略的实操建议混合优化器听起来复杂实现起来其实不难。核心思路是给参数分组二维矩阵参数比如 attention 的 Q、K、V、O 投影矩阵FFN 的权重矩阵用 Muon一维参数bias、LayerNorm 的 scale 和 shift和 embedding 层用 AdamW。分组的标准是参数维度。PyTorch 里可以遍历model.named_parameters()根据param.dim()判断。二维的归一组其他归另一组。然后创建两个优化器实例训练循环里分别 step。这里有个坑两个优化器的学习率要分别设。Muon 的学习率通常比 AdamW 小因为正交化后更新量更均匀不需要太大的步长。我一般 Muon 用 1e-4AdamW 用 3e-4。另外权重衰减也要分别设Muon 对权重衰减更敏感通常设小一点0.01 左右。还有个细节梯度裁剪要对所有参数一起做不能分开裁。因为裁剪是按全局梯度范数来的分开裁会破坏范数的全局性。所以流程是反向传播算出所有梯度全局裁剪然后分别喂给两个优化器 step。6. 常见问题与排查技巧实录6.1 loss 不下降或震荡的排查思路loss 不下降是最常见的问题原因可能有很多优化器只是其中一环。我的排查顺序是这样的先看学习率是不是太大。如果 loss 一开始就飞了或者剧烈震荡大概率是学习率太大。把学习率降一个数量级试试如果 loss 开始稳定下降那就是学习率问题。如果降了还是不动看下一步。再看 warmup 是不是不够。如果 loss 在前几百步正常下降然后突然飞了可能是 warmup 结束得太早模型还没适应大学习率。增加 warmup 步数或者降低 warmup 结束时的学习率。然后看梯度裁剪阈值。如果 loss 偶尔 spike 然后恢复可能是某些 batch 梯度太大。把裁剪阈值从 1.0 降到 0.5 试试。如果 spike 频繁可能是数据里有异常样本需要检查数据质量。最后看优化器状态。如果用的是 AdamW检查一阶矩和二阶矩的初始化是否正确偏差修正有没有开。如果自己实现优化器这些细节很容易漏。还有个容易被忽略的点混合精度训练时优化器状态要用 fp32 存。如果用 fp16 存一阶矩和二阶矩数值精度不够更新量会失真loss 下降会变慢甚至停滞。PyTorch 的 AMP 默认会把优化器状态转成 fp32但如果你手动管理精度一定要确认这一点。6.2 显存不足时优化器状态怎么省大模型训练显存不够是常态优化器状态是显存大户。AdamW 要为每个参数存一阶矩和二阶矩fp32 下就是每个参数 8 字节。7B 模型光优化器状态就 56GB加上参数本身和梯度轻松超过 100GB。省显存的办法有几个。最直接的是用 8-bit Adam把优化器状态量化到 8 位显存直接省四分之三。bitsandbytes 库有现成实现效果损失很小我实测 loss 差异在 0.01 以内。但 8-bit Adam 对学习率更敏感可能需要调小一点。另一个办法是 ZeRO 优化把优化器状态分片到多张卡上。DeepSpeed 的 ZeRO-2 和 ZeRO-3 都支持每张卡只存一部分优化器状态显存占用随卡数线性下降。但 ZeRO 会增加通信开销训练速度会慢一些这是用时间换空间。如果显存实在紧张还可以用 CPU offload把优化器状态放到 CPU 内存里需要时再搬到 GPU。但 PCIe 带宽有限频繁搬运会严重拖慢训练速度。我一般只在调试阶段用正式训练还是尽量用前两种办法。6.3 优化器相关的常见问题速查表问题现象可能原因排查方法解决方案loss 一开始就飞学习率太大降学习率 10 倍重试调小学习率增加 warmuploss 震荡不收敛batch size 太小或学习率太大增大 batch 或降学习率调 batch size调学习率loss 下降极慢优化器状态精度不够检查优化器状态是否 fp32用 fp32 存优化器状态训练后期 loss 反弹学习率没降到位检查学习率调度用 cosine 衰减确保后期学习率接近零显存爆了优化器状态太大算一下优化器状态显存占用用 8-bit Adam 或 ZeROEMA 没效果更新顺序错了或衰减率不对检查 EMA 更新时机和衰减率优化器 step 后再更新 EMA调衰减率梯度裁剪频繁触发学习率太大或数据异常统计裁剪触发频率降学习率检查数据质量这张表是我自己训练时总结的基本覆盖了优化器相关的常见问题。遇到问题先查表能省不少时间。6.4 几个容易踩的坑和实操心得第一个坑AdamW 的 weight_decay 不要加到梯度里。有些实现会把 weight_decay 乘以参数加到梯度上这是 Adam 的做法不是 AdamW。AdamW 的 weight_decay 是直接作用在参数上的不经过梯度。如果你用 AdamW 的接口但传了 L2 正则的参数效果会不对。第二个坑学习率 warmup 不要用线性增加到超过目标学习率。有些实现 warmup 结束后学习率会冲过头再降回来这会导致训练不稳定。warmup 的目标就是设定的最大学习率到了就保持或开始衰减。第三个坑EMA 的衰减率不要设成固定值。训练步数少的时候固定 0.9999 会导致 EMA 参数严重滞后。要根据总步数动态计算衰减率或者用 warmup 让衰减率逐渐增加。第四个坑混合优化器时梯度裁剪要在所有优化器 step 之前做。如果先让一个优化器 step 了再裁剪另一个优化器看到的梯度就不对了。第五个坑优化器状态保存和恢复要完整。断点续训时不仅要保存模型参数还要保存优化器状态一阶矩、二阶矩、step 数。如果只恢复模型参数不恢复优化器状态训练会有一个明显的性能下降因为优化器需要重新积累状态。这个坑我在早期做长训练时踩过续训后 loss 直接涨了一截花了半天才找到原因。7. 优化器选型的个人经验总结做了这么多训练任务我对优化器选型的体会是没有银弹只有权衡。AdamW 是目前最稳妥的选择适合绝大多数大模型训练场景参数好调、工程成熟、社区支持好。如果你刚开始做大模型训练直接用 AdamW把学习率、warmup、weight_decay 这三个调好基本能覆盖 90% 的情况。EMA 是性价比很高的技巧几乎不增加训练成本但能带来稳定的泛化提升。我现在的习惯是默认开 EMA除非显存实在不够。EMA 衰减率根据训练步数算不要拍脑袋设。Muon 是值得关注的新方向在特定场景下能比 AdamW 更快更好但工程复杂度和调参难度也更高。如果你在做研究性训练、追求极致效率可以花时间试试。如果是生产环境等 Muon 的生态更成熟再说。最后说一个我自己的习惯每次训练开始前我会先用小模型或小数据跑一个 short run验证优化器配置是否合理。看 loss 前几百步的下降曲线如果正常下降再上大规模训练。这个 short run 花不了多少时间但能避免在大规模训练上浪费几天才发现优化器配错了。这个习惯帮我省了至少几百个 GPU 小时推荐你也养成。
返回列表