optimizer自由替换，AdamW/Lion/RMSprop随心切换-编程实验室

optimizer自由替换，AdamW/Lion/RMSprop随心切换

在大模型训练的实践中，一个看似不起眼却深刻影响研发效率的细节浮出水面：换优化器到底有多麻烦？

过去，研究人员为了对比 Adam 和 AdamW 的微调表现，往往需要修改多处代码、调整参数绑定逻辑，甚至重构训练脚本。这种“一次实验，三行改码”的窘境，极大拖慢了技术迭代节奏。而如今，随着 ms-swift 这类全链路框架的成熟，我们终于可以做到——一行配置，切换优化器。

这不只是语法糖，而是一场工程范式的转变：从“写死逻辑”到“声明式配置”，从“重复造轮子”到“即插即用”。本文将围绕这一能力展开，深入剖析 AdamW、Lion 与 RMSprop 的底层机制，并揭示它们如何在 ms-swift 框架中实现无缝切换，为科研探索和工业落地提供真正的灵活性。

为什么我们需要多种优化器？

别忘了，优化器本质上是模型“学习方式”的编码器。不同的任务结构、数据分布、参数空间几何特性，决定了没有一种更新策略能通吃所有场景。

比如：

在 LLaMA 微调中，AdamW 凭借其对权重衰减的解耦控制，显著抑制过拟合；
而在 ViT-L 图像生成任务里，Google 提出的 Lion 以符号更新机制实现了更快收敛与更低显存占用；
至于某些在线学习或嵌入层训练场景，RMSprop 因其对梯度噪声的鲁棒性，反而比自适应更强的算法更稳定。

这就引出了核心问题：如何让同一个训练框架，既能跑通 BERT 的 AdamW 配置，又能轻松迁移到扩散模型的 Lion 实验？

答案就是——抽象接口 + 插件化设计。

AdamW：解耦正则化的胜利

要说现代大模型优化的起点，绕不开 AdamW。

它解决了一个曾被长期忽视的问题：原始 Adam 中的 L2 正则项其实是加在梯度上的，导致其实际作用受到自适应学习率的干扰。换句话说，你设了 weight decay，但它的效果被缩放过，不再是真正意义上的“参数衰减”。

AdamW 的突破在于将正则项从梯度更新中剥离出来，直接施加于参数本身。其更新公式如下：

