optim.step 之后 loss 不变

主线位置：训练主线的 optimizer step。更新后的参数在 burn 里是 step 的返回值，没接住，这一步就白做了。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备，autodiff）

flowchart TB
  h["train_step 拿 mut 引用"] --> e["E0507"]
  e -->|"照建议加 clone"| c["编译通过 无警告"]
  c --> d["新模型被丢掉"]
  d --> s["loss 不变"]
  d --> a["Adam 状态照样前进"]
  e -->|"写回 *model"| ok["正常更新"]

实验代码：PyTorch 侧 step_helper.py，burn 侧 crate pitb-opt 的 bin discard、othergrads、lint 和一段编译不过的 helper（feature std、flex、autodiff、optim，release）。模型、权重、数据同 A1（见一步）。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

PyTorch 里训练一步常被包成一个函数，模型和优化器传进去，opt.step() 原地改参数，什么都不用返回：

def train_step(model, opt):
  loss = nn.functional.cross_entropy(model(x), y)
  opt.zero_grad()
  loss.backward()
  opt.step()
  return loss.item()

model = a1_model()
opt = torch.optim.SGD(model.parameters(), lr=0.1)
print("losses", [train_step(model, opt) for _ in range(3)])

losses [1.1400575637817383, 1.105960726737976, 1.0757620334625244]

搬到 burn 的现象

照着搬，函数拿 &mut Mlp，把 *model 交给 optim.step：

fn train_step(model: &mut Mlp, optim: &mut ModuleOptimizer, device: &Device) -> f32 {
  let (x, y) = a1_data(device);
  let loss = ce(device).forward(model.forward(x), y);
  let v: f32 = loss.clone().into_scalar();
  let grads = GradientsParams::from_grads(loss.backward(), &*model);
  optim.step(0.1, *model, grads);
  v
}

error[E0507]: cannot move out of `*model` which is behind a mutable reference
  --> src/bin/helper.rs:11:19
   |
11 |   optim.step(0.1, *model, grads);
   |                   ^^^^^^ move occurs because `*model` has type `Mlp`, which does not implement the `Copy` trait
   |
help: consider cloning the value if the performance cost is acceptable
   |
11 -   optim.step(0.1, *model, grads);
11 +   optim.step(0.1, model.clone(), grads);
   |

For more information about this error, try `rustc --explain E0507`.

不包函数、直接在 main 里写 optim.step(0.1, model, grads); 再用 model，报的是 E0382，建议同样是 model.clone()（原文见一步的「忘了接返回值」）。照建议改成 optim.step(0.1, model.clone(), grads);，编译通过，cargo build 没有一条警告。跑三步（bin discard，sgd 模式；fixed 是下面「改法」里的写法）：

clone: losses [1.1400574, 1.1400574, 1.1400574]
fixed: losses [1.1400574, 1.1059608, 1.075762]

clone 版三步 loss 逐位相同，模型一次也没更新过。

换成 Adam 更难发现。下面三列：A 先用 clone 版跑 3 步，发现问题后改成正确写法再跑 5 步；B 从头用正确写法跑 5 步；C 和 A 一样先丢 3 步，但改正时顺手换了一个新的 Adam（lr 0.01，ε 1e-8，discard adam）。列出的是改正后 5 步各自更新前的 loss：

step  after 3 discarded  fresh Adam   discarded+new Adam
   1  1.14005744         1.14005744   1.14005744
   2  1.10198641         1.10198629   1.10198629
   3  1.06723177         1.06721926   1.06721926
   4  1.03499877         1.03404260   1.03404260
   5  1.00374413         1.00236607   1.00236607
fresh == discarded+new, bit for bit: true

PyTorch 同样的 Adam 5 步是 [1.14005756, 1.10198629, 1.06721914, 1.03404224, 1.00236571]，和 B 列差在 f32 的几个 ulp 内（这个差别的来历见优化器状态）。A 列却越走越偏，第 5 步差 1.4e-3：被丢掉的 3 步虽然没改模型，Adam 的矩和步数照样记下了。

原因

