PyTorch 训练框架内核深度解析

第10章 优化器与梯度更新

作者 杨艺韬 · 12,312 字 · 发布于 · 更新于

"Note that we default to foreach and pass False to use_fused. This is not a mistake--we want to give the fused impl bake-in time before making it the default, even if it is typically faster."

—— torch/optim/adam.py，adam() 函数内的注释（v2.11.0）

本章要点

  • Optimizer 基类是个简单的"参数容器 + step 方法"：核心数据结构是 param_groups（list of dict）和 state（per-param 状态）。基类不算梯度更新，子类（SGD / Adam / AdamW）实现 step()
  • 三档性能模式：single_tensor（朴素 for 循环）、multi_tensor / foreach（一次操作所有参数，用 torch._foreach_* 批量算子）、fused（一个融合 kernel 算完所有参数 + 状态更新；v2.11 里 CPU / CUDA / MPS 等都有实现）
  • capturable=True 让 optimizer 兼容 CUDA Graph：不依赖 Python 端 step counter，所有状态保留在 GPU 张量里
  • zero_grad(set_to_none=True) 是默认行为：把 param.grad 设为 None 而非 inplace 写零，省一次 zero kernel 调用
  • param_groups 是分组学习率的核心：每个 group 可以有自己的 lr / weight_decay 等。learning rate scheduler 只是在 group 上调 lr 字段
  • state_dict 与 ckpt：optimizer 的状态（如 Adam 的 momentum / variance）也能保存加载，但要留意 param_groups 的 ID 与新模型的 ID 对齐

10.1 一行代码引发的疑问

每次训练循环都长这样：

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch).sum()
    loss.backward()
    optimizer.step()

optimizer.step() 这一行做的事说起来简单 —— "用梯度更新参数"。但仔细想要回答：

  1. optimizer 怎么知道要更新哪些参数？
  2. Adam 的 momentum / variance 状态存在哪？
  3. step 内部怎么并行处理几百万个参数？
  4. 参数张量成百上千，为什么 step 不会成为训练的瓶颈？

答案在 torch/optim/optimizer.py（1193 行）和 torch/optim/adam.py（991 行）等具体优化器实现里。本章拆它的设计。

10.1.1 优化器在训练中的位置

把第 7-9 章的内容串起来，optimizer 是训练循环的最后一环：

graph LR
    F["forward<br/>model(batch) → loss"] --> B["loss.backward<br/>反向图执行<br/>填 param.grad"]
    B --> O["optimizer.step<br/>用 grad 更新 param"]
    O --> Z["optimizer.zero_grad<br/>清梯度"]
    Z --> F
    style O fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

optimizer 拿到的输入是已经填好的 param.grad（第 7 章 §7.6 AccumulateGrad 的产物），它把 grad 转成 param 的更新。optimizer 完全不参与 forward 与 backward —— 它就是个 "param 与 grad → param" 的纯函数（加上 momentum / variance 等状态）。这种"职责清晰"的拆分让 optimizer 可以独立演进、独立测试，与模型无关。

这种契约让 PyTorch 内置的十几种 optimizer（v2.11 的 torch/optim/__init__.py 导出 15 个：SGD/Adam/AdamW/Adagrad/Adadelta/RMSprop/LBFGS/RAdam/NAdam/...）都能用同一个 optim.step() 接口工作。第三方扩展（如 Apex 的 FusedLAMB）也都遵循这个契约 —— 不需要改 PyTorch 主仓就能集成。

10.2 Optimizer 基类：参数与状态的容器

torch/optim/optimizer.py:342 是基类：

class Optimizer:
    def __init__(self, params: ParamsT, defaults: dict[str, Any]) -> None:
        self.defaults = defaults                    # 全局默认参数 (lr, eps...)

        # 6 种 hooks，先于一切初始化
        self._optimizer_step_pre_hooks = OrderedDict()
        self._optimizer_step_post_hooks = OrderedDict()
        self._optimizer_state_dict_pre_hooks = OrderedDict()
        self._optimizer_state_dict_post_hooks = OrderedDict()
        self._optimizer_load_state_dict_pre_hooks = OrderedDict()
        self._optimizer_load_state_dict_post_hooks = OrderedDict()
        ...
        self.state: defaultdict[torch.Tensor, Any] = defaultdict(dict)
        self.param_groups: list[dict[str, Any]] = []

        param_groups = list(params)
        if len(param_groups) == 0:
            raise ValueError("optimizer got an empty parameter list")
        if not isinstance(param_groups[0], dict):
            param_groups = [{"params": param_groups}]    # 用户传 list 时包装成 group

        for param_group in param_groups:
            self.add_param_group(cast(dict, param_group))

核心数据结构两个：

  • param_groups 是 list[dict]。每个 dict 至少有 'params' key（参数列表）+ defaults 里的所有字段（lr / weight_decay 等）。支持每个 group 独立配置（如 backbone 用低 lr，head 用高 lr）
  • state 是 defaultdict[Tensor, dict]。每个参数张量映射到一个状态 dict（Adam 里包含 step / exp_avg / exp_avg_sq）。lazy 初始化：第一次 step 时才创建状态
graph TB
    Opt[Optimizer 实例]
    Opt --> PG["param_groups: list[dict]<br/>──────────<br/>group 0: {'params': [...], 'lr': 0.001, 'weight_decay': 0.01}<br/>group 1: {'params': [...], 'lr': 0.0001}"]
    Opt --> S["state: defaultdict[Tensor, dict]<br/>──────────<br/>p1 → {'step': tensor(100.), 'exp_avg': tensor, 'exp_avg_sq': tensor}<br/>p2 → {'step': tensor(100.), 'exp_avg': tensor, 'exp_avg_sq': tensor}"]
    Opt --> D["defaults: dict<br/>{'lr': 0.001, 'betas': (0.9, 0.999), ...}"]
    Opt --> H["6 种 hooks<br/>(step pre/post, sd pre/post, load_sd pre/post)"]
    style PG fill:#dbeafe,stroke:#3b82f6
    style S fill:#fef3c7,stroke:#f59e0b

10.2.1 为什么 param_groups 是 list[dict] 而不是简单 list

新手以为 optimizer 就是"一组参数 + 一个 lr"。实际生产代码常见写法是分组：

backbone_params = list(model.backbone.parameters())
head_params = list(model.head.parameters())

optimizer = torch.optim.AdamW([
    {'params': backbone_params, 'lr': 1e-5, 'weight_decay': 0.01},
    {'params': head_params,     'lr': 1e-3, 'weight_decay': 0.0},
])

backbone 通常已经预训练，用低 lr 微调；head 是新加的，用高 lr 训练。这种"按 group 分配 hyperparameter"在迁移学习几乎是标配。param_groups 的设计就是为这个场景。

每个 dict 里没指定的字段从 defaults 继承。所以你可以只写差异 —— PyTorch 自动补全。

10.2.2 为什么 state 用 defaultdict[Tensor, dict]

state 是 per-param 状态：每个参数 p 对应一个 state[p] dict。Adam 的 state[p] 包含：

{
    'step': torch.tensor(100.),   # 步数：0-d float32 张量；capturable/fused 时放在参数所在 device，否则在 CPU
    'exp_avg': tensor(shape=p.shape),       # 一阶动量 m
    'exp_avg_sq': tensor(shape=p.shape),    # 二阶动量 v
}

第一次 step 时 PyTorch 给每个参数初始化这些状态（exp_avg / exp_avg_sq 全零）。后续每步更新它们。

为什么用 defaultdict 而非普通 dict？因为 lazy 初始化：访问 state[p] 时如果没有，自动返回空 dict。这避免了在 __init__ 里就为每个参数预分配状态（那时还不知道用户会不会真的训练）。

10.2.3 state 的内存账