$$
\begin{aligned}
m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \
v_t &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \
\hat{m}t &= \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t} \
\theta_t &= \theta{t-1} - \alpha \left( \frac{\hat{m}t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta{t-1} \right)
\end{aligned}
$$

注意最后那个 $ \lambda \theta_{t-1} $ ——这才是纯粹的权重衰减，不依赖任何动量或学习率调度。

这个改动看似微小，实则深远。Hugging Face 官方已将其作为 Transformers 库的默认优化器，BERT、T5、RoBERTa 等主流模型均受益于此。

在 ms-swift 中，你可以这样启用：

from torch.optim import AdamW optimizer = AdamW( model.parameters(), lr=5e-5, weight_decay=0.01, betas=(0.9, 0.999), eps=1e-8 )

当然，更常见的是通过 YAML 配置一键指定：

training_args: optimizer: adamw learning_rate: 5e-5 weight_decay: 0.01 beta1: 0.9 beta2: 0.999

框架内部会自动完成类查找与实例化，无需手动导入模块。

Lion：极简主义的进化产物

如果说 AdamW 是“精雕细琢”，那 Lion 就是“大道至简”。

由 Google Research 通过神经架构搜索（NAS）风格的方法演化而来，Lion 放弃了二阶动量（即 $ v_t $），只保留一阶梯度的符号方向进行更新。它的核心流程非常干净：

$$
\begin{aligned}
m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \
\text{update} &= \text{sign}(m_t) \
\theta_t &= \theta_{t-1} - \alpha (\text{update} + \lambda \theta_{t-1})
\end{aligned}
$$

其中sign()函数将动量向量压缩为 ±1，意味着每次更新只关心“往哪走”，不在乎“走多远”。这种设计带来了几个关键优势：

显存节省约 50%：省去 $ v_t $ 缓冲区，在 7B 模型上可减少数 GB 显存占用；
计算更轻量：无开方、除法等高成本操作，GPU 利用率更高；
稀疏激励强：符号函数天然鼓励参数稀疏变化，在视觉-语言任务中泛化更好。

实验表明，在 JFT-300M 上训练 ViT-G 时，Lion 可比 AdamW 提升 1.5% Top-1 准确率，且训练速度提升 10%-15%。

虽然 PyTorch 原生未集成 Lion，但 ms-swift 通过插件机制支持第三方库接入：

from lion_pytorch import Lion optimizer = Lion( params=model.parameters(), lr=1e-4, betas=(0.95,), weight_decay=0.01 )

配合如下配置即可实现无感切换：

training_args: optimizer: lion learning_rate: 1e-4 weight_decay: 0.01 beta1: 0.95

值得一提的是，beta2字段在此无效，但在配置层面仍可保留（会被忽略），保证跨优化器配置的一致性。

RMSprop：老派选手的坚守

尽管被 Adam 系列光芒掩盖多年，RMSprop 并未退出历史舞台。

它最早由 Geoffrey Hinton 在课程中提出，用于应对 RNN 训练中的梯度爆炸问题。其思想朴素却有效：用历史平方梯度的指数平均来动态调节每个参数的学习率。

更新规则如下：

$$
\begin{aligned}
v_t &= \beta v_{t-1} + (1 - \beta) g_t^2 \
\theta_t &= \theta_{t-1} - \frac{\alpha}{\sqrt{v_t} + \epsilon} g_t
\end{aligned}
$$

这里 $ v_t $ 相当于一个“梯度强度计”，在过去变化剧烈的地方自动降低步长，防止震荡发散。

虽然缺少动量机制（可选添加），也不如 AdamW 泛化能力强，但它有几个不可替代的优势：

低内存开销：仅需维护一个滑动平均变量；
抗噪性强：适合小批量、流式或非平稳环境下的训练；
嵌入层友好：避免词向量因大幅更新导致语义漂移。

例如，在图文匹配任务中使用余弦相似度损失时，RMSprop 常表现出更好的稳定性。

在 ms-swift 中启用也极为简单：

import torch optimizer = torch.optim.RMSprop( model.parameters(), lr=1e-3, alpha=0.99, eps=1e-8, weight_decay=1e-4, momentum=0.0 )

对应配置项为：

training_args: optimizer: rmsprop learning_rate: 1e-3 alpha: 0.99 weight_decay: 1e-4

特别适用于资源受限设备（如单卡 A10/T4）上的微调任务，兼顾稳定与效率。

如何实现“自由切换”？架构背后的秘密

这一切的背后，是 ms-swift 对训练组件的高度抽象。

优化器不再是一个硬编码模块，而是Trainer组件中的一个可插拔单元。整体结构如下：

[Model] → [Dataset] → [Trainer] ↗ [Optimizer Plugin] ↘ [Scheduler + Gradient Clipping]

整个流程基于工厂模式构建：

配置解析：读取 YAML/JSON 文件中的optimizer字段；
注册表查找：根据名称映射到对应的构造函数（如"adamw"→torch.optim.AdamW）；
参数注入：提取通用字段（lr, weight_decay 等）并传入初始化函数；
实例化返回：生成符合torch.optim.Optimizer接口的对象供训练循环使用。

更重要的是，这套机制支持运行时动态切换。例如，可以在前 1k 步使用 SGD 进行 warm-up，之后再切换为 AdamW：

def on_step_end(step): if step == 1000: trainer.switch_optimizer("adamw", lr=2e-5)

这种灵活性在课程学习（curriculum learning）或多阶段微调中极具价值。

工程实践中的真实收益

场景一：快速验证新论文结果

当你看到一篇新论文声称“X 优化器在 Y 任务上超越 AdamW 3%”，传统做法是从头搭环境、找实现、调参。而在 ms-swift 中，只需：

安装对应插件（如pip install lion-pytorch）
修改配置文件中的optimizer: lion
启动训练，复现实验

整个过程无需改动一行训练逻辑代码。

场景二：适配不同模型的最佳实践

LLaMA 推荐 AdamW，Stable Diffusion 开始倾向 Lion，而某些语音模型仍在用 RMSprop。如果每个项目都要重写优化器部分，工程成本极高。

ms-swift 的解决方案是：一套模板，多种配置。

无论是 7B 还是 13B 模型，训练脚本完全一致，差异仅体现在配置文件中。这让团队协作、CI/CD 流水线部署变得异常高效。

场景三：显存敏感场景下的轻量化选择

在消费级 GPU 上微调大模型时，显存往往是瓶颈。此时可以选择 Lion 替代 AdamW，减少约 20% 显存占用，同时保持相近甚至更优的收敛曲线。

这对于边缘部署、私有化交付等场景尤为重要。

设计哲学：兼容、透明、开放

ms-swift 在优化器设计上坚持三个原则：

兼容性优先：所有优化器必须遵循torch.optim.Optimizer接口，确保与 DDP、FSDP、ZeRO 等分布式策略无缝协作；
默认值合理：为每种优化器预设推荐超参组合（如 AdamW 默认weight_decay=0.01），降低新手门槛；
日志透明化：训练日志中明确输出所用优化器及其全部超参，便于复现实验；
混合精度就绪：所有内置优化器均通过 AMP（自动混合精度）测试，支持 FP16/BF16 稳定训练。

此外，框架还预留了用户扩展接口：

@OPTIMIZER_REGISTRY.register('my_adam') class MyCustomAdam(torch.optim.Adam): def __init__(self, params, lr=1e-3, custom_param=0.5): super().__init__(params, lr=lr) self.custom_param = custom_param

注册后即可在配置中使用"optimizer": "my_adam"，真正实现社区共建、生态共荣。