☰
优化器管理与调参实战:从AdamW到Lion的工程化落地
2026/10/1 14:09:22 网站建设 项目流程

训练深度学习模型时,Model-Optimizer 这类优化器管理工具,往往是被忽视但又极其关键的一环。相信不少人都经历过类似的困惑:模型结构一模一样,数据也没问题,只是换了优化器或调了学习率,训练结果就天差地别,损失曲线要么不收敛,要么震荡得让人崩溃。我自己在做大规模模型训练时,被优化器的参数配置、状态保存、复现问题折腾了很多次之后,索性做了一个统一管理优化器的小项目,名字就叫 Model-Optimizer。这篇文章不是官方文档,也不是算法科普,就是想以我这个项目为主线,聊聊优化器到底在优化什么、我做这个工具时踩过的坑,以及从 AdamW 切换到 Lion 这类实操里真正有价值的细节。

如果你是正在训练自己的模型、被优化器调参困扰的工程师或研究人员,这篇文章应该能提供一些能直接落地的思路。

1. 优化器不是“换个名字”那么简单:一个被低估的训练变量

先说一个最容易被大家忽略的事实:在训练流程里,优化器绝对不是“从 torch.optim 里import一个类”这么简单的事。它决定了你的模型如何从当前状态走向下一个状态,决定了梯度下降的步伐、方向、历史信息的利用方式,甚至决定了这个训练任务能不能收敛到理想精度。

1.1 优化器的本质:它决定你如何“修正”网络

简单打个比方,训练模型就像一个人下山。损失函数是山的高度,模型参数是当前位置,梯度是脚下最陡的方向。但“最陡方向”并不总是直接走下去就最好,因为你可能正站在一个坑的边缘,或者前面是一段凹凸不平的路。这时就需要优化器来帮你决定:每一步走多长(学习率),要不要利用之前的运动惯性(动量),要不要对不同方向区别对待(自适应学习率)。

  • SGD 是最朴素的策略:只看当前的梯度方向,一步一个脚印地走,但很容易陷入局部坑洼或者走得很慢。
  • Momentum 引入了“惯性”:如果过去几个方向都朝同一个方向走,那就加快脚步;如果方向变化剧烈,就小心一点。
  • Adam 则进一步加入了“对不同方向历史信任度”的估计:历史上梯度小但稳定的方向可以走大一些,梯度忽大忽小的方向则要保守一些。

这些策略的差异直接体现在训练曲线上。我用同样的 ResNet 结构在 CIFAR 数据集上做过实验:SGD+Momentum 调好学习率之后精度能到 90% 以上,但如果直接照搬 Adam 的默认学习率 1e-3,前期收敛飞快,后期却很难达到同样的精度。这说明优化器不是“锦上添花”的模块,而是和模型结构、数据预处理处在同等重要的地位。

1.2 各框架API差异带来的真实痛点

在我实际开发和训练的过程中,最头疼的不是理解优化器本身,而是不同框架、不同版本之间 API 的割裂。PyTorch 里的torch.optim.AdamW、TensorFlow/Keras 里的AdamW虽然都叫 AdamW,但参数默认值、权重衰减的实现方式并不完全一致,尤其是“解耦权重衰减”和“L2正则化”这两个概念,在不同优化器实现里经常被混淆。

更麻烦的是参数名的差异。PyTorch 用lr,有些库用learning_rate;权重衰减有的叫weight_decay,有的叫wd,有的实现里还分decay_type。我在追求训练结果可复现时,经常因为优化器参数记错、写错而浪费一整天。我甚至遇到过这样的问题:同一个训练脚本,换了一台机器、换了一个依赖版本,优化器状态字典(state_dict)里的 key 结构居然不一样,导致断点续训直接报错。

1.3 为什么需要一个Model-Optimizer而不是直接用torch.optim

真正促使我做这个项目的场景是这样的:我在做一次多组对比实验,要依次验证 SGD、AdamW、LAMB、Lion 这四种优化器在同一个模型上的表现。如果直接用 torch.optim,我需要为每种优化器写一段几乎重复的训练循环代码,还要小心翼翼处理学习率调度器、梯度裁剪、EMA 这些外围配套逻辑。代码复制来复制去,很容易出现“上一个优化器的参数残留到下一个”的问题。

