一个训练步

主线位置：训练主线从头到尾，并把另外两条主线串起来：forward 是一串算子主线，loss.backward() 是一次求导主线，之后 GradientsParams::from_grads 和 ModuleOptimizer::step 是训练主线自己的部分。本页不用 Learner，手写循环，每一步对应的源码和细讲的页见怎么读这一卷的训练主线表。

sequenceDiagram
  participant L as 训练循环
  participant Mo as Mlp 模块
  participant AD as Autodiff
  participant GP as GradientsParams
  participant Op as ModuleOptimizer
  L->>Mo: forward x
  Mo->>AD: Linear Relu Linear 各记节点
  L->>AD: MseLoss 再 backward
  AD-->>L: Gradients 按节点 id
  L->>GP: from_grads 遍历模块
  GP-->>L: 按 ParamId 的梯度
  L->>Op: step lr model grads
  Op->>Mo: map 每个 Param
  Op-->>L: 新的 Mlp

程序

实验 crate intro-step，依赖和一次 matmul 相同（flex、metal、tracing，默认 feature 开着，所以 optim 和 autodiff 都在）。

use burn::nn::loss::{MseLoss, Reduction};
use burn::nn::{Linear, LinearConfig, Relu};
use burn::optim::{AdamConfig, GradientsParams};
use burn::prelude::*;
use burn::tensor::{DeviceKind, Distribution};

#[derive(Module, Debug)]
struct Mlp {
  l1: Linear,
  act: Relu,
  l2: Linear,
}

impl Mlp {
  fn new(device: &Device) -> Self {
    Self {
      l1: LinearConfig::new(4, 16).init(device),
      act: Relu::new(),
      l2: LinearConfig::new(16, 1).init(device),
    }
  }

  fn forward(&self, x: Tensor<2>) -> Tensor<2> {
    self.l2.forward(self.act.forward(self.l1.forward(x)))
  }
}

fn train(label: &str, device: Device, steps: usize) {
  device.seed(42);
  // Synthetic regression data: y = x · w + 0.5, 64 samples of 4 features.
  let x = Tensor::<2>::random([64, 4], Distribution::Uniform(-1.0, 1.0), &device);
  let w = Tensor::<2>::from_floats([[1.0], [-2.0], [0.5], [3.0]], &device);
  let y = x.clone().matmul(w).add_scalar(0.5);

  let mut model = Mlp::new(&device);
  let mut optim = AdamConfig::new().init();
  println!("{label}: device = {device:?}, params = {}", model.num_params());

  let t0 = Instant::now();
  let mut t_step = Instant::now();
  for step in 0..steps {
    let pred = model.forward(x.clone());
    let loss = MseLoss::new().forward(pred, y.clone(), Reduction::Mean);
    let loss_value: f32 = loss.clone().into_scalar();
    let grads = loss.backward();
    let grads = GradientsParams::from_grads(grads, &model);
    let n_grads = grads.len();
    model = optim.step(1e-2, model, grads);
    let dt = t_step.elapsed();
    t_step = Instant::now();
    if step < 5 || step % 50 == 49 {
      println!("  step {step:3}  loss {loss_value:.6}  grads {n_grads}  {dt:?}");
    }
  }
  device.sync().unwrap();
  println!("  {steps} steps in {:?}", t0.elapsed());
}

fn main() {
  let steps: usize = std::env::args().nth(1).map(|s| s.parse().unwrap()).unwrap_or(200);
  train("flex", Device::flex().autodiff(), steps);
  train("metal", Device::metal(DeviceKind::DefaultDevice).autodiff(), steps);
}

和 0.21 的写法比，Mlp 没有 <B: Backend>，字段是 Linear 而不是 Linear<B>，forward 收发的是 Tensor<2>。同一个 train 函数换一个 Device 就在另一个后端上训练，这正是 #4717 要的效果（见怎么读这一卷）。

❗IMPORTANT

动脑筋：输出里的 params 和 grads 各是几？

Linear 有一个 weight: Param<Tensor<2>>（形状 [d_input, d_output]）和一个 bias: Option<Param<Tensor<1>>>；Relu 没有参数。num_params() 数的是标量个数，GradientsParams::len() 数的是带梯度的参数张量个数。先算再看本页末尾的答案。