ModuleOptimizer::step 按值拿走模型，返回更新后的模型：

    /// Update the `module` parameters with the given `gradients`, advancing the optimizer state.
    pub fn step<M: Module>(
        &mut self,
        lr_module: impl Into<ModuleLearningRate>,
        module: M,
        grads: GradientsParams,
    ) -> M {
        // ...

函数上没有 #[must_use]，所以丢掉返回值不触发 unused_must_use 警告；origin/main 上这一行没有变。clone() 让编译器满意，因为 burn 的 clone 只加引用计数（见运算消费张量）：交进去的是一份共享缓冲的副本，step 给副本算出新参数，model 本身没人碰。

Adam 那一列的偏差来自同一个函数的另一半：更新逐个参数发生在 map_float 里，它把这一步的新状态按 ParamId 写回优化器，而 clone() 出来的副本和原模型 ParamId 相同。所以被丢弃的每一步都推进了原模型那几个参数的矩和步数。3 次看到的是同一个梯度 g₀（模型没变），改正后第一步的梯度还是 g₀，偏差修正后 m̂、v̂ 和全新的 Adam 第一步几乎一样，第 2 行只差 1.2e-7；之后历史里多了 4 份 g₀，新梯度的权重比全新 Adam 小，差距一步步拉开。C 列换了新优化器，状态清零，和 B 逐位相同。

PyTorch 里不存在这个问题：opt 构造时拿着 Parameter 对象，step() 原地改它们，模型对象一直是同一个（一步）。

梯度配错了模型：同样静默

另一个不报错的形态是 GradientsParams 和交给 step 的模型对不上，比如重新 init 过一份数值相同的模型（bin othergrads）：

let model = a1_model(&device);
let other = a1_model(&device); // 数值相同，ParamId 是新的
let loss = ce(&device).forward(model.forward(x), y);
let grads = GradientsParams::from_grads(loss.backward(), &other);
println!("GradientsParams::len() = {}", grads.len());
let other = optim.step(0.1, other, grads);

GradientsParams::len() = 0
fc2.bias changed: false

GradientsParams::from_grads 按 other 的 ParamId 去找梯度，一个都找不到；step 对没有梯度的参数原样返回。

改法

在 &mut 的函数里，把返回值写回 *model：

fn train_step(model: &mut Mlp, optim: &mut ModuleOptimizer, lr: f64, device: &Device) -> f32 {
  let (x, y) = a1_data(device);
  let loss = ce(device).forward(model.forward(x), y);
  let v: f32 = loss.clone().into_scalar();
  let grads = GradientsParams::from_grads(loss.backward(), &*model);
  *model = optim.step(lr, model.clone(), grads);
  v
}

这里的 model.clone() 不复制数据，只是为了从 &mut 后面拿出一个值；关键是 *model =。上面 fixed 那一行就是它，三步 loss 和 PyTorch 的差在 1.2e-7 以内（1.1400574 对 1.1400575637817383、1.1059608 对 1.105960726737976、1.075762 对 1.0757620334625244，f32 的 1 ulp 附近）。不包函数时就是 model = optim.step(lr, model, grads);。

已经用 Adam 丢过几步的，改正之后要换一个新的优化器（C 列），否则状态里留着那几步。

怎么检查

• rustc 有一个默认关闭的 lint unused_results，对任何没用上的非 () 返回值报警告。在训练代码的模块或 crate 顶上打开它（bin lint）：

  Copy

  #![warn(unused_results)]
  

  Copy

  warning: unused result of type `Mlp`
    --> src/bin/lint.rs:14:3
     |
  14 |   optim.step(0.1, model.clone(), grads);
     |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  ...
  

  它对所有函数都生效，噪音不小，适合临时打开扫一遍。

• 训练循环里留一条断言：debug_assert_eq!(grads.len(), 要训练的参数张量个数)。A1 是 4；0 说明梯度和模型配错了。

• 前几步打印 loss。同一批数据上 loss 逐位不变，就是更新没落到模型上。

• 搜代码里的 optim.step(，每一处左边都应该有一个赋值。

小结

• burn 的 optim.step(lr, model, grads) 返回更新后的模型；不接住，编译器报 E0382 或 E0507 并建议 clone()，照做后编译通过、没有警告、loss 一动不动。
• 用 Adam 时更隐蔽：被丢掉的步照样推进按 ParamId 存的矩，改正后的训练和全新 Adam 不一样（A1 上第 5 步差 1.4e-3）。改正时换一个新优化器。
• &mut 的函数里写 *model = optim.step(lr, model.clone(), grads);。
• ModuleOptimizer::step 没有 #[must_use]（origin/main 也没有）；临时打开 #![warn(unused_results)] 能扫出来。
• 梯度和模型 ParamId 对不上时 GradientsParams::len() 是 0，step 静默不做事；用它做断言。