Model-Optimizer 的核心思路,就是把“优化器的创建、参数配置、状态管理、与训练循环的配合”全部统一起来。它不是一个新优化器算法,而是一个优化器管理框架。你声明一次模型结构和训练需求,它负责把合适的优化器、学习率策略、状态保存逻辑都串起来,让实验对比变成“改配置而不是改代码”。这一点在跑对比实验阶段真的能省下大量时间。

2. Model-Optimizer 的核心设计:统一、可复现、能续训

既然要做成工具,就不能只是简单包一层torch.optim的 API。我在设计 Model-Optimizer 时,核心目标有三个:统一的参数入口、完整的状态管理、以及和主流训练组件(学习率调度、混合精度、分布式训练)的无缝配合。

2.1 统一注册与参数别名映射

第一件事是解决“同一个优化器在不同库参数名不一致”的问题。我采用注册表模式加参数别名映射来解决。简单说,所有优化器通过装饰器注册到内部工厂里,外部配置统一使用一套我定义的参数名,再由工厂内部映射到具体框架的 API。

# Model-Optimizer 的核心注册逻辑示意 OPTIMIZER_REGISTRY = {} def register_optimizer(name): def decorator(cls): OPTIMIZER_REGISTRY[name] = cls return cls return decorator def create_optimizer(name, model, lr, weight_decay=0.0, **kwargs): # 参数别名映射,例如 m_lr -> momentum, wd -> weight_decay normalized = normalize_hyperparams(lr=lr, weight_decay=weight_decay, **kwargs) optimizer_cls = OPTIMIZER_REGISTRY[name] return optimizer_cls(model.parameters(), **normalized)

这里要注意一个实际训练中非常常见的需求:不是所有参数都用同一个学习率。比如在微调 BERT 时,希望 backbone 用较小的学习率,而新增的 classification head 用较大的学习率。所以我提供了layers参数来指定参数组。

optimizer = create_optimizer( name="adamw", model=model, lr=2e-5, weight_decay=0.01, layers={ "backbone": {"lr": 1e-5}, "head": {"lr": 5e-5, "weight_decay": 0.0}, }, )

2.2 训练状态管理:断点续训与EMA

优化器状态管理,是 Model-Optimizer 里我觉得最值回票价的部分。很多人以为断点续训只需要保存模型权重,但在训练中途恢复时,如果优化器的动量、二阶动量状态丢失,学习率调度器的步数丢失,训练曲线会发生明显的“突变”。尤其是 Adam 系列,它的二阶动量记录的是历史梯度平方的滑动平均,一丢,重新训练时的有效学习率相当于被重置了,前面的训练节奏全乱。

我封装了一套save_training_state和load_training_state逻辑,统一保存以下内容:

  • 模型权重,这个不用说。
  • 优化器状态字典,包括每个参数的动量、二阶动量。
  • 学习率调度器状态,以及当前 epoch 和 global step。
  • EMA 权重,如果启用了 Exponential Moving Average 的话。
  • loss 曲线记录,方便恢复后对照。
def save_checkpoint(path, model, optimizer, scheduler, ema, epoch, global_step, loss_history): checkpoint = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict() if scheduler else None, "ema": ema.state_dict() if ema else None, "epoch": epoch, "global_step": global_step, "loss_history": loss_history, } torch.save(checkpoint, path) def load_checkpoint(path, model, optimizer, scheduler, ema): checkpoint = torch.load(path, map_location="cpu") model.load_state_dict(checkpoint["model"]) optimizer.load_state_dict(checkpoint["optimizer"]) # 恢复学习率调度器和EMA状态 ... return checkpoint["epoch"], checkpoint["global_step"], checkpoint["loss_history"]

这里有个很实务的点:保存 checkpoint 时建议先把优化器状态转到 CPU 再保存,否则会直接把显存占满。我在 8 卡训练时吃过这个亏,直接torch.save优化器状态把显存撑爆了,后来改成先把state_dict移动到 CPU,再序列化,问题就解决了。

