一步

主线位置：训练主线的最后一步：backward 算出梯度之后，优化器把梯度变成新的参数。本页把 A1 的一整步逐行对照，重点是这一步里每个张量去了哪。

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

flowchart TB
  subgraph pt ["PyTorch"]
    pz["opt.zero_grad"] --> pb["loss.backward"]
    pb -->|"写进"| pg["p.grad"]
    pg --> ps["opt.step 原地改 p"]
  end
  subgraph bn ["burn"]
    bb["loss.backward"] -->|"返回"| bg["Gradients 按节点"]
    bg --> bf["from_grads"]
    bf --> bp["GradientsParams 按 ParamId"]
    bp --> bs["optim.step 拿走模型"]
    bs -->|"交回"| bm["新模型 同一 ParamId"]
  end

实验代码：A1 那一对（train_step.py 和 crate train-step），以及 crate train-one 的 bin moves、noassign、discard 和 one_moves.py。burn 侧 feature 同模板（std、flex、autodiff、optim），release 构建。平台 M2 Max、macOS 27、rustc 1.99.0。

两份程序

模型、权重和输入在怎么读这一卷里已经给过。这里只看训练一步的那几行。PyTorch：

opt = torch.optim.SGD(model.parameters(), lr=0.1)
logits = model(x)
loss = nn.functional.cross_entropy(logits, y)
opt.zero_grad()
loss.backward()
opt.step()

burn：

let ce = CrossEntropyLossConfig::new().init(&device);
let logits = model.forward(x.clone());
let loss = ce.forward(logits.clone(), y.clone());
let grads = loss.backward();
let grads = GradientsParams::from_grads(grads, &model);
let mut optim = SgdConfig::new().init();
model = optim.step(0.1, model, grads);

逐行对过去：

PyTorch	burn	差别
`SGD(model.parameters(), lr=0.1)`	`SgdConfig::new().init()`	burn 的优化器建出来时不拿参数，也不拿学习率
`opt.zero_grad()`	没有	每次 `backward` 返回一个新容器，没有要清的东西（见 backward 与梯度）
`loss.backward()`	`let grads = loss.backward()`	梯度不写到参数上，作为返回值交出来
（`opt` 去 `p.grad` 里找）	`GradientsParams::from_grads(grads, &model)`	把按计算图节点存的梯度换成按参数 id 存
`opt.step()`	`model = optim.step(0.1, model, grads)`	学习率在这里给；模型按值交进去，更新后的模型作为返回值交回来

两份程序输出同样结构的 JSON，compare.py 逐个数比较：

logits                   n= 12  max|diff|=1.49e-08
loss                     n=  1  max|diff|=1.19e-07
grads.fc1.weight         n= 32  max|diff|=3.73e-09
grads.fc1.bias           n=  8  max|diff|=7.45e-09
grads.fc2.weight         n= 24  max|diff|=7.45e-09
grads.fc2.bias           n=  3  max|diff|=1.49e-08
params_after.fc1.weight  n= 32  max|diff|=0.00e+00
params_after.fc1.bias    n=  8  max|diff|=1.86e-09
params_after.fc2.weight  n= 24  max|diff|=1.49e-08
params_after.fc2.bias    n=  3  max|diff|=9.31e-10
loss_after               n=  1  max|diff|=1.19e-07
worst=1.19e-07 tol=1e-06 OK

最大误差 1.19e-7，是 1 附近的一个 f32 ulp，容差 1e-6 通过；不为 0 的原因（求和顺序）见怎么读这一卷。

这张表要求两边的输入逐位相同。输入和偏置在两边都先用 f64 算、最后舍入成 f32 一次。如果 Rust 侧直接写 0.1 * k as f32 - 0.5、step * o as f32，就是在 f32 里算：16 个输入里有 7 个和 Python 差一到三个 ulp（k = 3 时一边是 -0.19999998807907104，一边是 -0.20000000298023224），fc1.bias[5] 差一个 ulp（0.04999999701976776 对 0.05000000074505806）。那样比出来的 params_after.fc1.bias 是 3.73e-9，正是这个初值的差，量的已经不只是计算误差。最大误差碰巧仍是 1.19e-7，被两边计算本身的一个 ulp 盖住了；对照实验时先确认输入逐位相同。

优化器手里没有参数

