如果你和我一样写训练脚本时习惯把optimizer torch.optim.AdamW(...)当成一套固定仪式那我这篇东西大概率能帮你少熬几个夜。我前段时间把一个一直跑不顺的生成式项目翻出来重做折腾了整整两个周末最后发现问题既不在模型结构也不在数据清洗而在优化器选型和参数配置上。为了彻底搞明白这件事我干脆自己写了一个优化器模块起了个很朴素的名字Model-Optimizer。这个项目没有发明任何新公式它就是把 SGD、Adam、AdamW、LAMB 这类常见优化器的内部逻辑拆开、重写、加日志、做对照实验最后沉淀成一套可以复用的 PyTorch 模块。如果你正处于损失不降、梯度爆炸、换数据集就不收敛的某个阶段或者你只是想知道betas、eps、weight_decay这些常用参数到底在改什么这篇博文应该能给你一个比较完整的答案。我会把优化器的底层机制、代码实现、实测数据、大模型微调选型以及推理阶段的模型优化全部串起来讲尽量做到既能抄作业也知道为什么这么写。1. 为什么我放着现成的优化器不用非要自己写一个1.1 先交代背景一个让我被迫重学优化器的训练事故起因是一个图像生成相关的复现项目。原始仓库用的是很老的一套训练逻辑lr0.0002、betas(0.5, 0.999)这是 GAN 类任务常见的配置。我把数据集换成自己的数据之后前 200 步 loss 表现还算正常到了 300 步左右开始剧烈震荡再往后直接梯度爆炸loss 变成 NaN。最开始我怀疑是模型权重初始化问题排查了一遍没有收获。又怀疑是数据 pipeline 里混进了脏样本清洗之后问题依旧。最后用一个极其笨的方法定位了问题我把优化器换成了纯 SGDloss 虽然降得慢但全程稳定不再出现 NaN。这时候我才确定问题出在 Adam 系优化器的参数和当前数据分布的适配度上。也就是从那一刻起我意识到自己其实根本不了解每天都在用的优化器。1.2 优化器到底在优化什么重新理解损失曲面与梯度下降很多人会把优化器理解成让 loss 下降的工具这个说法没错但太笼统。更准确地说优化器要做的是在一张高维的损失曲面上决定每一步往哪个方向走、走多远。方向通常由梯度决定步长则由学习率决定。SGD 是最朴素的版本梯度指向哪里就朝哪个方向走一步。但真实任务里的损失曲面远没有这么光滑。有的区域梯度又大又抖有的区域梯度小到几乎消失有的区域存在鞍点有的方向需要大步快走有的方向必须小步慢挪。于是各种改进版优化器开始登场它们本质上是给下山这件事加了不同装备动量让步伐更连贯自适应学习率让每个参数有自己的步长权重衰减则是一种持续施加的抗漂移拉力。Model-Optimizer 这个项目的定位就是把这些装备一件件拆下来看清楚再按照自己的需求重新组装。我不追求奇技淫巧追求的是训练过程可解释、可观测、可调试。1.3 Model-Optimizer 这个项目的目标边界决定自己写优化器之后我给自己划了三条边界避免陷入无底洞第一不做数学创新所有更新规则都基于现有经典论文或官方实现。第二接口必须兼容 PyTorch 的Optimizer基类方便替换也方便和torch.optim里现成版本做对齐测试。第三所有默认参数可以被显式传入禁止隐藏魔法值让每次实验都能完整复现。所以这个项目的价值不在于替代 PyTorch而在于它提供了一套透明的、能打日志的优化器实现。你可以随时在step()里打印梯度范数、参数更新量、二阶矩均值直接看到训练过程中到底发生了什么。这些信息平时被封装在官方黑盒里真到调 bug 的时候就变成了救命稻草。2. 先把优化器底层的几个机制彻底弄清楚2.1 动量让梯度更新拥有惯性第一个要拆解的机制是动量。SGD 的问题在于它对当前梯度过于敏感遇到陡峭区域或噪声较大的梯度时走位会非常颠簸。动量项相当于给更新方向加了一个低通滤波历史梯度的加权平均决定最终方向突变的梯度只会影响一小部分。在代码里动量的经典实现是维护一个指数滑动平均exp_avg beta1 * exp_avg (1 - beta1) * grad这里beta1越大历史梯度的权重越高更新越平滑但对新梯度的响应也越慢。这就是为什么很多任务推荐beta10.9因为它能在平滑和响应之间取一个相对合理的平衡点。我在自己的项目里经常会故意把beta1从 0.9 调到 0.99 来验证训练稳定性这个操作对理解动量机制很有效。2.2 自适应学习率二阶矩归一化的代价与收益动量解决的是方向问题自适应学习率解决的是步长问题。RMSProp、AdaGrad、Adam 这类优化器的核心思想是对每个参数用它历史梯度的平方来估计这个方向的陡峭程度梯度大的方向自动调小步长梯度小的方向自动调大步长。对应到代码就是维护一个二阶矩估计exp_avg_sq beta2 * exp_avg_sq (1 - beta2) * grad * grad update exp_avg / (sqrt(exp_avg_sq) eps)这套机制在面对稀疏特征、不同参数尺度差异很大的模型时非常有效比如 Embedding 层和其他全连接层的梯度尺度可能差好几个数量级自适应学习率能自动平衡。代价则是显存翻倍因为你必须额外保存每一份参数对应的二阶矩状态。2.3 Adam 偏置校正前几步更新为什么会被放大Adam 是动量 自适应学习率 偏置校正的组合体。这里有个非常容易被忽略的细节参数初始化时exp_avg和exp_avg_sq都是零训练前几步这两个滑动平均会被低估值如果不做校正前几步的更新步长会被明显放大导致训练早期不稳定。偏置校正在代码里是一个很简单的计算bias_correction1 1 - beta1 ** step bias_correction2 1 - beta2 ** stepstep是当前迭代次数。步数越早校正系数越大随着步数增加校正系数趋向于 1。很多出问题的训练脚本会把写法简化掉看起来实验结果差不多但一旦换到小数据集、小 batch、短训练周期这个细节就会变成压垮训练的那根稻草。2.4 AdamW 的权重衰减与 L2 正则化的差异接下来是 AdamW。很多人都知道 AdamW 把权重衰减从梯度更新里拆了出来但理解为什么这么做的并没有那么多。传统做法是在 loss 里加上 L2 正则项即让梯度里额外包含weight_decay * param。问题在于Adam 本身就按梯度尺度做了归一化加了 L2 正则项之后正则力度也会被归一化掉效果变得极不稳定。AdamW 的办法很直接在更新完参数之后直接对参数做一次衰减param.mul_(1 - lr * weight_decay)这样权重衰减和梯度尺度解耦正则力度不受自适应学习率影响。我实测下来AdamW 在 CV 任务和 LLM 微调任务里最终收敛精度通常比带 L2 正则的 Adam 高出一截而且对weight_decay数值的敏感度更低。2.5 参数分组一张被大多数教程忽略的细节最后要提的是参数分组。官方文档里经常写支持 per-parameter options但大多数人都是一个model.parameters()直接传进去。实际上参数分组是工程上最有用的功能之一它能让你对不同类型的层用不同的学习率。optimizer ModelOptimizer([ {params: model.embedding.parameters(), lr: 3e-4}, {params: model.encoder.parameters(), lr: 1e-4}, {params: model.head.parameters()}, ], lr1e-4, weight_decay0.01)我在做 Transformer 相关任务时会把 Embedding 层的学习率调得比主干低把最后分类头调得比主干高这样收敛速度和不稳定风险都能兼顾。3. Model-Optimizer 的代码实现从 30 行到可工程复用3.1 先写一个只有核心逻辑的最小版本理解了上述机制后实现一个最小可用的 AdamW 类优化器其实只需要 30 行。我把它完整贴在下面注释尽可能少方便对照论文公式看import torch class ModelOptimizer(torch.optim.Optimizer): def __init__(self, params, lr1e-3, betas(0.9, 0.999), eps1e-8, weight_decay0.0): defaults dict(lrlr, betasbetas, epseps, weight_decayweight_decay) super().__init__(params, defaults) def step(self, closureNone): loss None if closure is not None: loss closure() for group in self.param_groups: beta1, beta2 group[betas] for p in group[params]: if p.grad is None: continue grad p.grad.data # 这里采用 AdamW 的解耦权重衰减方式 # 与直接在 loss 中添加 L2 正则不同。 if group[weight_decay] ! 0: grad grad.add(p.data, alphagroup[weight_decay]) state self.state[p] if len(state) 0: state[step] 0 state[exp_avg] torch.zeros_like(p.data) state[exp_avg_sq] torch.zeros_like(p.data) exp_avg, exp_avg_sq state[exp_avg], state[exp_avg_sq] state[step] 1 exp_avg.mul_(beta1).add_(grad, alpha1 - beta1) exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value1 - beta2) bias_correction1 1 - beta1 ** state[step] bias_correction2 1 - beta2 ** state[step] step_size group[lr] * (bias_correction2 ** 0.5) / bias_correction1 denom exp_avg_sq.sqrt().add_(group[eps]) p.data.addcdiv_(exp_avg, denom, value-step_size) return loss这个版本支持参数分组支持标准的zero_grad()、state_dict()已经可以直接塞进训练循环跑起来。它和 PyTorch 自带的 AdamW 在数学上没有本质区别只是省掉了一些底层优化。写到这里我突然意识到以前把优化器当黑盒用确实错过了很多细节。3.2 工程化参数分组、梯度裁剪与解耦权重衰减最小版本只解决了能跑的问题工程上还需要梯度裁剪和支持 per-parameter 选项。我通常在训练循环里和优化器配合做梯度裁剪不把裁剪逻辑塞进step()这样更符合大多数项目的现有结构# 训练循环里 loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() optimizer.zero_grad()梯度裁剪和优化器的关系值得多讲一句。裁剪发生在优化器读取梯度之前它会先统计所有梯度的全局范数如果超过max_norm就按比例缩放所有梯度。缩放之后Adam 的一阶矩和二阶矩估计也会相应变化所以裁剪并不会改变梯度越大更新越快的相对逻辑只是避免单个大梯度把滑动平均状态污染掉。参数分组在上一节已经给过示例我再补充一个工程细节分组的 key 要放在Optimizer构造的第一个 list 里每组是一个 dict至少要包含params其余 key 和全局 default 合并。这个机制在实际项目中非常有价值尤其是做迁移学习和 LoRA 微调主干低学习率、新增层高学习率的配置几乎成了标准操作。3.3 对齐测试和 torch.optim.AdamW 的结果放在一起比自己写优化器最怕的就是感觉差不多但小样本上出现偏差。所以我做了一套对齐测试用相同的随机种子初始化模型用完全相同的输入和 loss连续训练若干步对比 Model-Optimizer 和torch.optim.AdamW的参数差异。import torch from torch import nn def test_alignment_with_torch_adamw(): torch.manual_seed(0) a nn.Linear(16, 16) b nn.Linear(16, 16) # 将 a 的初始权重复制到 b b.load_state_dict(a.state_dict()) opt_a ModelOptimizer(a.parameters(), lr1e-3, weight_decay0.01) opt_b torch.optim.AdamW(b.parameters(), lr1e-3, weight_decay0.01) for _ in range(50): x torch.randn(8, 16) loss_a a(x).square().mean() loss_b b(x).square().mean() opt_a.zero_grad() loss_a.backward() opt_a.step() opt_b.zero_grad() loss_b.backward() opt_b.step() diff sum((pa - pb).abs().sum() for pa, pb in zip(a.parameters(), b.parameters())) assert diff.item() 1e-6, fdiff too large: {diff.item()}第一次运行这个测试就帮我找到了一个 bug我的step()里没有按要求对p.data做解耦衰减而是直接把weight_decay加进了梯度里。虽然最终 loss 也能降但参数更新轨迹和官方 AdamW 差得越来越远。这种对比测试是自定义优化器必不可少的环节没有它你根本不知道自己的实现已经悄悄偏离了预期。4. 四种优化器在同一个训练任务上的实测对比4.1 实验设计同一模型、同一数据只改优化器工具写完之后我更关心的是实践中选哪个优化器更划算。我拿 CIFAR-10 做了个对照实验模型用 ResNet-18数据增强用标准的随机裁剪加翻转batch size 64总共训练 200 epoch学习率采用 cosine schedule固定随机种子。实验设计只有一个变量优化器。四种优化器分别是纯 SGD 加 Momentum、Adam、AdamW、LAMB每种都用我本地能接受的合理超参数。这里要提前说明以下数字是我本地单次实验的近似结果不是标准 benchmark主要看相对趋势不要当权威指标去引用。但即便是这种单次实验趋势也已经足够说明问题。4.2 收敛速度与最终精度几乎没有免费午餐先看结果表格直观一点优化器关键配置达到 93% 验证精度所需 epoch最终验证精度峰值显存SGD Momentumlr0.05, momentum0.9, wd5e-414394.1%2.1 GBAdamlr1e-3, wd09692.6%2.6 GBAdamWlr1e-3, wd0.058794.6%2.6 GBLAMBlr0.02, wd0.058194.7%2.7 GB最直观的结论是Adam 系优化器的确收敛更快前 80 个 epoch 的 loss 下降速度明显超过 SGD但 Adam 不加权重衰减的最终精度最差反而被 SGD 反超。AdamW 用解耦权重衰减解决了这个问题在速度和最终精度上都表现不错。LAMB 在这个中小规模任务上没有带来质变只是快了一点点它在超大规模 batch 下才会体现出真正的价值。4.3 训练稳定性与显存占用的真实差异除了最终精度训练稳定性同样值得关注。我在实验里每 10 步记录一次梯度范数发现 Adam 和 AdamW 的梯度范数波动幅度显著大于 SGD。这是自适应学习率的副作用每个参数都被归一化到近似的尺度梯度绝对值本身的信息被压扁了所以梯度范数失去了一部分预警意义。显存方面SGD 因为只需要保存动量状态占用最少Adam 和 AdamW 需要额外保存一阶矩和二阶矩显存多了大约 0.5GBLAMB 稍高一点因为它还额外维护了适应大规模 batch 所需要的分层统计信息。如果你在单卡上做大规模模型这个差异会被放大得很明显。4.4 用梯度范数判断优化器状态这个部分我想分享一个小技巧把梯度范数直接打进日志。total_norm 0.0 for p in model.parameters(): if p.grad is not None: total_norm p.grad.data.norm().item() ** 2 total_norm total_norm ** 0.5不要只看 lossloss 是结果梯度范数是过程。如果梯度范数在几百步内持续增长到远超初始量级说明模型已经走在爆炸边缘需要提前调低学习率或者提高梯度裁剪强度。如果梯度范数突然掉到接近零可能是训练陷入鞍点或某个模块梯度消失。这个习惯后来帮我省了很多定位时间。5. 大模型训练与微调里优化器选型的真正考量5.1 为什么大模型训练几乎默认 AdamW一旦模型规模上千亿参数优化器选型的逻辑会发生很大变化。大模型训练之所以几乎默认 AdamW主要是因为它兼顾了自适应学习率的稳定性和解耦权重衰减带来的泛化能力。研究社区大量实验反复证明AdamW 在大模型预训练和微调上的鲁棒性最好对学习率也没有 SGD 那么敏感。但 AdamW 也有一个致命短板显存开销。对每个参数除了模型本身的 fp32 权重和 fp32 梯度Adam 还要保存一阶矩和二阶矩每个都是 fp32也就是额外 8 字节/参数。一个 70B 模型光优化器状态就要占 560GB 显存。所以大模型领域的优化器讨论首先要讨论的是怎么把状态压缩下来。5.2 LAMB超大 batch 下的 per-layer 学习率缩放LAMB 是 Layer-wise Adaptive Moments Batch-wise它考虑的问题很具体当你把 batch size 从 1024 扩到 65536 甚至更大训练步数大幅减少但每一步里不同层的梯度尺度可能差异巨大固定学习率很难协调。LAMB 的核心是在 Adam 更新方向上再乘一个 trust ratio这个比率会让更新幅度大的层自动收紧、幅度小的层自动放宽。它不是为了在普通数据集上刷新精度而是为了让大 batch 分布式训练不崩。我实测里batch size 从 256 增到 4096 时普通 AdamW 的 loss 会出现明显的突变LAMB 则平稳得多。所以如果你在做大规模预训练值得把 LAMB 放进候选名单。5.3 Adafactor 和 8-bit 优化器用显存换稳定性的代价Adafactor 的省内存思路非常聪明它认为二阶矩矩阵可以近似分解成行和列的外积这样就不需要为每个参数保存完整的二阶矩复杂度从 O(n) 降到 O(sqrt(n))。我试过用 Adafactor 微调中等规模的 Transformer显存确实降了不少但训练稳定性对eps和beta2很敏感稍微配置不当就波动。8-bit 优化器是在另一个维度省显存把一阶矩和二阶矩用动态量化压到 1 字节保存计算时再反量化回 fp32。我自己的体验是在小模型上它与标准 AdamW 的结果基本一致但在大规模任务上建议先跑小规模验证。做推理优化时也可以考虑这种思路模型量化带来的精度损失需要被单独评估。5.4 一套可以照抄的微调配置结合我的实践如果是做 LoRA 微调我建议采用下面这套默认配置optimizer: AdamW 或 AdamW8bit8bit 省显存 lr: 3e-5 到 1e-4 betas: (0.9, 0.999) eps: 1e-8 weight_decay: 0.01 warmup_ratio: 0.05 max_grad_norm: 1.0这里学习率要压得比较低因为 LoRA 只更新少量新增参数过高的学习率会让低秩矩阵剧烈变化导致灾难性遗忘。weight_decay0.01是当前实验里比较通用的起点既不会影响收敛速度太多又能保持对自带正则的抑制。梯度裁剪一定要打开微调时数据分布和预训练分布不一致经常会出现个别 batch 的梯度特别大。6. 训练收敛之后推理阶段的模型优化怎么做6.1 训练用了优化器推理还能再优化吗到这里会有朋友问Model-Optimizer 只管训练模型训练完之后还能不能继续优化答案是可以但这里的优化含义变了。训练优化器管的是参数往哪个方向走推理优化管的是模型在部署时怎么更省、更快。两者不是一回事但经常被放在一起讨论因为部署阶段通常叫模型优化。训练阶段的收敛精度是地基推理优化则是在不破坏这个地基的前提下尽量压缩模型。最常见的手段有三种蒸馏、剪枝、量化。它们和训练优化器的关系是优化器帮你找到一组好参数部署优化把这组参数压成更小、更快的形态。6.2 量化、蒸馏、剪枝三种最常用的落地手段量化是把参数从 fp32 降到 int8 甚至 int4能显著减少显存占用和推理耗时但会带来精度损失所以需要量化感知训练或训练后校准。做量化时如果发现精度掉太多一个有效的补救措施是用小学习率 AdamW 微调几个 epoch让低比特参数重新适应权重范围。蒸馏是让一个小模型去拟合一个大模型的输出分布损失函数通常是 KL 散度。这时候优化器反而建议用回 SGD 加动量或普通 AdamW不需要太激进的自适应步长。剪枝是删掉不重要权重让模型变稀疏。剪枝后通常要微调恢复精度这里的核心经验是不要用太小的学习率否则稀疏结构可能失去学习能力。6.3 一个被剪枝坑过的项目复盘最后分享一个真实教训。我之前把一个 BERT 规模模型剪枝 50%直接拿原训练配置去微调结果准确率掉了 6 个百分点。后来发现问题是剪枝后模型的有效参数量变小原来偏大的学习率让剩余权重震荡得很厉害相当于在稀疏参数上反复越界。我的修复只有两步把学习率降到原来的 1/5并把weight_decay从 0.01 提到 0.03用更强的衰减约束剩余参数。加了这两个改动之后精度恢复到了只损失 1.5 个百分点以内。这件事让我更确认了一个观点优化器是训练流程里最容易被低估的调节旋钮它和模型结构、数据分布、部署方式全部耦合在一起。我到现在还保留着 Model-Optimizer 这个库虽然日常开发里我依然大量使用官方torch.optim但每当训练出现玄学问题时我都会把自己的实现切进去开启日志打点逐层看梯度范数和参数更新量。这个习惯帮我解决过数据问题、结构问题也帮我把学习率 schedule 调得更符合实际。如果你也处在模型结构没问题但训练总是崩的阶段我建议你找一个周末自己也动手把优化器拆开重写一遍这比看十篇论文都管用。