state 占的内存惊人。Adam 每个 param 要 exp_avg + exp_avg_sq 两份与 param 同 shape 同 dtype 的张量。对 70B 模型：

  • params: 70B × fp32 = 280 GB（或 fp16 = 140 GB）
  • gradients: 70B × fp32 = 280 GB
  • exp_avg: 70B × fp32 = 280 GB
  • exp_avg_sq: 70B × fp32 = 280 GB

合计 1120 GB！没有任何单卡能装下。这是为什么 70B 训练必须用 FSDP/ZeRO 把这套 state 切到多卡。第 18 章 FSDP 章会展开 ZeRO-1/2/3 三档怎么分别分片"params / grads / optimizer state"。

理解 state 的内存占用，也就明白 SGD 在显存上的优势 —— SGD 的 momentum_buffer 只是参数量的 1 倍，不是 Adam 的 2 倍。在显存极致紧张的场景，"用 SGD 换显存"是一个可考虑的选项（代价是收敛行为与调参经验都要重来）。

10.2.4 add_param_group 的动态扩展

add_param_group(param_group) 让你在 optimizer 创建后动态加新参数组：

optimizer = torch.optim.AdamW(model.backbone.parameters(), lr=1e-5)
# 后来加了新 head, 想给它高 lr
optimizer.add_param_group({'params': model.head.parameters(), 'lr': 1e-3})

这种"事后追加"在 fine-tuning 场景常见 —— 加载预训练 backbone 后再加自己的 head。运行时的 state 以参数张量本身为 key（整数 ID 只在 state_dict() 里出现），新 group 的参数本来就没有 state，所以追加 group 不会破坏已有 state。

新加 group 时所有字段从 defaults 继承（除非用户在 dict 里显式指定）。这种"局部覆盖 + 全局兜底"的 dict 合并模式让用户可以只写差异、不重复指定相同字段，与 React Props 的 spread 操作思想一致。

但要注意：add_param_group 的参数必须不在已有任何 group 里，否则 PyTorch 会报错（避免参数被重复更新）。这种"参数唯一性"检查让 optimizer 不会出现"双更"灾难。

10.2.5 differentiable optimizer

各内置优化器传给 Optimizer.__init__ 的 defaults 里都有 differentiable: False。这个标志让 optimizer 自身的 step 操作可以被 autograd 跟踪 —— 用于 meta-learning 这种"对 optimizer 求二阶导"的场景。

正常训练 differentiable=False，所有 inplace 更新都在 with torch.no_grad(): 里跑（避免反向图爆炸）。Meta-learning 时 differentiable=True，step 内部允许构建反向图，让 outer loop 能对 inner step 求导。这是给研究场景留的扩展点。

10.3 step() 的全流程

step() 是子类实现的，但基类用 _patch_step_function 包装了它，加入 hook 调用：

# torch/optim/optimizer.py:564
def _patch_step_function(self) -> None:
    self._zero_grad_profile_name = (
        f"Optimizer.zero_grad#{self.__class__.__name__}.zero_grad"
    )
    hooked = getattr(self.__class__.step, "hooked", None)
    if not hooked:
        self.__class__.step = self.profile_hook_step(self.__class__.step)
        self.__class__.step.hooked = True

注意它打的补丁是打在类上（self.__class__.step = ...），所以用 hooked 属性做幂等标记，避免同一个 optimizer 类被创建多次时套上一层又一层 wrapper。包装体 profile_hook_step（torch/optim/optimizer.py:511）里第一件事就是开一个名为 Optimizer.step#{类名}.step 的 record_function —— 这正是你在 profiler trace 里看到 Optimizer.step#AdamW.step 的来源，也是量化 optimizer 占比的抓手。

调 optimizer.step() 时实际跑的是包装版：

sequenceDiagram
    autonumber
    participant U as 用户: optimizer.step()
    participant W as profile_hook_step (包装)
    participant Pre as step_pre_hooks
    participant S as 真实 step (子类实现)
    participant Post as step_post_hooks
    U->>W: step()
    W->>Pre: 跑所有 pre-hook
    Pre-->>W: 可能修改 args
    W->>S: 进入子类的 step (Adam.step / SGD.step)
    Note over S: 1. 遍历 param_groups<br/>2. 对每个 param 取出 grad<br/>3. 取出 state, lazy init<br/>4. 调用 _single/_multi/_fused 函数<br/>5. 写入 param 与 state
    S-->>W: 完成
    W->>Post: 跑所有 post-hook
    W-->>U: 返回

10.3.1 step 接受 closure 参数

Optimizer.step(closure=None) 的 closure 参数让某些 optimizer（如 LBFGS）可以多次重新算 loss：

def closure():
    optimizer.zero_grad()
    loss = compute_loss(...)
    loss.backward()
    return loss

optimizer.step(closure)

LBFGS 等"二阶"优化器需要在 line search 中多次评估同一点的 loss / grad，closure 提供了"如何重新算 loss"的回调。普通 SGD/Adam 不需要这个，把 closure 设为 None 即可。

理解 closure 的存在能让你看 PyTorch 源码时不困惑：每个 optimizer 的 step 签名都接受 closure，但 99% 的优化器代码里都是 if closure is not None: ... else: ... 简单分支。这是 PyTorch 给 LBFGS 一类特殊 optimizer 留的扩展点。

10.3.2 6 种 hooks 的工程价值

回顾 §10.2 的 hook 列表（step pre/post、state_dict pre/post、load_state_dict pre/post），这 6 种 hook 各有典型用途（下面除注明源码出处的两处外，都是常见用法举例）：

  • step_pre_hook：step 前做检查或统计；v2.11 里 torch/distributed/_tools/mem_tracker.py 就用全局 step pre/post hook 统计 optimizer step 期间的显存
  • step_post_hook：参数更新后做 EMA（exponential moving average）、log 参数 norm 等；torch/profiler/__init__.py 在设置了 KINETO_USE_DAEMON 时用全局 post-hook 给 Kineto 计步
  • state_dict_pre_hook：保存前做转换（如把 fp32 转 fp16 缩小 ckpt）
  • state_dict_post_hook：保存后注入额外 metadata
  • load_state_dict_pre_hook：加载前做版本迁移、key rename
  • load_state_dict_post_hook：加载后做后处理（如重置 step counter）

理解 hook 的设计意图，你写训练框架时能用 hook 优雅扩展，而不是 monkey-patch optimizer 类。这种"打开扩展点"的工程文化让 PyTorch 生态既稳定又能持续演进。

10.4 SGD：最简优化器

torch/optim/sgd.py 的核心 _single_tensor_sgd：

def _single_tensor_sgd(params, grads, momentum_buffer_list, *,
                       weight_decay, momentum, lr, dampening, nesterov, ...):
    for i, param in enumerate(params):
        grad = grads[i]
        if weight_decay != 0:
            grad = grad.add(param, alpha=weight_decay)

        if momentum != 0:
            buf = momentum_buffer_list[i]
            if buf is None:
                buf = grad.detach().clone()
                momentum_buffer_list[i] = buf
            else:
                buf.mul_(momentum).add_(grad, alpha=1 - dampening)
            if nesterov:
                grad = grad.add(buf, alpha=momentum)
            else:
                grad = buf

        param.add_(grad, alpha=-lr)   # 更新: param = param - lr * grad

——SGD 的全部数学就这一段：

  1. 应用 weight_decay：grad += weight_decay * param
  2. 应用 momentum：buf = momentum * buf + (1-dampening) * grad
  3. Nesterov（可选）：再用 momentum buf 修正
  4. inplace 更新：param -= lr * grad

add_ 等 inplace 操作在 with torch.no_grad(): 上下文里跑，避免参数更新被 autograd 记录。

注意 SGD 的状态只有 momentum_buffer（如果 momentum != 0），所以 SGD 的 state 比 Adam 少一半 —— Adam 的 m + v 是参数量的 2 倍，SGD 的 buf 只是 1 倍。