输出

release 构建，两个设备各 200 步（缓存已有 Metal 上这些形状的 autotune 结果）：

flex: device = Device<Autodiff { device: Flex(Cpu), checkpointing: Disabled }>, params = 97
  step   0  loss 4.287917  grads 4  7.013708ms
  step   1  loss 4.183877  grads 4  121.083µs
  step   2  loss 4.088831  grads 4  61µs
  step   3  loss 4.000997  grads 4  55.209µs
  step   4  loss 3.919291  grads 4  49.834µs
  step  49  loss 0.238806  grads 4  37.291µs
  step  99  loss 0.018378  grads 4  36.75µs
  step 149  loss 0.009130  grads 4  37.292µs
  step 199  loss 0.005537  grads 4  35.542µs
  200 steps in 14.576125ms
metal: device = Device<Autodiff { device: Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })), checkpointing: Disabled }>, params = 97
  step   0  loss 4.668396  grads 4  451.457625ms
  step   1  loss 4.542524  grads 4  389.202292ms
  step   2  loss 4.419348  grads 4  135.632125ms
  step   3  loss 4.298376  grads 4  8.576292ms
  step   4  loss 4.180090  grads 4  6.09875ms
  step  49  loss 0.101982  grads 4  3.973833ms
  step  99  loss 0.025170  grads 4  4.865458ms
  step 149  loss 0.016397  grads 4  5.996042ms
  step 199  loss 0.011650  grads 4  2.538042ms
  200 steps in 2.202439042s

loss 在两个设备上都从 4 点几降到 0.01 左右。同一个二进制再跑一次，每一步的 loss 逐位相同。计时（M2 Max、macOS 27、rustc 1.99.0）：

Flex 每步约 36 µs，第一步 7 ms。
Metal 前三步合计近 1 s，之后每步 2.5–6 ms。前几步里第一次出现的每串操作都要编译内核（编译产物默认不缓存，见一次 matmul）。这么小的模型，GPU 上每步的时间几乎全是调度、融合和每步一次读回 loss 的开销，所以比 Flex 慢上百倍；模型大了才轮到 GPU 占优。
debug 构建同样 200 步：Flex 约 0.11–0.13 s，Metal 约 3.1 s（空缓存的第一次 22 s，含设备测速和 autotune）。

📝NOTE

没有蠢问题

问：两个设备都先 device.seed(42)，为什么第 0 步的 loss 一个是 4.287917、一个是 4.668396？

答：种子交给的是各后端自己的随机数发生器。Device::seed 转给 Dispatch::seed，再按设备分到后端：Flex::seed 用它建一个 FlexRng 存进全局锁，CubeBackend::seed 调 cubek::random::seed(seed)，GPU 上的随机数由 cubek 的内核生成。算法不同，同一个种子给出不同的 x 和不同的初始权重，所以训练曲线不同；各自在同一个后端上是可重复的。Device::seed 的文档也提醒：种子可能是全局的，不一定只作用于这一个设备。

逐步看这一步

1. 模块：derive 出来的遍历

#[derive(Module)] 为 Mlp 实现 Module trait。和训练主线相关的是两个遍历方法：visit（只读，访问每个参数）和 map（按值消耗模块，返回换过参数的新模块）。num_params 就是一次 visit。派生宏按字段逐个递归：l1、l2 是 Linear，其中的 Param 带一个 ParamId；act: Relu 没有参数。宏展开的样子和 Param 的懒初始化见 Module 派生；LinearConfig::new(4, 16).init(device) 这种「Config 再 init」两段式见 Config。

2. forward：一串算子主线

Linear::forward 查一下输入的最后一维，再调函数式的 linear(input, weight, bias)（它怎么落到后端的 ModuleOps，见 nn 模块）；Relu 是逐元素的激活。模型在 .autodiff() 设备上建，参数是要梯度的叶子，于是每个算子都经过 Autodiff：先登记一个反向步骤，再交给内层后端（Flex 上当场算，Metal 上 Fusion 先记成 IR）。MseLoss::forward(.., Reduction::Mean) 也是几个张量算子，结果是 Tensor<1>。

