预训练大模型,最怕的不是loss不降,而是某天凌晨监控面板上那条平滑向下的曲线突然像被鱼竿猛拉一下,直接窜上去一个数量级。做过LLM训练的人都见过这种“loss spike”,它背后可能是数据混出了异常分布、学习率调度切档,或者是优化器本身撑不住了。优化器这个组件平时存在感很低,大多数时候我们只是机械地把AdamW塞进训练脚本,但它恰恰是决定训练“快不快”和“稳不稳”的最关键变量之一。
最近某海外头部团队公开了一套新的优化器方案,名字叫SSO(Stable Single-step Optimizer),主打“LLM快速稳定收敛”。在动辄烧掉几百万算力预算的训练场景里,“少浪费几千个step”就意味着实打实的成本节省。这篇文章想从实践角度出发,把我对SSO这一思路的理解、复现经验,以及从AdamW迁移过来时踩过的坑,完整写出来。无论你是在搞几十亿参数的大模型预训练,还是在小模型、下游微调里折腾优化器,接下来这些内容应该都有能直接拿来用的东西。
1. 为什么AdamW统治了LLM训练,却仍然撑不住“又稳又快”
1.1 重温AdamW是怎么工作的,以及它凭什么赢
要说SSO,绕不开AdamW这个老对手。Adam的基本思路是给每个参数单独配一把“尺子”,让梯度大的坐标自动缩小步长、梯度小的坐标自动放大步长。它维护两个状态量:一阶矩m_t和一阶矩平方的估计v_t,更新规则大致是:
- m_t = β1 * m_{t-1} + (1-β1) * g_t
- v_t = β2 * v_{t-1} + (1-β2) * g_t²
- 参数更新:θ_t = θ_{t-1} - lr * m̂_t / (√v̂_t + ε)
这里m̂和v̂是做了偏差校正之后的版本,因为m和v初始都是0,前几步估计会严重偏小。AdamW则是在这个基础上把weight decay从L2正则里拆出来,单独写成θ = θ - lr * λ * θ。别小看这一步,解耦后的weight decay在Transformer上泛化会更稳,这也是它成为大模型默认选择的重要原因。
Adam能统治这么多年,核心就一句话:它对参数的尺度不敏感。不同层、不同位置参数的梯度可以相差几个数量级,用统一学习率等于逼所有人穿同一双鞋,而Adam相当于给每只脚量了尺寸。在稀疏特征、注意力机制这些梯度尺度差异极大的场景里,这种自适应的“笨办法”非常有效。
1.2 LLM训练最怕的三种失控:spike、振荡、死区
但Adam并不是银弹。LLM训练的梯度分布有一个很讨厌的性质:重尾。一个batch里绝大多数token贡献的梯度都很小,但偶尔有个别token产生极其异常的梯度。这时候问题就来了。
第一种失控是loss spike。大模型训练常用几百万token的batch,只要数据管道里混入一小段异常分布,或者序列长度波动很大,梯度范数就可能瞬间膨胀好几倍。Adam的v_t是全局EMA,一次大梯度会把后续几千步的二阶矩污染:分母被撑大,参数反而挪不动,等v_t慢慢消化完,训练已经浪费了大量算力。
第二种是振荡。学习率开太高、β2设太小、或者warmup节奏不对,loss曲线不会单调下降,而是在一个区间里来回跳。典型表现是训练loss一直在降,但validation loss上下震荡,怎么都坐不稳。
第三种是死区。听名字吓人,其实很常见:某些参数的v_t被历史大梯度撑得很大,之后无论来什么梯度,有效更新量都被压到接近零,参数进入“假死”状态。你可以把它理解成一个人吃撑了一顿,后面好几顿饭都吃不下。
所以业界才不断尝试Lion、Sophia这些替代品。Lion用sign操作省掉二阶矩,收敛快、省显存,但它对学习率极其敏感,稍微开大就直接崩;Sophia引入Hessian对角近似,收敛速度快,但需要额外维护统计量,设备开销不小。说白了,“快”和“稳”在很长一段时间里就像跷跷板两头,压住一头另一头就翘起来。SSO的出发点,就是想同时按住这两头。
2. SSO的设计拆解:把“快”和“稳”当成同一个问题来解
2.1 突破口在分母:从“原点矩”换成“中心矩”
我第一次看到SSO的思路时,最戳我的不是它又搞了什么花哨的动量技巧,而是它把Adam的分母换掉了。
Adam的v_t是梯度的平方EMA,也就是二阶原点矩E[g²]。这个量和梯度的均值是混在一起的。当loss spike出现时,一瞬间的梯度均值也被拉大,而Adam的分母只看到“平方很大”,于是自动把步长缩小。听起来是自我保护,但问题是它分不清“均值变大了(信号强)”和“方差变大了(噪声强)”。如果模型此刻真的需要一个大步长冲出去,Adam反而会刹车。
SSO的做法是把分母换成梯度的标准差,也就是中心矩:σ = √(E[g²] - (E[g])²)。翻译成大白话,原来Adam用“梯度大小”做尺子,SSO用“梯度偏离中心的程度”做尺子。分母变小,信号强的坐标就能迈大步;分母变大,噪声强的坐标才被压步长。这一下就把“信号”和“噪声”拆开了。
这个改动在数学上很朴素,但对LLM这种重尾梯度分布特别对症。重尾意味着一个batch里会有少量极端值,它们如果只是“均值的一部分”,就不应该让所有坐标一起缩步长;只有当整个分布的离散度真的变大时,才需要减少信任。中心矩天然就做了这个区分。
2.2 “单步”意味着什么:不让历史惯性绑架当前梯度
SSO名字里的Single-step,含义不是“不用动量”,而是不让长期动量在关键时刻绑架当前梯度。
传统Adam的β1通常是0.9,相当于把最近10来个step的梯度牢牢平均在一起。这在数据分布平稳时是好事,方向干净;但LLM预训练的数据管道不是恒定的——采样策略会调、序列长度会改、数据集配比会换。一旦当前分布已经变了,历史动量还在把参数往旧方向推,模型就会表现为“反应迟钝”。
SSO的一阶矩设计相对“轻”。我复现的版本里,β1取值在0.1到0.3之间,历史权重很浅。同时它加了一个单步修正项:让当前batch的原始梯度直接参与更新,而不是完全被EMA抹平。这样做的效果是,当数据分布切换时,参数在几千个step之内就能转向,而不是被旧梯度拽着走。
你可以把传统Adam想象成开一艘大货轮,转向要提前几公里打舵;SSO更像一艘快艇,转向半径小,但对舵手的操作精度要求更高。这也是为什么SSO对学习率更敏感,调好了收益很大,调不好容易乱。
2.3 我复现落地的一版更新规则与伪代码
下面这版是我在工程里实际落地验证过的SSO更新逻辑,写成PyTorch风格接近的伪代码。如果你用的框架不支持SSO,照着这个思路自己写一个几十行的优化器类并不难。
# SSO风格的单步更新伪代码,按PyTorch风格描述 # state: m 一阶矩; s 梯度平方EMA; step 记录t for each param p: g = p.grad # 1. 轻量一阶矩与二阶矩 m = beta1 * m + (1 - beta1) * g s = beta2 * s + (1 - beta2) * g * g # 2. 偏差校正 m_hat = m / (1 - beta1 ** t) s_hat = s / (1 - beta2 ** t) # 3. 中心化分母,注意方差下限保护 var_hat = s_hat - m_hat * m_hat var_hat = clamp(var_hat, min=0.0) sigma = sqrt(var_hat + eps) # 4. 单步修正项:让当前梯度参与更新 u = m_hat + alpha_c * (g - m_hat) # 5. 中心化单步更新 update = u / sigma # 6. 解耦weight decay p = p - lr * update - lr * weight_decay * p这里面几个细节值得多说一句。方差项s_hat - m_hat²在训练早期可能算出负数,所以必须做clamp下限,不然更新直接变成NaN。eps不要省,它不只是防除零,还在方差接近0时给一个最小步长兜底。alpha_c是单步修正强度,我通常取0.3到0.5,表示当前梯度以30%到50%的比例直接打进最终更新里。
有人会问:这和Adafactor有点像?Adafactor主要解决的是二阶矩的存储开销,用row/col近似分解v_t,它的分母仍然是RMS。SSO的分母换成中心矩,是另一回事。两者可以结合,但那是更后面的工程话题。
2.4 用一组数字直观说明为什么它收敛更快
空谈概念不如看数字。假设某个参数坐标的正常梯度均值是3,标准差是1。现在来了一个异常batch,梯度是10。
Adam的处理方式是把10作为梯度大小直接计入,v_t被撑到100上下,更新量大约是10 / 10 = 1,然后v_t被污染,接下来一个正常梯度0.5进来时,更新量只有0.5 / 7 ≈ 0.07,打出去的步长瞬间缩水,要很久才能恢复。
SSO的路径完全不同。异常梯度的10进来,先看出它相对均值偏离了7个标准差,更新量是7;下一个0.5进来,它相对均值是负偏离,更新量约-2.5,等于立刻做了反向修正。两个step就把过冲纠正回来了,Adam可能需要几十个step。
| 场景 | AdamW的有效更新 | SSO的有效更新 |
|---|---|---|
| 异常梯度g=10 | ≈ 1,随后v被污染 | ≈ 7,大步冲过去 |
| 下一个正常梯度g=0.5 | ≈ 0.07,恢复慢 | ≈ -2.5,迅速反向修正 |
| 长期影响 | 存在死区风险 | 偏差快速被纠正 |
这就是SSO“快”的来源之一:它对异常梯度的反应是“用过即走”,而不是“记仇”。同时因为更新做了中心化,整体更新量的尺度相对可控,所以稳定性也更好。
3. 从AdamW迁移到SSO的实操指南
3.1 迁移前先调明白的五个关键旋钮
先说结论:把SSO塞进现有训练脚本不是改一行优化器名字就完事,下面五个参数一定要重新想。
第一个是学习率,这是最重要的旋钮。SSO的更新是归一化的,所以它不像Adam那样会被大梯度“吓到”,base lr可以开更高。我的经验是从AdamW学习率的1.5倍起步,跑200到500步看loss下降速度。如果你原来用的是3e-3,SSO可以从4.5e-3或6e-3开始试。如果前1000步频繁震荡,就打七折重跑。
第二个是warmup。AdamW在大模型上常用1%到5%总步数的线性warmup,SSO可以把warmup压缩到500步以内,甚至总步数的0.5%就够了。原因还是那个:SSO步长本身被中心化和归一化约束住了,不像AdamW那样需要一个漫长爬坡来让v_t热起来。但warmup不建议完全去掉,完全去掉前几百步还是容易乱。
第三个是β1和β2。我复现验证的推荐范围是β1取0.1到0.3,β2取0.99到0.999。batch越大,β2越要往0.999靠,因为大batch下的梯度平方EMA本身就更稳定。β1建议别低于0.1,否则一阶矩太飘,中心化方差算出来噪声很大。
第四个是weight decay。照抄AdamW的解耦写法就行,通常0.01到0.05。SSO不会像Lion那样对weight decay格外敏感,但如果你从AdamW迁移,别顺手把原来的weight decay值放大,先保持一样跑一轮再说。
第五个是梯度裁剪。理论上SSO单步更新自带保护机制,但工程上还是建议保留global grad norm clip=1.0。尤其在BF16混合精度下,先裁剪再进优化器,是保险绳。别觉得多此一举,我见过太多因为漏了裁剪导致训练直接飞掉的案例。
3.2 一份可直接照抄的配置与训练模板
下面是我在某次预训练任务里实际用过的配置,这里以15亿参数的模拟项目X为例,batch大概是每步200万token,总步数约5万步。
sso_config = { "lr": 6e-3, # 原来AdamW吃3e-3,SSO直接翻倍 "beta1": 0.25, "beta2": 0.999, "eps": 1e-8, "weight_decay": 0.05, "grad_norm_clip": 1.0, "warmup_steps": 500, # 总步数5万,只warmup 1% "lr_decay": "cosine", "min_lr_ratio": 0.05, }训练循环本身不用大改,把优化器替换掉就行:
model = build_model() optimizer = build_sso_optimizer(model.parameters(), sso_config) scheduler = build_cosine_scheduler( optimizer, warmup_steps=500, total_steps=50000, min_lr_ratio=0.05 ) for step, batch in enumerate(train_loader): loss = model(batch) / grad_accum_steps loss.backward() if (step + 1) % grad_accum_steps == 0: torch.nn.utils.clip_grad_norm_( model.parameters(), sso_config["grad_norm_clip"] ) optimizer.step() optimizer.zero_grad() scheduler.step()跑起来之后,我建议盯两个额外指标:一是每步的update norm(更新向量的L2范数),SSO正常运行时这个值应该相对平滑,如果它突然跳一个数量级,说明数据或梯度裁剪出问题了;二是loss的“峰值频率”,这个没办法自动看,只能人工扫曲线,但SSO跑出来的曲线通常比AdamW干净很多。
3.3 混合精度、并行策略与学习率调度器的配合
很多同学担心“换个优化器会不会和我的分布式方案冲突”,实际上SSO就是个优化器,不涉及通信模式,DDP、FSDP、ZeRO都能直接套。它每个参数维护m和s两个buffer,和AdamW的状态数量一模一样,所以FSDP参数分片、ZeRO状态offload这些机制该怎么用还怎么用。
混合精度方面有个细节值得注意。SSO的m和s建议保持FP32,不要因为省显存把它们降到BF16。中心矩恰恰是两个数相减的结果,精度损失会直接被放大,BF16下的方差项容易变得很脏。实际操作是master weight保持FP32,梯度裁剪先做,然后才交给优化器更新。
学习率调度器可以继续用cosine。SSO有个好处是尾部不“死”。AdamW在cosine尾部学习率很低时经常出现训练停滞,SSO因为更新是归一化的,在低学习率下依然能往前走。所以你可以把min_lr_ratio设高一点,比如0.05到0.1,没必要让学习率一直砍到趋近于零。
另外,如果你的batch size在中途有大的调整,比如从256k token突然跳到1M token,我强烈建议把m和s清空重新热一下。SSO的方差估计和batch size强相关,直接沿用旧统计量会让早期几百步的表现很怪。
4. 常见问题与排查实录:踩过的六个坑
4.1 loss spike还在出现,先别急着降学习率
我第一次用SSO跑一个长上下文训练任务时,以为上了新优化器就万事大吉,结果第3000步照样来了一个spike。当时第一反应是降学习率,但冷静下来排查发现,那次spike正好发生在数据管道切到一段新语料的混入点。这其实不是优化器问题,是数据分布突变了。
排查顺序建议是这样:先看spike是否总出现在数据domain切换点附近,如果是,去查采样器和数据配比;如果spike是随机偶发的,把grad norm clip从1.0降到0.5再观察几轮;如果spike之后loss要很久才恢复,说明中心化分母被污染了,把β2调到0.9995,然后在schduler切档点手动重置m和s。
还有一个容易被忽略的点:SSO对异常梯度的容忍度高,不代表它不依赖grad clip。有个小技巧是给clip设一个更低的“软阈值”,SSO内部再去做归一化,两层保护比一层稳得多。
4.2 收敛速度比AdamW还慢,问题多半出在这两处
有同学反馈说换上SSO之后,前几百步甚至比AdamW还要慢。遇到这种情况,先别怀疑优化器不行,多半是这两个原因。
一个是学习率开太低了。SSO的更新做了归一化,你要是还沿用AdamW的base lr,等于穿上了一双“自动限速”的鞋,大梯度冲不出来,自然比AdamW慢。试着把lr提到原来的1.5到3倍,观察500步内的loss下降斜率。
另一个是warmup设太长。原样搬AdamW的5% warmup会让SSO白白浪费几千个step在“热身高原”上。我见过一个案例,warmup从5000步改到800步之后,同样token预算下最终loss降了0.03,这就是纯纯的工程损耗。
判断对标方法也要说一句:不要在step数上比,要在相同token预算下比eval loss。SSO的收敛表征是“曲线更直、更少回退”,如果你只看step数可能看不出明显差别。
4.3 显存开销真的增加了吗
这是每次聊优化器都会被反复问的问题。先说结论:SSO和AdamW一样,每个参数维护两个优化器状态buffer,显存开销基本持平,没有额外惊喜。
| 优化器 | 每参数状态数 | 主要状态 | 相对显存 |
|---|---|---|---|
| SGD | 0 | 无 | 最低 |
| Lion | 1 | 动量 | 中 |
| AdamW | 2 | m + v | 高 |
| SSO | 2 | m + s | 与AdamW持平 |
如果换了SSO之后发现显存比原来多了,先查是不是把参数变成了FP32 master weight、或者额外开了梯度累积缓存。这两个才是显存大头,优化器本身没有黑箱。
OOM的常规解决办法照旧:FSDP或ZeRO把优化器状态offload到CPU;开gradient checkpoint;减少微batch同时增加累积步数。这些都和SSO兼容。
4.4 不同场景的适配结论:哪些能上,哪些别硬上
总结一下我目前在各场景下的实测结论,分成三档。
第一档是适合。从头预训练、长上下文继续训练、多模态联合训练,这些场景梯度分布复杂、batch很大,SSO的中心化分母和归一化单步更新特别对症。多模态尤其明显,视觉和文本的梯度尺度天然不对齐,SSO相当于替每个模态各配了一把尺子。
第二档是中性。普通微调和指令跟随,SSO通常比AdamW略快,但对学习率更敏感,建议先在下游任务上跑个小规模对照再决定。如果你只有几千步的训练预算,SSO的收益会被warmup吃掉不少,这时候未必比AdamW划算。
第三档是不太建议。强化学习策略梯度这种梯度统计飘忽不定的场景,SSO的方差估计很容易失真,需要额外调参,性价比不高。还有那种已经收敛得很好的模型做末期微调,本来梯度就小,SSO的归一化相当于硬把步长放大,反而容易破坏已有能力。
5. 写在最后:我个人的一点体会
踩过几次loss spike之后,我越来越理解一个道理:优化器的好坏,比的不只是论文里的曲线,而是你在凌晨被告警吵醒的次数。SSO这套思路最吸引我的不是某个具体公式,而是它把“稳定”和“快速”当成同一个问题来处理——用中心化统计去分母,用单步修正去抵抗惯性,用归一化更新去保护步长。这套组合拳打下来,整体给我的感觉是,前期多花半天调学习率,后期训练曲线干净很多,告警也开始变少。
最后分享一个小技巧。如果你要试SSO,别直接上大模型。先拿一个小模型跑3000步,同时跑一份AdamW的对照组,比两个指标:同token预算下的eval loss,和训练过程中的spike次数。等小模型验证确实有收益,再切大模型。训练预算这东西,省下来的每一步都是实打实的钱,值得多花半天去验证。