更深一层：常见的经验说法是 Adam 早期收敛快，SGD with momentum 收敛慢但在一些视觉任务上泛化更好。于是有研究提出先用 Adam、再切 SGD（如 Keskar & Socher 2017 的 SWATS，Improving Generalization Performance by Switching from Adam to SGD）。PyTorch 让这种"切 optimizer"非常容易：

opt_adam = torch.optim.AdamW(params, ...)
# train phase 1 with adam
opt_sgd = torch.optim.SGD(params, ...)
# train phase 2 with sgd

两套 optimizer state 互不影响，因为各自维护自己的 state dict。理解这种"optimizer 独立 state"是深度学习训练流水线设计的基础。

10.5 Adam：状态优化器的代表

torch/optim/adam.py:347 的 _single_tensor_adam（精简版）：

def _single_tensor_adam(params, grads, exp_avgs, exp_avg_sqs, ..., state_steps, *,
                       beta1, beta2, lr, weight_decay, eps, ...):
    for i, param in enumerate(params):
        grad = grads[i]
        exp_avg = exp_avgs[i]      # 一阶动量
        exp_avg_sq = exp_avg_sqs[i] # 二阶动量
        step_t = state_steps[i]

        step_t += 1                 # 步数 +1（源码里这一步在 weight_decay 之前）
        step = _get_value(step_t)   # 非 capturable 路径：取成 Python 数

        if weight_decay != 0:
            if decoupled_weight_decay:      # AdamW 走这条
                param.mul_(1 - lr * weight_decay)
            else:                           # Adam 走这条
                grad = grad.add(param, alpha=weight_decay)

        exp_avg.lerp_(grad, 1 - beta1)              # m = beta1*m + (1-beta1)*g
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1-beta2)   # v = beta2*v + (1-beta2)*g^2

        bias_correction1 = 1 - beta1 ** step
        bias_correction2 = 1 - beta2 ** step

        denom = (exp_avg_sq.sqrt() / sqrt(bias_correction2)).add_(eps)
        param.addcdiv_(exp_avg, denom, value=-lr/bias_correction1)
        # 等价 param -= lr * (m_hat / (sqrt(v_hat) + eps))

——Adam 的核心更新公式实现。注意：

  • 大量 inplace 操作（exp_avg.lerp_、exp_avg_sq.mul_）：避免分配新张量
  • bias correction 修正初期偏估计：开始几步 m / v 接近 0，要除以 1 - beta^step 修正
  • addcdiv_ 是融合算子（a += value * b / c），一次做加 + 除 + scale，比拆开三步快

每步对每个参数都跑一次这套逻辑。对 70B 模型，参数量 7e10，每个元素十来次逐元素运算，单步 Adam 就是 10^11–10^12 量级的运算，而且 single_tensor 版还要为每个参数张量逐个发起这些算子调用。这就是为什么需要后面的优化模式。

10.5.1 与 SGD 的 Adam 字段对比

直观对比：

字段 SGD Adam
step counter 不需要 需要 (用于 bias correction)
momentum buffer momentum_buffer (1x param) exp_avg (1x param)
二阶动量 无 exp_avg_sq (1x param)
显存倍率 1x 2x
超参数 lr, momentum（常调 1-2 个；另有 weight_decay / dampening / nesterov） lr, beta1, beta2, eps, weight_decay (5 个)

SGD 的简洁让它适合"显存极致紧张"或"超参数极致敏感"的场景。Adam 的丰富 state 让它在初始化早期收敛更快、对 lr 选择更宽容。两者各有优劣，今天大模型训练用 AdamW 居多，但 SGD 在某些 vision 任务（如 ImageNet 上的经典 ResNet 训练）仍很常用。

10.5.2 fused 算子的 Apex 渊源

NVIDIA 早期在 Apex 库里提供过 FusedAdam、FusedLAMB 等 CUDA 融合算子，那时 PyTorch 主仓还没有 fused 实现。后来主仓吸收了同样的思路（2022 年引入 CUDA fused Adam 的提交 5f26df0345，标题就是 "APEX style Fused Adam"），做成 torch.optim.Adam(fused=True) 这个内置选项 —— v2.11 里它对应 torch/optim/adam.py:803 的 _fused_adam，底层是 torch._fused_adam_ / torch._fused_adamw_ 这两个 ATen 算子。

这条历史让你理解：今天的 fused optimizer 不是 PyTorch 团队从零写的，是站在 Apex 几年优化经验的肩膀上。Apex 仍然存在并提供更激进的 fused 算子（如 LAMB、FusedSGD with Nesterov），但生产代码越来越多直接用 PyTorch 主仓的内置 fused —— 因为不依赖外部库更省心。

10.5.3 bias_correction 的数学解释

为什么 Adam 需要 bias correction？因为 exp_avg 和 exp_avg_sq 初始为零，前几步的指数滑动平均严重偏低估真实统计量。考虑只跑一步：

m_1 = beta1 * 0 + (1 - beta1) * g_1 = (1 - beta1) * g_1 = 0.1 * g_1   # 假设 beta1 = 0.9

m_1 只有 g_1 的 10%！如果直接用 m_1 做更新，前几步会非常迟缓。bias_correction1 = 1 - beta1^t 在 t=1 时是 0.1，把 m_1 除以这个值就还原成 g_1，消除偏置。

随着 t 增大，beta1^t 与 beta2^t 指数衰减到接近 0，两个 bias_correction 都趋近 1，修正就自动退场了。衰减有多快只取决于 beta：beta1=0.9 时几十步就基本归位，beta2=0.999 时要几千步。所以 bias correction 影响的是训练最开头那一段，而且 beta2 那一路的影响远比 beta1 持久 —— 这正是 warmup 常被认为与 Adam 早期方差估计不稳定有关的原因之一（这是社区里被广泛讨论的假说，不是本章能给出的证明；训练曲线上的其他现象也别一股脑归到 bias correction 头上）。

10.5.4 关于 inplace 操作与 autograd 的协调

注意 Adam 内部的 exp_avg.lerp_(...)、param.addcdiv_(...) 都是 inplace。如果不在关掉梯度的上下文里跑，对 requires_grad=True 的叶子参数做 inplace 更新会被 autograd 直接拒绝，抛 "a leaf Variable that requires grad is being used in an in-place operation."（CPU 实跑确认）。

关梯度的不是 profile_hook_step（它只负责 record_function 和 hooks），而是内置优化器 step 上的另一个装饰器 _use_grad_for_differentiable（torch/optim/optimizer.py:59）：它在调用 step 前执行 torch.set_grad_enabled(self.defaults["differentiable"])，结束后恢复原状态。differentiable=False（默认）时 step 内部就等价于 no_grad，所以函数体里可以放心地用 inplace 操作。

基类不会替子类关梯度。如果你写自定义 optimizer 忘了给 step 加 @torch.no_grad()，第一步就会报上面那个错。这个错误在 PyTorch 错误信息里足够明显，但新手不熟悉时会困惑半天。

10.5.5 AdamW vs Adam：weight decay 的微小差别

v2.11 里 AdamW 直接是 Adam 的子类（torch/optim/adamw.py:20，自 v2.7 起如此）：它只是把 decoupled_weight_decay=True 传给 Adam.__init__，所以 Adam(..., decoupled_weight_decay=True) 与 AdamW 的更新完全一致（CPU 实跑逐元素相等）。另一个容易忽略的差别是默认值：AdamW 的 weight_decay 默认 0.01，Adam 默认 0。两者在数学上的区别是 weight_decay 应用方式：

# Adam: weight_decay 加到 grad 里 (污染 momentum)
grad = grad + weight_decay * param

# AdamW: weight_decay 直接从 param 里减 (不污染 momentum)
param = param - lr * weight_decay * param