2.3 与学习率调度器、AMP、DDP的协作逻辑

优化器不是孤立存在的。实际训练时,我们几乎总是会叠加学习率调度器(如 cosine decay、OneCycleLR)、混合精度训练(AMP)和分布式数据并行(DDP)。Model-Optimizer 在这块设计上遵循几个原则:

  • 学习率调度器统一由训练框架控制,不在优化器内部“偷偷修改 lr”,这样打印日志时看到的 lr 才是可信的。
  • 混合精度训练下,优化器更新的是 FP32 的 master weight,梯度则是 FP16 的。这时如果自己手动实现优化器,很容易在 warmup 阶段出现数值偏差。所以我默认依赖 PyTorch 自带 AMP 接口,Model-Optimizer 只保证优化器step()被正确调用,以及scaler.scale(loss)的缩放因子能正确作用于梯度裁剪。
  • DDP 场景下,每个进程各持有一份优化器状态,这没问题,但 checkpoint 保存和加载要以 rank 0 为准,然后在加载后broadcast给所有进程。我踩过一次这样的坑:单卡保存、多卡恢复,前几个 batch 的 loss 对不上,就是因为优化器状态没同步。

3. 从 AdamW 切换到 Lion 的完整实操:配置、验证与调参

聊完设计,来点实打实的操作。以我自己用的比较多的一个场景为例:把 Baseline 的 AdamW 切换成 Lion。为什么要拿 Lion 举例?因为在 CV 分类、目标检测这类任务上,Lion 经常比 AdamW 有更好的收敛性和精度,而且在显存占用上有明显优势——它不需要保存二阶动量。

3.1 为什么拿Lion举例:节省显存不代表白嫖性能

先说原理。Lion 是谷歌在 2023 年提出的一种优化器,核心是“符号动量”策略。它和 AdamW 最大的区别在于更新规则:AdamW 会根据梯度的一阶动量m和二阶动量v来计算更新量,而 Lion 只用了一阶动量的符号,配合固定的学习率权重来更新参数。简化来看,Lion 的更新方向是sign(beta1 * m + (1 - beta1) * gradient)。这个“只保留符号、抛弃幅值”的做法有点反直觉,但它在很多任务上确实能带来更好的泛化性。更实在的好处是,Lion 只用维护一个动量,所以省掉了 Adam 里那整整一份二阶动量参数的显存。

不过“省显存”不等于“白嫖性能”。Lion 的默认超参非常敏感,尤其是学习率和 weight decay。我实测的结论是,从 AdamW 切到 Lion,学习率要先缩小到原来的 1/3 到 1/5 再开始调,比如 AdamW 用 3e-4,Lion 可以从 1e-4 起步。

3.2 实际配置代码与切换步骤

在一个典型的分分类任务里,原来的训练配置是这样的:

optimizer = create_optimizer( name="adamw", model=model, lr=3e-4, weight_decay=0.05, betas=(0.9, 0.999), )

切换到 Lion 时,我采用如下的配置:

optimizer = create_optimizer( name="lion", model=model, lr=1e-4, weight_decay=0.01, betas=(0.9, 0.99), )

注意 weight_decay 我直接从 0.05 降到了 0.01。这是我在复现 Lion 论文和实际调参时总结出的经验:Lion 的 weight decay 对最终精度的影响比 AdamW 更加剧烈,过大的 weight decay 很容易把模型“压”得太死,导致验证精度上不去。如果你是从 AdamW 切过来,建议先保持和学习率同样的缩放逻辑,观察 5 到 10 个 epoch 的曲线再做调整。

还要强调一点:warmup 在 Lion 里非常重要。Lion 的更新量依赖于动量的符号,训练早期如果学习率直接拉到最大,很容易出现非常大的参数跳变,loss 会在前几百步里冲上天。我个人的做法是配合 5% 到 10% 的 warmup steps,让学习率线性爬升到目标值。不要觉得这是多此一举,尤其在 imageNet 规模的数据集上,这一步直接决定模型会不会在 epoch 1 就崩掉。

3.3 切换后训练曲线的评估方法:不要只看第一个epoch