loss.clone().into_scalar() 是一次读数据：在 Metal 上它会让融合流 drain、提交这一步 forward 攒下的内核、等 GPU 读回一个数。这是每步一次的同步点。

3. backward：一次求导主线

Tensor::backward 进入 burn-autodiff：从 loss 的节点出发收集步骤、按顺序执行每个 op 的反向，梯度累加进 Gradients，按节点 id 存。这一段的机制在 04 自动微分：节点和 require_grad 见 Autodiff 装饰器，matmul 的反向怎么写见一个算子的反向，执行顺序和 Gradients 见 backward 与 Gradients。

4. 从节点 id 到参数 id

优化器认的是参数，不是计算图节点。GradientsParams::from_grads 用一个 ModuleVisitor 走一遍模块，对每个要梯度的浮点参数 grad_remove，再按 ParamId 登记：

    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<D>>) {
        if let Some(filter) = self.filter.as_ref()
            && !filter.contains(&param.id)
        {
            return;
        }

        // Prevents `grad_remove` from panicking if one of the tensor is not autodiff.
        if !param.is_require_grad() {
            return;
        }

        let Some(grad) = param.val().grad_remove(self.grads) else {
            return;
        };

        self.grads_params.register(param.id, grad);
    }

（burn-optim/src/optim/visitor.rs。）from_grads 按值吃掉 Gradients，取过的梯度用 grad_remove 从里面删掉。

5. 优化器：按参数 id 更新

AdamConfig::new().init() 返回的不是 Adam，而是一个 ModuleOptimizer：AdamConfig::init 先 build() 出逐张量的 Adam，再 ModuleOptimizer::from(..) 包一层。ModuleOptimizer::step 对模块调 map，在 map_float 里对每个参数：

按 ParamId 从 GradientsParams 里取出梯度，没有梯度的参数原样返回；
按 ParamId 从自己的 HashMap<ParamId, OptimizationContext> 里取出上一步的状态（Adam 的一阶、二阶矩），第一次没有；
调 Adam::step，得到新张量和新状态，状态存回表里；
Param::from_mapped_value(id, tensor, mapper)：参数 id 不变，值换成新张量。

参数 id 不变，是下一步还能找到同一份 Adam 状态的原因。Optimizer trait（逐张量的 step）和 ModuleOptimizer（类型擦除的状态、参数分组、梯度裁剪）的分工见优化器；Learner 把这个循环连同数据加载、指标、checkpoint 包起来，见 Learner 与数据。

动脑筋答案

flex: device = Device<Autodiff { device: Flex(Cpu), checkpointing: Disabled }>, params = 97
  step   0  loss 4.287917  grads 4  7.013708ms

params = 97：l1 的权重 4×16 = 64、偏置 16；l2 的权重 16×1 = 16、偏置 1；64 + 16 + 16 + 1 = 97。Relu 不贡献参数。
grads 4：带梯度的参数张量是两个权重、两个偏置，一共 4 个。GradientsParams 按参数张量存，不按标量。
每一步都是 4，说明 map 之后的新参数仍然要梯度、仍然用同一个 ParamId：否则下一步 from_grads 会找不到它们。

小结

模型代码没有后端类型参数：struct Mlp { l1: Linear, .. }，换 Device 就换后端，同一个函数在 Flex 和 Metal 上都训练收敛。
一个 step = forward（一串算子主线，每个 op 经过 Autodiff 记节点）→ backward（求导主线，梯度按节点 id）→ GradientsParams::from_grads（按参数 id）→ ModuleOptimizer::step（map 遍历模块，按参数 id 取梯度和 Adam 状态，换上新值、id 不变）。
AdamConfig::init() 返回 ModuleOptimizer，它把逐张量的 Adam 包起来并保存每个参数的状态。
同一个种子在不同后端上给出不同的随机数；每步读一次 loss 在 GPU 上是一次同步。小模型在 Metal 上每步几毫秒、Flex 上几十微秒，前几步在 Metal 上要编译内核。