PyTorch 的优化器在构造时拿到 model.parameters()，手里存着这几个 Parameter 对象；step() 时挨个看它们的 .grad。burn 的 SgdConfig::new().init() 什么参数都没拿，返回的是一个 ModuleOptimizer。它怎么知道要更新谁？靠 GradientsParams 的键。

GradientsParams::from_grads 遍历模型的每个参数，从 backward 返回的 Gradients 里取出它的梯度，按这个参数的 ParamId 登记。ModuleOptimizer::step 再遍历一次模型，每遇到一个参数就按它的 ParamId 去 GradientsParams 里找梯度，找到了才更新。moves 打出这一步的键和值：

let loss = ce.forward(model.forward(x.clone()), y.clone());
let grads = loss.backward();
let grads = GradientsParams::from_grads(grads, &model);
println!("GradientsParams.len()    {}", grads.len());
for (name, id) in [
  ("fc1.weight", model.fc1.weight.id),
  ("fc2.weight", model.fc2.weight.id),
] {
  println!("  {name} id {id}  get::<2> is_some = {}", grads.get::<2>(id).is_some());
}
let id_fc2_bias = model.fc2.bias.as_ref().map(|b| b.id).expect("bias");
println!("  fc2.bias   id {id_fc2_bias}  get::<1> = {:?}", grads.get::<1>(id_fc2_bias).map(vec1));

GradientsParams.len()    4
  fc1.weight id r8un4uuklgu46  get::<2> is_some = true
  fc2.weight id h4994t0up6peq  get::<2> is_some = true
  fc2.bias   id lq72i0r3o8egc  get::<1> = Some([0.113273114, 0.038419843, -0.15169296])

PyTorch 那边同一个梯度：

fc2.bias.grad            [0.11327312141656876, 0.03841985762119293, -0.1516929566860199]

四个参数张量，四个梯度，键是 ParamId。ParamId 是一个随机生成的 u64（ParamId::new），打印成 base32，每次运行都不一样。get::<2> 的 2 是秩：GradientsParams 里存的是擦掉了秩的张量，取的时候要说回来。

学习率也不在优化器里。optim.step 的第一个参数是 impl Into<ModuleLearningRate>，传一个 f64 就是所有参数同一个学习率；每一步都可以给不同的值，调度器就是这么接上的（见优化器状态）。

一步之后，张量去了哪

PyTorch 的 opt.step() 原地改参数：参数还是那个对象，数据还在那块内存。burn 的 optim.step 按值拿走整个模型，在 map_float 里给每个参数算出新张量，装进一个 ParamId 不变的新 Param，组成新模型交回来。两边各在更新前留一份 fc2.bias：

opt = torch.optim.SGD(model.parameters(), lr=0.1)
alias = model.fc2.bias
snap = model.fc2.bias.detach().clone()
ptr = model.fc2.bias.data_ptr()
print("fc2.bias before         ", model.fc2.bias.tolist())
opt.zero_grad()
nn.functional.cross_entropy(model(x), y).backward()
opt.step()
print("fc2.bias after step     ", model.fc2.bias.tolist())
print("alias taken before step ", alias.tolist())
print("clone taken before step ", snap.tolist())
print("same object, same buffer:", alias is model.fc2.bias, ptr == model.fc2.bias.data_ptr())

let before = model.clone();
let b2 = |m: &train_one::Mlp| vec1(m.fc2.bias.as_ref().map(|b| b.val()).expect("bias"));
println!("fc2.bias before          {:?}", b2(&model));
// ... forward、backward、from_grads，同上一节
model = optim.step(0.1, model, grads);
println!("fc2.bias after step      {:?}", b2(&model));
println!("clone taken before step  {:?}", b2(&before));
let same = model.fc2.bias.as_ref().map(|b| b.id) == Some(id_fc2_bias);
println!("fc2.bias ParamId unchanged: {same}");
println!("param still require_grad: {}", model.fc2.weight.val().is_require_grad());

❗IMPORTANT

动脑筋：alias 和 before 在更新之后打印的是旧值还是新值？

PyTorch 的 alias 是 model.fc2.bias 这个对象本身，snap 是 detach().clone()；burn 的 before 是 model.clone()，而 burn 的 clone 只加引用计数、不复制数据（见运算消费张量）。先想好三份各打印什么，答案在本页末尾。