切换到新优化器后,最忌讳的就是拿第一个 epoch 的 loss 来下结论。我之前犯过这个错误:用 AdamW 时第一个 epoch loss 从 2.0 降到 0.8,切到 Lion 后第一个 epoch 只从 2.0 降到 1.4,我差点就回滚了。但实际上,再往后跑几个 epoch,Lion 的下降速度会逐渐追上,最终收敛精度反而更高。

正确做法是设定一个固定的评估窗口,比如在相同的训练步数下比较验证集的精度曲线。我在 Model-Optimizer 里专门加了一个compare_experiments的小工具,把两次实验的 loss history、验证精度 history 对齐到相同 step 上,画在同一张图里。这个工具本身代码量不大,但能避免很多“感觉好像更差”的错觉。

4. 训练优化阶段的排错链路:Loss不降、显存溢出与结果不一致

工具做得再顺,训练过程中还是会出现各种异常。这一节是我在真实训练中被折磨过很多次之后整理出的排查链路,希望能帮你少走一些弯路。

4.1 Loss不降或震荡:先检查优化器状态,再怀疑超参数

遇到 loss 完全不下降,很多人的第一反应是把 learning rate 调大或者调小,但我建议先按这个顺序排查:

  1. 确认优化器是否真的拿到了正确的参数组。如果使用分层学习率,检查一下param_groups里的参数数量对不对,而不是只看日志里打印的 lr。我遇到过layers配置的 key 和模型模块名对不上,结果 head 那层根本没被单独设置,还在使用默认 lr 的情况。
  2. 确认是否存在梯度为 0 的情况。如果某个模块没有参与 loss 的计算,优化器会一直对它的参数做 weight decay,这会导致这部分参数被推到接近 0,但 loss 一点不动。
  3. 确认是否需要梯度裁剪。特别是切换到 Lion 这类“签名更新”优化器后,梯度范数的分布会有变化。我一般会把max_norm=1.0的先加上,再看 loss 曲线。
  4. 最后才调学习率。而且调学习率不要每次都从头跑,建议用 warmup + cosine 的形式,用一次完整训练的数据来比较。

Loss 震荡的话,最常见的原因是 batch size 太小,或者学习率在后期仍然过大。这时候可以用学习率调度器把后期学习率降下来,而不是手动减小初始学习率。

4.2 显存OOM与梯度爆炸:优化器参数组的排查顺序

显存溢出(OOM)这块,优化器是最大的隐性占用者之一。很多人在排查 OOM 时只看模型大小和中间激活,忘了 Adam 需要保存exp_avg和exp_avg_sq两份状态。以 ViT-Base 为例,参数量约 8600 万,FP32 下一个参数占 4 字节,Adam 的两份状态就是 2 * 8600万 * 4,大约是 688MB。这还没算梯度和 master weight。所以一整层 embedding 的显存代价里,优化器状态往往占了一半以上。

如果 OOM 经常发生在optimizer.step()附近,可以考虑两个方向:一是升级到 Lion、Adafactor 这类省状态的优化器;二是降低 master weight 的精度,比如通过 AMP 把优化器状态保持在 FP32 但梯度用 FP16。Model-Optimizer 里我增加了一个“显存预检”功能:在初始化优化器时,估算各状态张量的大小,并在日志里打印出来,提前预警而不是等训练到一半才崩溃。

4.3 复现不一致:状态检查点、随机种子与优化器状态字典的坑

这个问题在多人协作、多机训练时特别突出。你在一台机器上调好的训练脚本,到另一台机器上跑,理论上结果应该一致,但经常出现细微偏差。排查时我发现最常见的因素,就是优化器状态字典没有正确恢复。

举个例子:训练中断后,你是从上一步的 checkpoint 恢复的,但如果优化器状态没有一起恢复,Adam 的二阶动量相当于从零开始,等效于“学习率突然变大”,收敛曲线自然就变了。如果只是恢复模型权重而优化器没恢复,那结果差异可能不大,但学习率调度器的步数如果也丢了,warmup 就会从头开始,训练前期的 loss 曲线几乎不可能对齐。