数学上不等价 —— Adam 里 L2 项混进梯度后，会和梯度一起被 sqrt(v_hat) 归一化，历史梯度大的参数实际受到的正则反而更弱；AdamW 的衰减绕开这层归一化，对所有参数按同一比例 lr * weight_decay 收缩。论文 Decoupled Weight Decay Regularization（Loshchilov & Hutter，ICLR 2019；arXiv 2017）的实验表明，解耦后 weight decay 与学习率的最优取值更加独立，Adam 的泛化性能也明显改善。注意 PyTorch 的实现把衰减系数乘上了 lr（param.mul_(1 - lr * weight_decay)），论文算法里衰减项只乘调度系数 η_t、不乘步长 α，两边的 weight_decay 数值不能直接照搬。

今天公开的大模型训练配置大多用 AdamW（如 LLaMA 论文写明使用 AdamW）。Adam 在某些 fine-tuning 场景仍有用，但作为默认选择已经被淘汰。

10.6 三档性能模式：foreach / fused / capturable

考虑大模型场景：一个几十层的 transformer，每层有 attention 的 QKVO 四个投影加 MLP 的两三个投影，再算上各类 norm 的 weight/bias —— 参数张量的个数通常是几百到上千。single_tensor 循环里每个张量都要跑 5-6 个算子，dispatcher 调用因此是"张量数 × 算子数"的量级。每次 dispatcher 调用本身很便宜（第 5 章 §5.8 讨论过），但乘上这个系数之后，optimizer.step 就能变成训练 step 里一个肉眼可见的固定开销。

PyTorch 提供三档优化路径：

graph LR
    Single["single_tensor (朴素)<br/>for p in params: update(p)<br/>开销 ∝ 张量数 × 算子数"]
    Foreach["multi_tensor / foreach<br/>torch._foreach_add_(params, grads, ...)<br/>开销 ∝ 算子数 (张量数被吃进一次调用)"]
    Fused["fused (CUDA kernel 融合)<br/>一个 kernel 处理所有 param + state<br/>开销 ∝ 常数"]
    Single -->|每 op 调一次 dispatcher| ForeachSlow[N 次 dispatcher]
    Foreach -->|一次 dispatcher 跑所有张量| ForeachFast[1 次 dispatcher + 内部 loop]
    Fused -->|一次 kernel launch 完成所有数学| FusedFast[1 次 kernel launch]
    style Single fill:#fee2e2,stroke:#ef4444
    style Foreach fill:#fef3c7,stroke:#f59e0b
    style Fused fill:#dcfce7,stroke:#22c55e

10.6.1 foreach (multi_tensor) 模式

torch._foreach_* 是一组特殊算子，一次操作多个张量：

# 朴素: N 次 dispatcher
for p, g in zip(params, grads):
    p.add_(g, alpha=-lr)

# foreach: 一次 dispatcher
torch._foreach_add_(params, grads, alpha=-lr)

_foreach_add_ 在 CUDA 上跑 multi-tensor 融合 kernel，对张量列表做批量操作，dispatcher 开销除以 N；CPU 上它注册的是逐张量循环的 slow path（foreach_tensor_add_list_kernel_slow_），这也是 CPU 默认不走 foreach 的原因。

_multi_tensor_adam（torch/optim/adam.py:554）就是把 _single_tensor_adam 里所有 inplace 操作换成 foreach：

# 简化版
torch._foreach_add_(device_state_steps, 1)     # 所有 step += 1
torch._foreach_lerp_(device_exp_avgs, device_grads, 1 - beta1)   # 所有 m = lerp(m, g)
torch._foreach_mul_(device_exp_avg_sqs, beta2)                    # 所有 v *= beta2
torch._foreach_addcmul_(device_exp_avg_sqs, device_grads, device_grads, value=1-beta2)
...

它把"N 个张量各调一次算子"压成"一次算子调用内部循环 N 个张量"，省下的正是那个 N 倍的框架开销 —— 参数张量越多、每个张量越小，省得越明显。v2.11 里 Adam / AdamW / SGD 等在满足条件时默认就走 foreach（见下面 §10.6.4）。

_foreach_* 算子的 CUDA 实现（aten/src/ATen/native/cuda/ForeachBinaryOpList.cu 等）用模板加宏（如 FOREACH_BINARY_OP_LIST）为常见 binary / unary 操作批量实例化 multi-tensor 版本。每个 _foreach 内部用一个 fused CUDA kernel 一次处理几百个张量，把 dispatcher overhead 与 launch overhead 都摊到一次 launch 上。

注意 foreach 有个隐藏约束：所有张量必须在同一 device、同一 dtype。Optimizer 的 multi_tensor 实现在跑 _foreach_* 之前会按 (device, dtype) 分组，每组单独 foreach。这种分组细节让 mixed-precision 训练（如 bfloat16 weight + fp32 master copy）能正确工作。

10.6.2 fused 模式：把整个 step 压成一个 CUDA kernel

fused 模式更激进 —— 把整个 Adam step（十几个 foreach 操作）压成 一个融合 kernel（每个 (device, dtype) 组还另有一次 _foreach_add_ 给 step 计数加 1）：

# fused 的伪代码 (实际是 C++/CUDA 实现)
torch._fused_adam_(
    params, grads,
    exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
    state_steps,
    lr=lr, beta1=beta1, beta2=beta2, weight_decay=weight_decay, eps=eps,
    amsgrad=amsgrad, maximize=maximize,
)

相对 foreach，fused 进一步省下两样东西：

  • kernel launch 次数：foreach 版 Adam 仍然是十几次 _foreach_* 调用（每个数学步骤一次），fused 版把它们合成一次
  • 中间张量的显存往返：foreach 每一步都要把 exp_avg / exp_avg_sq 从显存读出、算完写回；fused 在一个 kernel 内把整条链算完，中间量留在寄存器里

启用方式：optimizer = torch.optim.Adam(params, fused=True)。但 fused 有限制：

  • 只支持部分优化器（v2.11 里是 Adam / AdamW / SGD / Adagrad）和浮点参数；设备上 CUDA、CPU、MPS 等都有 kernel（native_functions.yaml 里 _fused_adam_ 注册了 CPU / CUDA / MPS）
  • 不支持 differentiable=True
  • 数值上与 foreach 可能有末位差异（运算顺序与融合方式不同）

注意 fused 并不会被自动打开（原因见 §10.6.4），要显式传 fused=True。它在大模型训练里是常规优化，但收益占总 step 时间的比例完全取决于你的 optimizer.step 本来占多少 —— 先 profile，再决定值不值得。

10.6.3 capturable：兼容 CUDA Graph

capturable=True 是另一档优化。它解决一个特殊问题：CUDA Graph 录下的是一串固定的 kernel 与参数，replay 时 CPU 端的 Python 代码根本不会再跑；任何在 CPU 上算好、再作为常量塞进 kernel 的值，都会被固化成录制那一刻的值。

非 capturable 的 Adam 恰好有这样的代码：

step_t += 1                         # step_t 是放在 CPU 上的 0-d 张量
step = _get_value(step_t)           # .item() 取成 Python 数
bias_correction1 = 1 - beta1 ** step    # 在 CPU 端算 bias correction，再作为标量传给 kernel

这种 "CPU 端算标量" 让 step 无法被 CUDA Graph 正确捕获。capturable=True 把 step 改成 GPU 张量 + 用 GPU 算子算 bias correction：

state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)   # GPU 上的 0-d 张量
# step += 1 改成 _foreach_add_(state_steps, 1)
# bias_correction = torch.pow(beta1, step)  在 GPU 上算

代价是少量 GPU 算子开销。收益是 整个 step 能进 CUDA Graph，配合 torch.compile 可以让训练 hot loop 几乎零 CPU 开销。

第 15 章 §15.7.2 讲 CUDA Graph 捕获时，就是把 forward + backward + optimizer.step() 整段录进 graph。

10.6.4 一个看起来微小但重要的细节：_default_to_fused_or_foreach

