torch.optim 是 PyTorch 的优化器包：不负责求导，只消费 torch.autograd 算好的梯度、按既定算法更新 torch.nn 模块的参数，并配套 param groups、lr scheduler、梯度裁剪与断点续训机制。

torch.optim

职责边界

优化器与 autograd 的分工是理解本篇的前提：

• autograd 负责求导：loss.backward() 沿动态图反传，把 $\partial L/\partial \theta$ 写进每个参数的 .grad 属性
• optimizer 负责更新：step() 读取 .grad，按算法（动量、自适应学习率等）计算增量并写回参数

优化器本身不感知 loss、不感知计算图，它只是一组按算法规则消费梯度的状态机。这一解耦使任何自定义损失都能接入同一套优化器。

训练闭环

for inputs, targets in loader:
    optimizer.zero_grad(set_to_none=True)  # 梯度置 None 而非清零，省一次写内存，也更省显存
    loss = criterion(model(inputs), targets)
    loss.backward()                        # 梯度累积进 param.grad
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 可选：裁剪
    optimizer.step()                       # 按算法更新参数

三个动作各司其职：zero_grad() 清掉上一轮梯度（PyTorch 梯度默认累加，不清空会混入历史）；backward() 属于 autograd；step() 才是优化器的工作。set_to_none=True 自 PyTorch 2.0 起已是默认值，值得显式写出以表明意图。

torch.cuda 与混合精度下的 AMP 训练需将 step() 包在 GradScaler 中，且梯度裁剪要在 scaler.unscale_(optimizer) 之后进行，否则裁的是缩放后的梯度。

优化器族系

优化器	更新规则要点	状态份数	适用场景
SGD	$w\leftarrow w-\eta g$	0	学习率调度配合好的精调阶段；泛化常优于自适应方法
SGD + momentum	累积速度 $v\leftarrow \mu v+g$，按 $v$ 更新	1（速度）	CV 大 batch 训练（ResNet 系标配）
Adam	一阶/二阶矩指数滑动平均 + 偏差修正	2（$m$、$v$）	快速实验、NLP 小模型、稀疏梯度的默认起点
AdamW	Adam + 解耦权重衰减	2（$m$、$v$）	Transformer/LLM 事实标准
RMSprop	只用二阶矩做梯度缩放	1（$v$）	RNN 时代遗产，现多被 Adam 取代

状态与参数同形状、驻留显存——Adam 族因此固定多占 2 倍参数量的显存，是 分布式训练总览 中显存墙估算与 ZeRO-1 分片的首要对象。

Adam 单步更新：

${\hat{m}}_t=\frac{m_t}{1-{\beta }_1^t},{\hat{v}}_t=\frac{v_t}{1-{\beta }_2^t},w_t=w_{t-1}-\eta \frac{{\hat{m}}_t}{\sqrt{{\hat{v}}_t}+ϵ}$

其中 $m_t$、$v_t$ 分别是梯度及其平方的指数滑动平均（默认 ${\beta }_1=0.9$，${\beta }_2=0.999$，$\eta =10^{-3}$）。偏差修正使训练初期估计无偏，这是 Adam 冷启动稳定的关键。

AdamW 与 Adam 的权重衰减差异

Adam 把 weight decay 实现为 L2 正则：将 $\lambda w$ 加进梯度后再经过 $\hat{m}/\sqrt{\hat{v}}$ 的自适应缩放。后果是衰减强度被梯度历史绑架——梯度大的参数衰减被压缩、梯度小的参数衰减被放大，违背了正则化的初衷。

AdamW 将两者解耦：梯度照常走自适应通路，权重衰减直接独立施加：

$w_t=\text{Adam 更新}-\eta \lambda w_{t-1}$

每个参数获得与梯度量级无关的均匀衰减。这对权重本应保持小量级的 Transformer 尤其重要：AdamW + weight decay 0.01–0.1 几乎是预训练的固定搭配，混用 Adam + L2 会带来可测的泛化损失。

param groups 与分组学习率

构造优化器时传入的其实是一组 dict，每个 dict 是一个 param group，拥有独立的超参数：

optim.AdamW([
    {"params": model.backbone.parameters(), "lr": 1e-5},   # 预训练骨干：小步
    {"params": model.head.parameters(), "lr": 1e-3},       # 新增头部：大步
    {"params": model.norm.parameters(), "weight_decay": 0.0},  # norm/bias 不衰减
])

微调、禁用 norm 层权重衰减、冻结/解冻分层训练都靠它实现。运行期改学习率的正确方式是 optimizer.param_groups[i]["lr"] = new_lr——optimizer 本身没有 set_lr 方法；更规范的做法是交给 scheduler。

lr scheduler

scheduler 持有 optimizer 引用，按 step 计数改写 param group 里的 lr：

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
for epoch in range(epochs):
    for batch in loader:
        ...  # optimizer.step()
    scheduler.step()  # 每 epoch 调用；OneCycleLR 则每 batch 调用

Scheduler	曲线	典型用途
StepLR	每 N 步乘 $\gamma$	简单分段衰减
CosineAnnealingLR	余弦退火至 eta_min	长训练的主流选择
OneCycleLR	warmup 升至 max_lr 再退火	快速收敛，自带 warmup，每 batch step
LinearLR + SequentialLR	线性 warmup 拼接其他曲线	Transformer 训练冷启动稳定

warmup 的动机：训练初期 Adam 二阶矩估计噪声大、大 lr 易震荡，先线性升 lr 再进入主曲线可稳定前期。规则：scheduler.step() 必须在 optimizer.step() 之后调用（PyTorch 1.1.0 起强制，顺序颠倒会告警且计数错位）。

梯度裁剪

clip_grad_norm_ 在 backward() 之后、step() 之前调用，按全局 L2 范数缩放梯度：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# 若全局范数超过 1.0，所有梯度等比例缩放到 1.0；否则不动

它不改变梯度方向，只压制模长，是对梯度爆炸的廉价保险，RNN/Transformer 训练的常规操作。另有按元素截断的 clip_grad_val_，较少使用。

断点续训

优化器内部状态（Adam 的 $m$、$v$，SGD 的动量）不在模型权重里，resume 时若只加载 model.state_dict()，动量清零会导致 loss 突跳。标准做法是把三者一起存：

torch.save({
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),      # 动量/二阶矩 + param_groups 超参数
    "scheduler": scheduler.state_dict(),      # 已进行的 step 计数
    "epoch": epoch,
}, "ckpt.pt")

optimizer.load_state_dict() 按 param 顺序索引匹配参数（不按名字），因此加载前不能改变模型参数的注册顺序。这些状态的体积是参数量的 2–3 倍，是 大模型检查点分片策略（ZeRO/FSDP）主要优化对象；保存格式细节见模型权重保存。

torch.optim 在 PyTorch 架构总览中位于 autograd 之上的一层：向下消费 .grad，向上被训练循环与 scheduler 驱动，是连接求导与收敛的最后一环。