最近优化器圈子又冒出一个让我很感兴趣的工作某头部科技公司放出一个面向大语言模型训练的新优化器方案名字起得相当克制就叫 SSO。做 LLM 训练的朋友都知道优化器这个话题平时不太起眼但真正跑过几次大规模训练之后谁都不会小看它。SSO 的全称我倾向于理解为 Stable Step Optimizer核心目标可以概括成两个词快速稳定。翻译得更直白一点就是让模型在训练过程中收敛得更快同时尽量不出现 loss spike、梯度爆炸这类让人血压飙升的问题。这玩意儿解决的是实实在在的痛点。现在跑一个中等规模的语言模型动辄几千张卡跑几周每一步都在烧钱。同样的数据量AdamW 可能要 30 万步才能收敛如果换一个优化器能压到 25 万步省下来的就是几百万的算力成本。反过来如果一味追求快训练到一半 loss 突然飞了前面的工作全部作废这比慢一点更难受。SSO 想做的是在速度和稳定性之间找一个更好的平衡点而不是像某些激进优化器那样只盯着收敛步数一个指标。这篇文章我会从三个角度展开先拆解 LLM 训练里优化器到底卡在哪里再分析 SSO 这类方案的设计思路最后给出一个可以直接照着改代码的落地参考。不论你是正在用 AdamW 调模型的算法工程师还是在自研训练框架的工程团队这篇都值得花十分钟看完。1. 为什么 LLM 训练会把优化器问题放大得这么明显1.1 收敛慢不是玄学是优化器状态在拖后腿很多人觉得 LLM 训练慢主要是因为模型大、数据多优化器只是配角。这个理解在早期小模型时代没错但到了十亿以上参数规模优化器对收敛速度的影响会呈指数级放大。以目前最常用的 AdamW 为例它维护两组状态变量一阶动量 m 和二阶动量 v每个参数都要额外存两份和参数等大的浮点数。一个 7B 参数的模型光是优化器状态就要占 56GB 显存按 BF16 计这还没算模型参数本身和梯度。更大的问题是AdamW 的更新量近似等于 m 除以 sqrt(v) 的平方根这个机制本质上是对每个参数的学习率做归一化。听起来很完美但实际训练中语言模型的梯度分布极其不均匀有些参数长期梯度很小有些参数梯度大得离谱。AdamW 的固定 beta2 只能对所有参数用同一个滑动窗口去估计梯度量级导致它对学习率非常敏感需要精心设计 warmup 和 decay 才能跑稳。LR 稍微调大一点训练前期直接炸调小一点训练速度慢如蜗牛。我见过不少团队在 7B 模型上花了整整一周时间最后发现只是在调学习率真正的模型结构一个都没改。这种“调参玄学”的背后本质上就是优化器自适应能力不够。1.2 稳定性问题loss spike 一次几小时白跑LLM 训练的第二个痛点是稳定性。什么叫不稳定就是 loss 曲线在某个 step 突然从 2.3 飙到 8.7之后要么慢慢降回来要么彻底发散回不去。这种情况通常出现在刚刚加大学习率、更换数据分布、或者某个 batch 里混入了异常样本的时候。最坑的地方在于loss spike 往往毫无预兆。你盯着训练日志看了三天一切正常然后凌晨三点一个 spike 出现等工作早上醒来发现已经在这个错误状态上继续跑了几百步既浪费了算力还可能污染了 checkpoint。传统的对抗手段无非是梯度裁剪和更低的学习率但这两个方案本质上都是“自残”把学习率调低训练速度肉眼可见地下降把梯度裁剪阈值调紧又会让小梯度参数的学习信号被压制。在 CV 领域模型相对小数据相对干净稍微不稳定一点还能容忍。但 LLM 的训练成本是 CV 模型的几十倍一次 spike 的代价可能就是别人一个月的工资这逼着优化器必须具备更强的自稳定能力。1.3 显存账本优化器状态也是大头第三个容易被忽略的问题是显存效率。现在大模型训练普遍用 BF16模型参数和梯度勉强能塞进显存但 AdamW 额外的两组状态经常成为压垮骆驼的最后一根稻草。训练 13B 模型时activation checkpointing 已经把激活内存压到极致结果优化器状态一加batch size 就只能从 2048 降到 1024吞吐量直接打对折。有些团队尝试用 Adafactor 或者 SM3 这类省内存的优化器但省下来的显存往往靠牺牲收敛质量来换。8-bit 优化器虽然好一点但离线——也就是超参和实现没做好的时候数值溢出问题能把人折磨疯。所以任何新优化器如果真的能做到“更快更稳”同时还附带一点显存优势对工程团队来说吸引力是非常大的。2. SSO 想解决什么从设计目标反推技术路线2.1 名字透露的信息稳定、快速、迭代式我不太确定 SSO 官方全称是什么但从命名气质来看大概率离不开三个关键字Stable稳定、Step步长、Optimizer优化器。有些类似的工作还会把 Second-Order 塞进去变成 Stabilized Second-Order Optimizer那就更激进了直接对标 Sophia 这类用二阶信息加速收敛的路线。不管全称是什么SSO 的设计目标是可以从名字反推出来的。第一它必须守住稳定的底线不能像某个知名激进优化器一样收敛是快但十次训练有五次发散。第二它要比 AdamW 收敛得更快不然没有任何换掉 AdamW 的理由。第三它要适合“迭代式训练”这种场景——不是在一个固定数据集上从头训一遍就行而是模型要反复接数据、洗数据、扩展词表每一步都可能改变梯度分布优化器需要有足够鲁棒性。这就决定了 SSO 不会是一个简单的加减法而是要把多套机制组合在一起在一个整体框架里做权衡。2.2 快速收敛单靠动量不够得让更新方向更聪明AdamW 已经证明了一阶动量加上二阶矩归一化可以做出不错的基准线。想再快要么从方向入手要么从步长入手通常两条腿一起走。方向层面Sophia 已经给出了一个思路在梯度更新时引入对角 Hessian 信息相当于知道每个参数附近的曲率大小遇到陡峭的地方自动绕开。这样做的好处是能减少优化器在“山谷”里来回震荡的步数。SSO 如果走了这条路线大概率会采用某种简化的曲率估计比如通过梯度的滑动窗口方差来近似曲率信号而不是真去算二阶导数——真算 Hessian 的话内存和算力都吃不消。步长层面AdamW 的不足在于它对每个参数一视同仁地做归一化导致实际步长受限于梯度的局部爆发力。SSO 可以做得更聪明引入一个“方向一致性”判断当短周期动量方向和长周期趋势方向高度一致时说明当前走的方向很可能是对的允许放大步长当两者背离时说明环境变化剧烈立刻压缩步长避免踩空。2.3 稳定机制方向分歧检测与步长压缩这就是 SSO 和 AdamW 最核心的区别。AdamW 的更新公式里没有任何“环境判断”机制它只相信历史梯度的滑动平均不管当前局面的梯度方向是否发生了根本性变化。而 LLM 训练中数据分布可能因为某个语料桶的切换而突变模型某一层的输出分布可能因为前面层参数的累积变化而漂移这些都会导致梯度信号在几百步内剧烈转向。如果优化器在这个时刻依然保持原有步长结果就是大角度转向时过冲loss spike 随之而来。SSO 的稳定性机制可以设计为维护一个慢尺度动量比如 beta0.99 的滑动平均让它代表“长期趋势”再维护一个快尺度动量beta0.9代表“短期方向”。每次更新之前计算这两个方向向量的余弦相似度。高度一致时说明趋势清晰步子可以迈大一些高度背离时说明模型当前处于方向转折点步子必须收小甚至几乎原地踏步等待方向明朗。这个思路跟人开车很像。你走在一条笔直的公路上可以放心踩油门但前面突然出现急转弯第一反应肯定是先减速等看清弯道方向再加速。SSO 本质上就是给优化器装了一个“急转弯检测器”。2.4 和 AdamW、Lion、Sophia 的关系把 SSO 放在现有优化器坐标系里看会更清楚。AdamW 的特点是稳健但中庸是行业默认基线Lion 的特点是省内存、收敛快但稳定性和泛化表现经常忽上忽下Sophia 的收敛速度最快但实现复杂度高需要每若干步就估算一次 Hessian工程端并不友好。SSO 的位置应该介于 AdamW 和 Sophia 之间。它很可能不像 Sophia 那样引入完整的二阶估计而是借二阶方法的曲率感知思想用一种更便宜的方式获得类似效果。同时它保留了 AdamW 的稳定底座只是在关键节点增加了方向判断和步长压缩机制。这样做的好处是实现在工程上不会过度复杂又能有效规避 Lion 那种“一根筋加速”的风险。3. 实操落地在训练脚本里替换成 SSO3.1 先跑小规模对照实验再做决定不管 SSO 的宣传效果多好我都不建议你直接在一个大模型训练任务上做 AB 测试。优化器替换的验证成本极高一旦训练到一半发现效果不行既浪费了算力又搞乱了进度。比较务实的路径是先在小规模语言模型上做对照实验用同样的数据、同样的模型结构把 AdamW、Lion、SSO 三组跑起来对比 loss 下降曲线和稳定性指标。以 500M 左右的小模型为例用几十亿 token 的小语料跑三组实验单机四卡大概两三天就能完成。这个规模的结论基本能预判大模型场景的表现尤其是优化器对梯度分布的适应能力规模之间的差异远没有数据分布和超参设置的影响大。做完对照再决定是否在大规模训练中切换成本完全可控。3.2 一段可用的 SSO 示意实现在真正拿到官方源码之前我按照上面的设计思路写了一个简化版 SSO 优化器方便在 PyTorch 里直接跑起来感受效果。先说清楚这是示意实现不是官方实现核心逻辑是“快尺度动量 慢尺度动量方向一致性检测 更新范数限制”用来帮助你理解优化器内部发生了什么。以下代码基于 PyTorch 2.x适配torch.optim.Optimizer接口可以无缝替换现有训练脚本中的AdamW调用位置import torch from torch.optim import Optimizer class SSO(Optimizer): 简化版稳定步长优化器SSO 核心思想用快慢两套动量判断当前更新方向是否可靠 方向一致时放大步长方向背离时压缩步长并限制单步更新范数。 def __init__( self, params, lr3e-4, betas(0.9, 0.95, 0.99), eps1e-8, weight_decay0.1, max_update_norm1e-3, ): if not 0.0 lr: raise ValueError(fInvalid lr: {lr}) if not 0.0 eps: raise ValueError(fInvalid eps: {eps}) if not 0.0 max_update_norm: raise ValueError(fInvalid max_update_norm: {max_update_norm}) defaults dict( lrlr, betasbetas, epseps, weight_decayweight_decay, max_update_normmax_update_norm, ) super().__init__(params, defaults) torch.no_grad() def step(self, closureNone): loss None if closure is not None: with torch.enable_grad(): loss closure() for group in self.param_groups: beta1, beta2, beta3 group[betas] for p in group[params]: if p.grad is None: continue grad p.grad state self.state[p] if len(state) 0: state[m] torch.zeros_like(p) # 快尺度动量 state[s] torch.zeros_like(p) # 慢尺度动量 state[v] torch.zeros_like(p) # 梯度二阶矩 state[step] 0 m, s, v state[m], state[s], state[v] state[step] 1 # 快尺度动量捕捉短期梯度方向 m.mul_(beta1).add_(grad, alpha1 - beta1) # 慢尺度动量捕捉长期趋势 s.mul_(beta3).add_(grad, alpha1 - beta3) # 二阶矩估计保留 AdamW 的自适应能力 v.mul_(beta2).addcmul_(grad, grad, value1 - beta2) # 方向一致性快慢动量余弦相似度 m_flat m.view(-1) s_flat s.view(-1) cos_sim torch.dot(m_flat, s_flat) / ( m_flat.norm() * s_flat.norm() group[eps] ) stability 0.5 * (1.0 cos_sim.clamp(-1.0, 1.0)) # stability 越接近 1方向越一致 # 更新方向 denom v.sqrt().add_(group[eps]) update m.div(denom) # 动态步长方向一致时接近满步长背离时压缩 update.mul_(0.2 0.8 * stability) # 限制单步更新范数防止极端情况下参数跳变 up_norm update.norm() max_norm group[max_update_norm] * (p.numel() ** 0.5) if up_norm max_norm: update.mul_(max_norm / up_norm) # 权重衰减decoupled style p.mul_(1.0 - group[lr] * group[weight_decay]) p.add_(update, alpha-group[lr]) return loss这段代码可以直接放进训练脚本里替换原有的optimizer torch.optim.AdamW(model.parameters(), lr...)其余训练逻辑完全不用动。从实测感受来看简化版最明显的变化是loss 下降曲线比 AdamW 更早进入平稳期同时梯度范数的波动幅度明显小了。3.3 关键超参怎么设lr、beta、clip、warmupSSO 的超参虽然看起来多但大部分可以从 AdamW 的配置迁移过来。最需要仔细调的是 lr、beta3 和 max_update_norm 这三项。下面是我在多个模拟实验里总结出的经验值不同模型和数据场景需自行微调。超参建议范围作用说明lr1e-4 ~ 3e-4全局步长LLM 场景建议从 3e-4 往下调beta10.9快尺度动量维持 AdamW 常用值beta20.95 ~ 0.98二阶矩滑动平均建议比 AdamW 的 0.999 小一些beta30.99 ~ 0.999慢尺度动量越大趋势越保守weight_decay0.01 ~ 0.1解耦权重衰减沿用 AdamW 设置max_update_norm1e-3 ~ 3e-3单层参数更新范数上限防止过冲beta2 这里特别说一句。AdamW 喜欢把 beta2 设成 0.999但这个值对语言模型来说往往反应太慢梯度量级变化时二阶矩估计要很久才能跟上。SSO 里方向检测已经承担了一部分稳定性压力beta2 可以稍微降低一些让自适应学习率更敏捷。warmup 阶段不要省。不管用什么优化器LLM 训练头几千步都处于预热期梯度统计量还没建立好直接满步长跑很容易出问题。SSO 在 warmup 阶段会自动处于“慢尺度动量尚未成型”的状态方向一致性偏低实际步长会被压缩所以看起来起步偏慢但这是正常的保护机制。3.4 需要盯紧的四个监控指标换优化器的时候只盯着 loss 曲线是不够的loss 是最滞后的指标等你看到它出问题的时候参数已经坏了好几轮。建议在训练日志里同时记录以下四个指标每 100 步打一次点全局梯度范数反映训练稳定性的早期信号梯度范数持续飙升往往是 loss spike 的前兆。实际更新范数也就是lr * ||update||的全局均值它告诉你模型参数每一步实际变化多大。更新范数突然增大说明优化器认为自己进入了高速通道。方向一致性系数SSO 内部计算的 stability 值长期低于 0.5 说明训练方向反复横跳需要检查数据是否出现重复或分布突变。异常 batch 占比当 loss 超过滑动平均值三倍以上时记为一次异常统计频率可以量化稳定性的改善效果。我建议你在训练脚本里加一个简单的日志回调把上面四个指标输出成 JSON 文件然后用现成的可视化工具画到同一张图上。这样能很早定位到问题发生在哪个 step是数据问题还是优化器问题不用每次都翻整页日志。4. 真实复现中踩过的坑与排查速查表4.1 现象一前几步 loss 纹丝不动换上 SSO 之后最常见的第一个现象就是跑了三四千步loss 基本没降比 AdamW 肉眼可见的慢。很多人的第一反应是“这优化器不行”但大概率不是优化器的问题而是 warmup 和 beta2 的配合出了问题。beta2 偏低会让二阶矩估计快速变大自适应学习率的分子被压得很小而 warmup 阶段全局 lr 本身就在爬升两个因素叠在一起有效步长会变得极低。排查方法很简单把 beta2 调回 0.98 以上或者稍微加大 max_update_norm看 loss 是否在两三千步后开始启动。如果还不行就把 lr 从 3e-4 提到 5e-4 试试。注意一次只改一个参数两个参数同时动出了问题根本分不清是谁的锅。4.2 现象二loss 掉得快但突然 spike如果你已经调通 SSO发现它的收敛速度确实比 AdamW 快但某个 step 突然跳高了一下不用太紧张。先看梯度范数日志如果 spike 发生在梯度范数急剧升高的同一 step通常问题出在 batch 数据上而不是优化器本身。解决方案分两步。第一步降低 max_update_norm 一个量级看看 spike 幅度是否变小这能验证是不是单步更新过冲导致的。第二步检查 trigger 位置的数据分布是否有重复 text、异常 padding、乱码内容混进 batch 中。数据问题引发的 spike 没有任何优化器能免疫SSO 能做的只是帮你更快定位不至于一切都要从头排查。4.3 现象三显存不降反升说好的省内存呢如果你看到 SSO 的宣传说比 AdamW 省内存实际一跑发现多了不少这就要看具体实现差异了。AdamW 只维护 m 和 v 两组状态SSO 如果加了慢尺度动量 s就等于多了一整份参数大小的张量显存自然更高。你看到的现象不是 bug而是设计取舍。应对方法有三个。一是把 s 变量放在 CPU 上每若干步同步一次用通信换显存二是直接去掉慢尺度动量用历史梯度的平均值替代节省一组状态三是换用 8-bit 量化存储慢尺度变量精度损失对方向判断的影响通常可以忽略。要记住新优化器的显存表现从来不是自动变好而是要结合实现细节具体分析。4.4 稳定性排查速查表这里给一张我实际排查问题时会拿来对照的速查表基本覆盖了我用 SSO 这类优化器遇到的大部分状况现象优先排查方向推荐处理手段loss 完全不降warmup 过长、beta2 偏低缩短 warmup 到 1000 步以内beta2 调回 0.98后期 loss 缓慢上升lr 过高、weight decay 过大lr 减半weight_decay 降到 0.05 以下梯度范数周期性飙升数据分布突变检查数据 shuffling增加 batch 内数据多样性更新范数持续为 0max_update_norm 设置过小调大到 5e-3或去掉范数限制再看现象训练早期频繁 spikedata pipeline 存在重复样本检查 sampler 是否出现漏 shuffle多卡训练不收敛梯度 allreduce 不完整检查是否使用了正确的平均方式不要二次归一化下游任务效果不如 AdamW优化器只优化了训练 loss对比训练结束时的 checkpoint 数量SSO 可能过拟合 train loss需要提前早停4.5 另外一个容易被忽视的坑学习率调度器的配合很多人换了优化器之后忘记了 scheduler 可能也要跟着换。AdamW 常用 cosine decay 加线性 warmupSSO 因为收敛更快、更稳往往可以搭配更短的 warmup 和更低的最小学习率。实验中发现给 SSO 配上 AdamW 原样的调度器后期 loss 会一直在一个较高位置浮动怎么压都压不下去原因是 cosine 的初始额外余量不够学习率衰减速度跟不上模型收敛节奏。我的建议是换 SSO 时把 warmup 步数缩到原来的 1/3 到 1/2最终学习率比例从0.1 * lr改成0.05 * lr甚至更低。这能让 SSO 在训练后期释放潜力最后一段的收敛速度尤其明显。5. 换了这么多次优化器我沉淀下来的一点心得我自己在项目里换过 AdamW、Lion也自己魔改过几版优化器每次结论都差不多优化器本身只是训练系统里的一小块拼图但它的影响会辐射到学习率调度、数据管线、checkpoint 策略等方方面面。SSO 这类方案真正让我觉得有价值的不是某一个具体技巧而是它把“稳定性”变成了一等公民来设计而不是事后补救。我的建议是如果你手头没有特别着急的大规模训练任务完全可以拿一个小模型试一版 SSO跑通整个验证流程记录好稳定性指标和收敛曲线。等下一次大任务要启动的时候就有足够的数据支撑决定要不要换掉 AdamW。盲目跟风换优化器是大忌但如果它能用并帮你省下 10% 的训练预算那这时间花得非常值。