torch/optim/optimizer.py 里有个 _default_to_fused_or_foreach 函数。它不在 __init__ 里调用 —— __init__ 只把 foreach=None、fused=None 原样存进 defaults；每次 step 时，函数式入口（如 adam()、sgd()）发现两者都是 None，才调它决定这一步走哪条路：

fused, foreach = _default_to_fused_or_foreach(params, differentiable, ...)

它的签名是 _default_to_fused_or_foreach(params, differentiable, use_fused=False)（torch/optim/optimizer.py:164），逻辑是：

  • torch.jit.is_scripting() 或 differentiable=True → 两档都关，走 single_tensor
  • use_fused 为真且所有 param 的 device 在 _get_fused_kernels_supported_devices() 里且是浮点 → fused
  • 否则所有 param 的 device 在 _get_foreach_kernels_supported_devices() 里 → foreach
  • 都不满足 → single_tensor

这里有个关键事实经常被记错：fused 不会被自动打开。 torch/optim/adam.py:939 的调用点写死了 use_fused=False，源码注释解释得很清楚：

# Respect when the user inputs False/True for foreach or fused. We only want to change
# the default when neither have been user-specified. Note that we default to foreach
# and pass False to use_fused. This is not a mistake--we want to give the fused impl
# bake-in time before making it the default, even if it is typically faster.
if fused is None and foreach is None:
    _, foreach = _default_to_fused_or_foreach(params, differentiable, use_fused=False)

torch/optim/sgd.py:283 同理。所以 Adam(params, lr=1e-3) 在 GPU 上默认拿到的是 foreach，不是 fused —— 想要 fused 必须自己传 fused=True。反过来，想强制走 single_tensor（如调试时），传 foreach=False, fused=False。

10.6.5 _foreach 算子的内部实现

_foreach_* 的 CUDA 实现散在 aten/src/ATen/native/cuda/Foreach*.cu 这一组文件里（ForeachBinaryOpList.cu、ForeachBinaryOpScalar.cu、ForeachPointwiseOp.cu、ForeachReduceOp.cu、ForeachTernaryOp.cu、ForeachUnaryOp.cu 等），共用的"一次 launch 处理一批张量"骨架在 aten/src/ATen/native/cuda/MultiTensorApply.cuh。思路可以用这样一段伪代码表达：

// 极简版 _foreach_add_ 内部 (简化)
template <typename scalar_t>
__global__ void multi_tensor_apply_kernel(
    void** params, void** grads, int* sizes, scalar_t alpha) {
    int t = blockIdx.x;   // 第几个张量
    scalar_t* p = (scalar_t*) params[t];
    scalar_t* g = (scalar_t*) grads[t];
    int n = sizes[t];
    int i = threadIdx.x;
    while (i < n) {
        p[i] += alpha * g[i];
        i += blockDim.x;
    }
}

真实的 MultiTensorApply.cuh 比这段伪代码复杂：它把张量地址与尺寸打包进一个有大小上限的 TensorListMetadata 结构（因为 CUDA kernel 参数区容量有限），张量太多时分成多批、每批一次 launch；每个 chunk（不是每个张量）对应一组 block，这样大小悬殊的张量之间也能负载均衡。

理解这层实现，你能预判 foreach 的适用边界：它省的是"每张量一次框架调用"的固定开销，所以张量数量多时收益最大；如果你的模型只有寥寥几个巨大的参数张量，foreach 与 single_tensor 差别就不明显了。

10.6.6 多 device 并发的微妙性

multi_tensor 模式有个隐藏前提：所有张量都在同一 device 才能 fused。如果你的模型一部分在 cuda:0、一部分在 cuda:1（如 model parallel 场景），foreach 内部要按 device 分组分别跑：

分组的粒度实际是 (device, dtype) 两级 —— 混合精度训练里同一个 device 上同时有 bf16 与 fp32 张量，也必须分开。每个 (device, dtype) 组各跑一遍完整的 _foreach_* 序列，所以组数越多，"省掉 N 倍框架开销"这件事的分母就越小。

10.6.7 三档模式的开销结构

不背数字，只看每档的调用次数怎么随规模增长（设参数张量数为 T、Adam 每步的数学步骤数为 K ≈ 十几个）：

模式 算子 / kernel launch 次数 随 T 增长
single_tensor T × K 线性
foreach (multi_tensor) K × (device, dtype) 组数 与 T 无关
fused 约 1–2 × (device, dtype) 组数（融合 kernel + step 计数的 _foreach_add_） 与 T 无关，且常数更小
fused + capturable + CUDA Graph graph replay，CPU 侧接近 0 与 T 无关

三档之间是数量级的差别，但差多少绝对时间完全取决于 T、K、张量大小与你的 GPU —— 换台机器结论就变。判断自己该走哪档的正确姿势只有一个：用 torch.profiler 抓一段训练，看 Optimizer.step#Adam.step 这个 record_function 区间占了整个 step 的多少。如果它只占百分之几，换 fused 的收益上限就是那百分之几。

10.6.8 三档模式的选择决策树

实战里怎么选？决策树：

flowchart TD
    Start[选 optimizer 模式]
    Start --> CPU{CPU 还是 GPU?}
    CPU -->|CPU only| Single[默认 single_tensor<br/>Adam/AdamW/SGD/Adagrad 可显式 fused=True]
    CPU -->|GPU| HasFused{支持 fused?<br/>Adam/AdamW/SGD/Adagrad?}
    HasFused -->|是| Fused[fused=True<br/>最快]
    HasFused -->|否, 如自定义 optimizer| Foreach[foreach 模式]
    Foreach --> CG{要 CUDA Graph?}
    CG -->|是| Capturable[capturable=True]
    CG -->|否| ForeachOK[默认 foreach 已够]
    style Fused fill:#dcfce7,stroke:#22c55e

简化版：

  • 什么都不传 → CUDA 上是 foreach，CPU 不在 foreach 默认设备列表里、回落到 single_tensor，不会自动变 fused（CPU 实跑：_default_to_fused_or_foreach 返回 (False, False)）
  • 想要 fused → 显式 AdamW(params, fused=True)；它与 foreach=True 互斥，同时传会报 "fused and foreach cannot be True together."
  • 要 CUDA Graph / 编译 optimizer → 再加 capturable=True
  • 调试数值问题 → foreach=False, fused=False 退回 single_tensor，最容易逐行对照
  • 自定义优化器（如 LAMB）→ 自己实现 foreach 版本才能享受这层加速

10.6.9 一个真实场景：torch.compile + capturable optimizer

torch.compile 默认编译的是 forward 与（经 AOTAutograd 的）backward，optimizer.step() 不在其中。把 step 也包进 torch.compile 是可行的，配合 capturable=True 能让整个 step 变成编译产物：

optimizer = torch.optim.AdamW(params, fused=True, capturable=True)

@torch.compile
def opt_step():
    optimizer.step()
    optimizer.zero_grad()

# 在训练循环里
for batch in loader:
    loss = model(batch).sum()
    loss.backward()
    opt_step()

这种"编译 optimizer step" 是大模型训练 squeezing performance 的最后一刀。配合 CUDA Graph，整个 forward + backward + step 可以在 GPU 上完全异步执行，CPU 端几乎没有任何 dispatcher 开销。

第 15 章 torch.compile 端到端章会展开 CUDA Graph 这条路径。capturable 的意义不只在性能：torch/optim/optimizer.py 的 _disable_dynamo_if_unsupported 规定，编译时若 capturable=False 而 step 计数在 CUDA 上，就不追踪这段优化器函数、退回 eager 执行 —— 对 fused 这类把 step 计数放在 GPU 上的配置，想让 optimizer.step 真正进入编译器栈，capturable=True 是前提。

10.6.10 fused / capturable 与 dynamic shape