两边都保持了参数的身份：PyTorch 是同一个 Python 对象，burn 是同一个 ParamId，而且新参数仍然要梯度。下一步的 from_grads 和优化器状态都靠这个身份找到它。更新本身在内层后端上算，不进计算图（tensor.inner()，见 burn 卷的优化器）。

拿错了模型：两边都不报错

PyTorch 里一个老问题：优化器在 A 模型上建，训练的却是 B 模型（比如中途 copy.deepcopy 了一份）。opt.step() 只看自己手里那几个参数的 .grad，B 上的梯度它看不见：

other = copy.deepcopy(model)
before = other.fc2.bias.detach().clone()
opt.zero_grad()
nn.functional.cross_entropy(other(x), y).backward()
opt.step()
print("other fc2.bias changed by opt.step():", not torch.equal(before, other.fc2.bias.detach()))
print("model params have grad:", [p.grad is not None for p in model.parameters()])

other fc2.bias changed by opt.step(): False
model params have grad: [False, False, False, False]

burn 里对应的错是把梯度和另一份模型配对。下面的 other 数值和 model 一样，但每个 Param 是新建的，ParamId 是新的；前向和反向用的是 model：

let other = a1_model(&device);
let loss = ce.forward(model.forward(x.clone()), y.clone());
let grads = GradientsParams::from_grads(loss.backward(), &other);
println!("from_grads(.., &other).len() = {}", grads.len());
let other_bias = b2(&other);
let other = optim.step(0.1, other, grads);
println!("other fc2.bias changed by step: {}", b2(&other) != other_bias);

from_grads(.., &other).len() = 0
other fc2.bias changed by step: false

from_grads 按 other 的参数 id 去 Gradients 里找，一个也找不到，得到一个空的 GradientsParams；step 遇到没有梯度的参数原样返回。没有警告，没有 panic。GradientsParams::len() 是现成的检查：训练循环里它应该等于要训练的参数张量个数，这里是 4。model.clone() 出来的副本共享 ParamId，不会出这个问题；会出问题的是重新 init、或者用 Param::from_tensor 重新装参数的模型。

忘了接返回值

从 PyTorch 搬过来，最顺手的写法是 optim.step(...) 不接返回值。noassign：

let model = a1_model(&device);
let mut optim = SgdConfig::new().init();
let loss = ce.forward(model.forward(x.clone()), y.clone());
let grads = GradientsParams::from_grads(loss.backward(), &model);
optim.step(0.1, model, grads);
let loss2 = ce.forward(model.forward(x), y);

error[E0382]: borrow of moved value: `model`
  --> src/bin/noassign.rs:15:26
   |
10 |   let model = a1_model(&device);
   |       ----- move occurs because `model` has type `Mlp`, which does not implement the `Copy` trait
...
14 |   optim.step(0.1, model, grads);
   |                   ----- value moved here
15 |   let loss2 = ce.forward(model.forward(x), y);
   |                          ^^^^^ value borrowed here after move
   |
help: consider cloning the value if the performance cost is acceptable
   |
14 |   optim.step(0.1, model.clone(), grads);
   |                        ++++++++

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

编译器挡住了，但它给的建议是错的方向。照着加上 .clone()（discard），编译通过、没有任何警告，更新后的模型被直接丢掉：

optim.step(0.1, model.clone(), grads);
let loss2: f32 = ce.forward(model.forward(x), y).into_scalar();
println!("loss before {loss1}  loss after {loss2}");

loss before 1.1400574  loss after 1.1400574

ModuleOptimizer::step 没有标 #[must_use]（origin/main 上也没有），所以丢掉返回值不会触发 unused_must_use 警告。正确的写法只有一种：model = optim.step(lr, model, grads);。用 Adam 这类有状态的优化器时还更糟：map_float 照样把这一步的矩存进优化器，模型却还是旧的，状态和参数从此对不上。

动脑筋答案

PyTorch：

fc2.bias before          [-0.0, -0.019999999552965164, -0.03999999910593033]
fc2.bias.grad            [0.11327312141656876, 0.03841985762119293, -0.1516929566860199]
fc2.bias after step      [-0.01132731232792139, -0.023841984570026398, -0.02483070269227028]
alias taken before step  [-0.01132731232792139, -0.023841984570026398, -0.02483070269227028]
clone taken before step  [-0.0, -0.019999999552965164, -0.03999999910593033]
same object, same buffer: True True