所以复现一致性的最低要求是:固定随机种子,初始化模型前固定torch.manual_seed、numpy.random.seed,然后确保 checkpoint 里同时保存模型、优化器、调度器、EMA 四件套。Model-Optimizer 的save_checkpoint就是按这个标准设计的,强烈建议不要为了省空间只存模型权重。

5. 不同任务下的优化器选型表与进阶玩法

最后一块,是很多刚接触训练优化的人最关心的:我到底该用哪个优化器?我基于实际跑过的任务,整理了一张选型表,不敢说绝对正确,但至少能提供一个可靠的起点。

5.1 按任务类型和模型规模选优化器

任务/场景常用优化器学习率参考范围weight_decay参考备注
小规模CNN分类(CIFAR等)SGD + Momentum0.01 - 0.11e-4 - 5e-4调参空间大,收敛稳定
标准Transformer分类/微调AdamW2e-5 - 5e-5(微调时)0.01 - 0.1最稳妥的默认选择
大规模视觉模型LAMB / LARS线性缩放策略0.01 - 0.1适合大批量分布式训练
大规模NLP预训练AdamW / Adafactor1e-4 - 3e-40.01Adafactor更省内存
CV任务尝试新优化器Lion1e-4 - 3e-40.01 - 0.05对warmup敏感需注意
大模型低资源微调Adafactor / 8-bit优化器1e-4 左右0.01主要是为了省显存

这里要特别解释一下为什么大批量训练会用 LAMB 而不是 AdamW。AdamW 在 batch size 很大(比如上万)时,每个 step 的梯度方向更加平滑,导致到达同一个泛化点的“有效步数”变少,模型精度通常会有明显下降。LAMB 通过逐层归一化来扩大学习率,让大批量训练可以走得更快。如果你只是单卡或 2-4 卡训练,就老老实实用 AdamW,没必要追求这些大批量优化器带来的复杂度。

5.2 分层学习率与参数组:同一个模型用不同优化器的操作

优化器选型还有一个高级操作:对模型的不同部分采用不同的优化策略。最典型的场景是迁移学习微调,特征提取层(backbone)已经经过充分预训练,需要小学习率慢慢地适应新数据;而分类头是随机初始化的,需要大学习率快速学习。Model-Optimizer 的layers参数就是直接服务这个需求的。

还有一种更极端的玩法:backbone 用 SGD 或者 LARS 来保持稳定的特征更新,head 用 Adam 来快速自适应。这个做法在一些对抗训练和域适应任务里有效,但不太建议新手上来就试,因为它让训练调试的难度显著上升。我的建议是,先统一用 AdamW 把流程跑通,再根据瓶颈逐步引入分层优化策略。

5.3 更现代的优化器:什么时候值得尝鲜

除了 AdamW、Lion 之外,最近还出现了 Adan、Sophia 这类在特定任务上报出不错结果的优化器。我的态度是:当 Baseline 已经稳定、且你有多余训练资源做对比实验时,才值得试新优化器。尝鲜要控制变量,只改优化器,其他一切保持不动。

如果在尝试新优化器时使用了 Model-Optimizer 这类统一工具,切换成本会大幅降低,否则你会陷入大量的重复代码调整中。我个人的经验是,新优化器带来的收益通常在 0.5% 到 2% 的精度提升范围,但前提是它和你的任务分布匹配。比如 Lion 在视觉任务表现好,但在部分语言模型任务上效果不如 AdamW。不要因为某一篇论文的结论就盲目全量切换。

最后再分享一个小技巧:在换优化器之前,先跑一个固定迭代数(比如 5000 步)的基线实验,把 loss 曲线和验证指标记录好。切换之后,你只需要调整学习率和 warmup,其他参数先保持不变,通过对比曲线来快速判断这个优化器在你的任务上值不值得继续调下去。这样既能控制实验成本,又能避免“感觉新优化器更差”这种没有数据支撑的主观判断。

做模型训练优化这件事,本质上是在不确定性里找规律。Model-Optimizer 并没有改变任何优化算法本身,它只是把复杂的配置和状态管理变得有条理,让我能用更少的时间去验证更多想法。希望这篇文章里的一些设计思路和排错经验,能给你自己的训练流程带来一点启发。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询