值得注意的边界情况：动态参数集（"某些层在某些 batch 跳过"、运行时增减参数）。在 eager 下 foreach / fused 并不怕它 —— 每次 step 都从 p.grad is not None 的参数现拼张量列表（见 Adam._init_group），某层这一步没有梯度就不进列表，没有需要失效的缓存，也不会切换到 single_tensor。真正怕变化的是"录下来"的路径：CUDA Graph 录的是固定的一串 kernel 与固定地址，参数集一变就得重录；torch.compile 编译的 step 对输入列表有 guard，变化会触发重编译。生产代码建议固定模型架构，少用动态层 —— 这也是 torch.compile / FSDP 等特性的基础假设。

10.7 zero_grad：清梯度的两种语义

torch/optim/optimizer.py:1028：

def zero_grad(self, set_to_none: bool = True) -> None:
    foreach = self.defaults.get("foreach", False) or self.defaults.get("fused", False)
    ...
    with torch.autograd.profiler.record_function(self._zero_grad_profile_name):
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is not None:
                    if set_to_none:
                        p.grad = None
                    else:
                        if p.grad.grad_fn is not None:
                            p.grad.detach_()
                        else:
                            p.grad.requires_grad_(False)
                        if not foreach or p.grad.is_sparse:
                            p.grad.zero_()
                        else:
                            per_device_and_dtype_grads[p.grad.device][p.grad.dtype].append(p.grad)
        if foreach:
            for per_dtype_grads in per_device_and_dtype_grads.values():
                for grads in per_dtype_grads.values():
                    torch._foreach_zero_(grads)

set_to_none=True 是默认行为：直接把 p.grad 设 None。下次 backward 时 AccumulateGrad 看到 grad 是 None，会创建新张量；如果不是 None，会 inplace add。

set_to_none=False 是老行为：跑一次 zero kernel inplace 把 grad 清零。

差异：

维度 set_to_none=True set_to_none=False
zero kernel 调用 0 N (每个 param 一次)
显存 第一次 backward 后才分配 grad grad 一直分配
性能 更快 更慢
与某些第三方代码兼容 可能不兼容（期望 grad 永远不是 None） 兼容

省下的是 N 次（N = 参数张量数）zero kernel，参数越碎省得越多；第 7 章 §7.6.5 已经讨论过，这里不重复。

另外注意 v2.11 的 zero_grad 在 set_to_none=False 这条路上也做了 foreach 优化：如果 optimizer 的 defaults 里 foreach 或 fused 为真，它会先按 (device, dtype) 把所有 grad 收集起来，最后一次 torch._foreach_zero_(grads) 清完，而不是逐个 p.grad.zero_()。

注意一个容易被忽略的语义细节：set_to_none=True 让下一次 backward 第一次给 param.grad 赋值时，PyTorch 会新分配一个张量而不是 inplace 加。也就是说：

  • set_to_none=True：每个 batch 的 grad 是新对象
  • set_to_none=False：所有 batch 共享同一个 grad 对象（inplace clear）

如果你的代码持有 g = param.grad 引用、期望它在下一个 batch 仍指向有效梯度，set_to_none=True 会让你拿到旧值：g 仍是上一轮那个张量，新梯度在另一个新张量里（CPU 实跑：g is param.grad 为 False，g 保持旧值）；set_to_none=False 时 g 与 param.grad 是同一对象，内容随新 backward 更新。这种隐性差异在用 hook 检查梯度的代码里偶尔出现。生产代码里不要 cache param.grad 引用跨 batch 用。

10.8 state_dict 与 ckpt

optimizer 也有 state_dict() / load_state_dict()，跟 nn.Module 类似但简单。一个 Adam optimizer 的 state_dict 大约：

{
    'state': {
        0: {'step': tensor(100.), 'exp_avg': tensor(...), 'exp_avg_sq': tensor(...)},
        1: {'step': tensor(100.), 'exp_avg': tensor(...), 'exp_avg_sq': tensor(...)},
        ...
    },
    'param_groups': [
        {'lr': 0.001, 'betas': (0.9, 0.999), ..., 'params': [0, 1, 2, ...]}
    ],
}

注意 state 的 key 不是 Tensor 对象（不能序列化），而是整数 ID。param_groups 里的 params 列表也是整数 ID。这套 ID 是 Optimizer.state_dict 里那个 pack_group 闭包（torch/optim/optimizer.py:744）现算的：它维护一个 param_mappings: dict[id(p) -> 序号] 和一个跨 group 累加的 start_index，按 group 顺序、group 内参数顺序依次编号。

具体的 ID 编号规则：optimizer 在 save 时按 param_groups 的顺序、每个 group 内 params 的顺序，给每个 param 分配从 0 开始的递增整数。所以 ID 0 是第一个 group 的第一个 param、ID 1 是第二个，依此类推。这种简单递增编号让 ckpt 在 resume 时只要参数顺序不变就能正确匹配。

如果你想跨"不同模型架构"加载 optimizer state（如 fine-tuning 时模型加了新层），就得用名字匹配 + 手动重建 state。HuggingFace 的 save_pretrained / from_pretrained 内部就有这套基于参数名的 state 迁移逻辑。PyTorch 这边只走到一半：如果 optimizer 是用 named_parameters() 构造的，add_param_group 会把名字记进 group 的 param_names（torch/optim/optimizer.py:1132-1144），state_dict 随 group 一并存下；add_param_group 还会校验所有 group "要么都带名字、要么都不带"（:1186）。但 load_state_dict 的文档写明名字不影响加载，真正的 state 匹配仍然是按 ID 顺序 zip 的，state_dict 文档里那句 "match state WITHOUT additional verification"（:705）说得很明白。

加载时 PyTorch 重新匹配新 optimizer 的参数与 ckpt 里的 ID：

new_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
new_optimizer.load_state_dict(saved_state_dict)
# 每个 param 的 state 自动恢复

这个 ID 匹配机制有个微妙坑：新模型的参数顺序必须与保存时一致。如果你重新构造 model 时改了 layer 顺序，ID 对不上，state 会被张冠李戴地装到别的参数上。HuggingFace 等框架用 named_parameters() 的名字匹配规避这个坑；PyTorch 内置 optimizer 即便存了 param_names，匹配仍按 ID 顺序走。

10.8.1 optimizer ckpt 的体积

很多新手发现 torch.save(optimizer.state_dict()) 文件比 torch.save(model.state_dict()) 还大。这不是 bug —— Adam 每个 param 的 state（exp_avg + exp_avg_sq）大小与 param 一致，两个 state 张量加起来是 param 体积的 2 倍。所以 Adam optimizer ckpt ≈ model ckpt × 2。

这就是为什么生产代码里通常只保存 model state_dict 不保存 optimizer，除非要 resume training：

# 仅保存模型 (用于推理 / 部署 / 模型分享)
torch.save(model.state_dict(), 'model.pt')

# 保存完整训练状态 (用于 resume)
torch.save({
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'scheduler': scheduler.state_dict(),
    'epoch': epoch,
}, 'checkpoint.pt')

按 §10.2.3 那笔账算：一个 70B 参数的模型，若权重与 Adam 状态都是 fp32，model 部分约 280 GB、optimizer 部分约 560 GB —— 完整训练 ckpt 落到 TB 量级。这是纯算术推导（实际数字随 dtype、是否存 master weights、是否分片而变），但足以说明为什么大模型训练的 ckpt 存储本身就是个工程课题。

10.9 Learning Rate Scheduler：在 group 上调 lr

torch.optim.lr_scheduler 提供一组 scheduler：StepLR、ExponentialLR、CosineAnnealingLR、OneCycleLR、ReduceLROnPlateau 等。它们的实现都很简单 —— 改 param_group['lr'] 字段：

# 简化示意：真实的 StepLR 继承 LRScheduler，由 get_lr() 计算新 lr，
# last_epoch 从 -1 起步、构造时先做一次初始 step
class StepLR:
    def __init__(self, optimizer, step_size, gamma):
        self.optimizer = optimizer
        self.step_size = step_size
        self.gamma = gamma
        self.last_epoch = 0

    def step(self):
        self.last_epoch += 1
        if self.last_epoch % self.step_size == 0:
            for group in self.optimizer.param_groups:
                group['lr'] *= self.gamma

这套设计让 scheduler 与 optimizer 解耦 —— scheduler 不需要知道 optimizer 内部，只需要修改 param_groups。这种"在共享 dict 上协作"是 Python 库间互操作的优雅模式。

注意调用顺序：

optimizer.step()        # 先用当前 lr 更新参数
scheduler.step()        # 再调整 lr 给下一次用

PyTorch 1.1.0 之前的约定与此相反；自 1.1.0 起改成"先 optimizer.step()、再 scheduler.step()"。顺序写反会跳过 schedule 的第一个 lr 值，LRScheduler.step 检测到"scheduler 先于 optimizer step"时会发 UserWarning："Detected call of lr_scheduler.step() before optimizer.step()..."（CPU 实跑确认）。

10.9.1 chained / sequential schedulers

实际训练经常组合多个 scheduler（如 warmup + cosine decay）：

warmup = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=0.01, total_iters=1000)
cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100000)
scheduler = torch.optim.lr_scheduler.ChainedScheduler([warmup, cosine])

ChainedScheduler 让多个 scheduler 串联，每个 scheduler 各自在 param_groups 上累积修改 lr。这套组合让"先线性 warmup 1000 步、再余弦退火"的复杂 schedule 用一两行代码搞定。

SequentialLR 是另一种组合：根据当前 step 数选不同 scheduler。比如前 1000 步 warmup，之后交给 cosine：

scheduler = torch.optim.lr_scheduler.SequentialLR(
    optimizer,
    schedulers=[warmup, cosine],
    milestones=[1000])

理解这两种组合，你可以构造任意复杂的 lr schedule，不需要自己写 scheduler 子类。

10.9.2 LR Scheduler 与 optimizer 的隐式契约

scheduler 和 optimizer 之间有一个约定俗成的契约：scheduler 假设 optimizer 的 param_groups 里有 'lr' 字段。所有 PyTorch 内置 scheduler 都直接读写 group['lr']。

如果你写自定义 optimizer 没有 'lr' 字段（比如某个论文里的 "scale-invariant" optimizer），用 PyTorch scheduler 就要崩。修法：

class MyOptimizer(Optimizer):
    def __init__(self, params, my_step_size=0.1):
        # 把 my_step_size 别名为 'lr', 让 scheduler 能改它
        defaults = dict(lr=my_step_size)
        super().__init__(params, defaults)

这种"用约定俗成的字段名"是 Python 库间互操作的常见模式 —— 没有显式接口，但大家都知道某些字段名是 contract。

10.9.3 ReduceLROnPlateau：另一种 scheduler

ReduceLROnPlateau 与其他 scheduler 不同 —— 它的 step 接受一个 metric 参数：

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, patience=10)

for epoch in range(100):
    train(...)
    val_loss = validate(...)
    scheduler.step(val_loss)   # 注意传入 val_loss

它会监控 val_loss，连续超过 10 个 epoch 没有改善（即第 11 个 epoch 仍无改善）时把 lr 减半（CPU 实跑确认）。这种"基于性能动态调整"在某些训练场景比固定 schedule 更好用。

接口上 ReduceLROnPlateau.step(metric) 与其他 scheduler.step() 不一致，是 PyTorch 历史遗产之一。如果你写 wrapper 库，要兼容这两种 step 签名。训练监控怎么接入，见第 21 章 §21.33。

10.9.4 一些 optimizer 与训练流水线的关键交互

1. DDP 与 optimizer step：DDP 在 backward 后做 AllReduce 同步梯度，完成后才能调 optimizer.step()。如果你 step 在 AllReduce 完成前调，每个 rank 的更新会基于本地梯度而非全局平均梯度，训练完全失效。DDP 通过 with model.no_sync(): 提供"梯度累积"模式让用户控制何时同步，但默认每 backward 都同步。第 17 章会展开。

2. 混合精度训练：fp16 训练时 grad 容易 underflow 成 0，导致 optimizer.step() 实际不更新参数。GradScaler 先把 loss 放大来防下溢；GradScaler.step(optimizer) 在更新前 unscale 并检查 grad 是否有 inf/nan，有就跳过这次更新，随后 update() 调小 scale。这套机制让 fp16 训练可以"动态调整 loss scale"维持训练稳定。第 20 章 量化与混合精度会展开。

3. ZeRO/FSDP 与 optimizer 分片：FSDP 把参数分片到多 rank，每个 rank 只持有自己的 1/N 参数与对应 1/N optimizer state。step 时每个 rank 只更新自己那份分片，下次 forward 用到时再 AllGather 拼回完整参数。这种分片让 params / grads / optimizer state 这几项的显存降到约 1/N。第 18 章会展开。

理解 optimizer 与这三个系统（DDP / GradScaler / FSDP）的协作，你就理解了大模型训练的工程全景。optimizer 看起来是单点，实际是整个训练流水线的串联点。

10.10 跨专栏关联

  • 《vLLM 推理内核深度解析》专栏第 7 章（模型加载）：vLLM 推理只加载 model state_dict，完全不需要 optimizer state —— 对照 §10.8.1 那笔账就知道，推理侧的 ckpt 天然比训练 ckpt 小一大截
  • 《Tokio 源码深度解析》专栏第 12 章（异步 Mutex / RwLock / Semaphore）：optimizer 与 scheduler 通过共享 param_groups 这个可变 dict 协作，与 Tokio 里多个 Task 共享配置的模式相通 —— 都是"约定字段名而非显式接口"
  • 《MCP 协议设计与实现》专栏第 4 章（生命周期与能力协商）：optimizer 的 per-param state 与 MCP server 的 per-session state 思路类似 —— 都是"按某个 key 维护独立状态"
  • 《Serde 元编程》专栏第 12-13 章（派生宏代码生成）：自定义 optimizer（LAMB、Sophia 等）新增 state 字段后也要能存进 state_dict，这与 Serde 派生的"按字段定义序列化"思想一致

10.11 几条工程经验

实战 optimizer 相关：

1. 想要 fused 必须自己传：v2.11 的默认是 foreach 而非 fused（§10.6.4 有源码注释为证）。GPU 训练想吃这层收益就写 fused=True，但先 profile 确认 Optimizer.step#Adam.step 这段确实占了可观比例

2. foreach 与 fused 互斥：两者同时传 True 会在构造时直接报 RuntimeError（"fused and foreach cannot be True together."）

3. param_groups 用具名 dict 更可读：{'name': 'backbone', 'params': [...], 'lr': 1e-5} 加 name 字段，调试时打印能立刻看出哪个 group

4. resume training 时：load_state_dict(opt_state) 后 lr 会被恢复成 ckpt 时的值，不是你新设的 default。如果想重置 lr，要在 load 后手动改 for g in opt.param_groups: g['lr'] = new_lr

5. 学习率与 batch size 同步缩放：经典的 linear scaling rule（lr 与 batch size 成正比，见 Goyal et al. 2017 Accurate, Large Minibatch SGD）是针对 SGD 提出的；Adam 类优化器不宜照搬线性缩放，放大 batch 后的 lr 需要实验确定

6. AdamW 的 weight_decay 与 grad clipping 顺序：先 clip 再 step，weight_decay 在 step 内自动应用。如果手动写了 param -= weight_decay * param * lr，要小心不要重复 decay

7. CUDA Graph 训练务必 capturable=True：忘了这条会让 graph capture 失败或行为不一致

8. mixed precision 训练用 GradScaler 配合：fp16 训练时梯度容易下溢成 0，需要 torch.amp.GradScaler('cuda') 把 loss 放大、反向后再缩小（老写法 torch.cuda.amp.GradScaler 在 v2.11 已挂 @deprecated，见 torch/cuda/amp/grad_scaler.py:18）。这个 wrapper 内部也跟 optimizer 协作（在 step 前检查梯度是否 inf/nan）

9. 给 optimizer 加 hook 比 wrap 优于继承：要给 optimizer 加自定义行为（如 grad 监控、lr 自适应），用 register_step_pre_hook / register_step_post_hook 而非继承重写。这种 hook 模式让你的扩展不会破坏 internal state 的演进

10. 不要在训练循环中 optimizer = MyOpt(...)：每次新建都会丢掉已积累的 state（momentum、exp_avg 等从零开始），下一次 step 还要重新分配全部 state 张量。optimizer 在训练前创建一次即可

10.12 自定义 optimizer 实战

把整章串起来，写一个 toy 带 momentum 的 SGD optimizer 演示完整套路：

class MySGD(torch.optim.Optimizer):
    def __init__(self, params, lr=0.01, momentum=0.9):
        defaults = dict(lr=lr, momentum=momentum)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group['lr']
            momentum = group['momentum']
            for p in group['params']:
                if p.grad is None:
                    continue
                state = self.state[p]
                if 'momentum_buffer' not in state:
                    state['momentum_buffer'] = torch.zeros_like(p)
                buf = state['momentum_buffer']
                buf.mul_(momentum).add_(p.grad)
                p.add_(buf, alpha=-lr)

        return loss

二十几行就实现了带 momentum 的 SGD（CPU 实跑：与 torch.optim.SGD(lr=0.1, momentum=0.9) 逐元素一致）。所有 PyTorch 内置 optimizer 都是这个套路：

  1. __init__ 定义 defaults，调 super().__init__
  2. step() 在关梯度的上下文里跑（内置优化器用的是 §10.5.4 的 _use_grad_for_differentiable，自定义时用 @torch.no_grad() 即可）
  3. 遍历 param_groups → 遍历 params → lazy 初始化 state → 计算更新

接 closure 参数是为了支持 LBFGS 等需要"算 loss 多次"的优化器。普通 optimizer 忽略它即可。

理解这套模板，你可以快速实现各种学术论文里的新 optimizer（Sophia、Adan、Lion 等都是这个套路实现）。

如果你想让自定义 optimizer 也享受 foreach 加速，要再实现一个 _multi_tensor_my_sgd 函数，把所有 inplace 操作换成 torch._foreach_* 版本，然后在 step 里根据 foreach flag 选择路径。这是 PyTorch 内置 optimizer 的标准实现方式，在 torch/optim/sgd.py / adam.py 等文件里能看到完整模板。

写一个"生产级"optimizer 的代码量远超 toy 版：v2.11 的 torch/optim/adam.py 是 991 行、sgd.py 是 545 行，多出来的部分几乎全在 single_tensor / multi_tensor / fused 三套实现，加上 capturable、differentiable、complex、AMP 的 grad_scale/found_inf、TorchScript 兼容等边界处理。对学术 prototype 来说这太重了，所以很多论文的 reference implementation 只写 single_tensor 版本。如果你的论文 optimizer 被广泛采用，社区会有人贡献 multi_tensor 与 fused 版本。这是 PyTorch 优化器生态的演进路径。

10.13 一些常被忽略的 PyTorch optimizer 实现细节

把整章看完，几个不常被提及但在源码里反复出现的工程细节：

1. _to_scalar(x) 的处理：lr / betas 等超参既可以传 float，也可以传 1 元素 Tensor（配合 capturable）。_to_scalar（torch/optim/optimizer.py:231）把非 0 维的 1 元素 Tensor squeeze 成 0-d Tensor，float 原样返回，统一成后续算子可直接用的形式。这是 PyTorch 在"标量 vs 张量统一接口"上的工程胶水

2. differentiable=False 的 fast path：step 实现里散布着 if differentiable: ... else: ... 分支。non-differentiable 路径走纯 inplace + no_grad；differentiable 路径多做 clone 等处理以保留反向图，更慢但允许 meta-learning

3. grad_scale / found_inf 参数：fused optimizer 会设置 _step_supports_amp_scaling = True，GradScaler 看到后把 grad_scale 与 found_inf 两个张量挂到 optimizer 上（torch/amp/grad_scaler.py），Adam.step 再把它们传进 fused kernel，让 unscale 与 inf 处理也在 kernel 内完成，省掉单独的 unscale 遍历

4. complex 张量支持：has_complex 标志让 optimizer 在复数参数（罕见但支持）场景走专用代码路径。绝大多数代码走默认 real-only 路径

5. JIT script 的兼容性 hack：源码里大量 if torch.jit.is_scripting(): 分支处理 TorchScript 限制（如不能传 Tensor 当 lr）。这些是历史遗产，今天 TorchScript 已经被 torch.compile 替代但兼容代码仍在

理解这些细节，你看 _single_tensor_adam / _multi_tensor_adam 的几百行代码就不会被那么多 if 分支搞晕 —— 它们都对应"为某个特殊场景留的口子"。

10.14 几条 optimizer 设计的"通用启示"

把 optimizer 思想抽象到任何"状态化算子"系统：

第一：defaults + groups 双层配置：让用户既能"一刀切"（defaults）又能"按组定制"（param_groups）。这套模式在数据库连接池、HTTP 客户端、日志器等场景同样适用

第二：lazy 初始化 state：开始不知道用户怎么用，等第一次访问时按需创建。defaultdict 是 Python 实现这种 pattern 的最佳工具

第三：ID 而非对象作为持久化 key：state_dict 的 key 是整数 ID 而非 Tensor 对象，避免序列化复杂对象。这种"用整数代理对象"的模式在 ORM、缓存系统都常见

第四：性能 / 通用性多档：single_tensor → foreach → fused 三档对应"通用慢"到"专门快"。允许用户按场景选择，而不是"一种实现包打天下"

第五：hooks 让框架可扩展：optimizer step 前后开放 hook，让显存追踪、profiler 计步等工具能 hook 进来（§10.3.2）。这是与 nn.Module hook 思想一致的扩展点设计

第六：inplace 数学操作 + no_grad 上下文是底层数值更新的标配：避免分配中间张量、避免污染反向图，是任何"状态化更新"系统的最佳模式。游戏引擎的物理更新、数据库的索引更新都用这套思想

第七：显存不只是模型 — optimizer state 占同样级：开始一个新训练任务时除了算 model 显存，必须算上 optimizer state（Adam 是 2x 参数量）+ gradients（1x）+ activations（depends on batch size）。这个"四件套"显存账是大模型训练的基础数学

第八：closure 接口是给少数高级 optimizer 留的：99% 优化器不用 closure，但接口里始终预留它。这种"为少数边缘场景留扩展点"是稳定 API 设计的常见折衷

10.15 一段反思：optimizer 的"小而美"哲学

最后回顾一下 PyTorch optimizer 设计哲学。它的核心是让"如何更新参数"成为一个可替换的小模块：

  • 接口窄：step() + zero_grad() + state_dict() 几乎是全部公开 API
  • 状态简单：param_groups (list of dict) + state (defaultdict)
  • 与其他系统解耦：optimizer 不知道 model 是什么，只看到 params。model 不知道 optimizer 在做什么，只期待 step 后参数被更新

这套"窄接口 + 简单状态"让 optimizer 是 PyTorch 整个训练栈中最稳定的部分。Module、autograd、dispatcher 都重构过几次，optimizer 的接口从 v0.x 到 v2.11 几乎没变 —— 多年前写的自定义 optimizer 代码，今天大多还能直接跑。

这种"内部反复演进、接口长期稳定"的 trade-off 是 PyTorch 工程文化的体现。理解这条哲学，你看任何"长寿命接口设计"问题时都能更清楚地选择 —— 哪些放出去、哪些藏起来、哪些预留扩展点。

下一章拆 DataLoader：多进程数据流水线、pin_memory、collate_fn 怎么协作。