Rust in Everything
基于 burn 0.22.1

模型 to_device 之后没有梯度

主线位置：训练主线开始之前，把模型放到训练设备上的那一步。Module::to_device 和 fork 逐字段落到 Param 上，区别只在 fork 多做一次 detach() + require_grad()；机制见 Module 派生 的「to_device 与 fork」。之后的 backward 和优化器一步见 backward 与 Gradients、优化器。

flowchart TB
  init["参数已在 CPU 上初始化"] --> move{"换设备用哪个"}
  move -->|to_device| node["参数是 to_device 的输出"]
  node --> s0["第 0 步 梯度 0 个"]
  s0 --> s1["第 1 步 tape consumed"]
  move -->|fork| leaf["参数是新设备上的叶子"]
  leaf --> ok["正常训练"]

实验 crate exp-burn/pb-pit，bin todevice，目标设备 Device::metal(DeviceKind::DefaultDevice).autodiff()，debug，M2 Max、macOS 27、rustc 1.99.0。

现象

一个 Linear(4, 1)，先在 CPU 上建好并读过一次参数（从文件加载、打印、做过一次推理都会让懒参数初始化），再搬到 Metal 上，用 Adam 训练 5 步。每步打印 loss 和 GradientsParams::len()：

fn train(label: &str, mut model: Linear, device: &Device) {
  device.seed(1);
  let x = Tensor::<2>::random([32, 4], Distribution::Uniform(-1.0, 1.0), device);
  let y = x.clone().sum_dim(1);
  let mut optim = AdamConfig::new().init();
  print!("{label:<26} require_grad {:<5} loss(grads):", model.weight.is_require_grad());
  for _ in 0..5 {
    let loss = MseLoss::new().forward(model.forward(x.clone()), y.clone(), Reduction::Mean);
    let value: f32 = loss.clone().into_scalar();
    let grads = GradientsParams::from_grads(loss.backward(), &model);
    print!(" {value:.4}({})", grads.len());
    std::io::stdout().flush().unwrap();
    model = optim.step(1e-1, model, grads);
  }
  println!();
}

fn initialized(device: &Device) -> Linear {
  let model = LinearConfig::new(4, 1).init(device);
  let _ = model.weight.val(); // e.g. loaded from a file, printed, or used once
  let _ = model.bias.as_ref().unwrap().val();
  model
}

七种搬法，每种一个进程（初值没有固定种子，各进程的 loss 不可比，看的是走势和括号里的梯度个数）：

### todevice direct
built on the GPU           require_grad true  loss(grads): 2.0250(2) 1.6701(2) 1.3634(2) 1.1043(2) 0.8904(2)
exit=0
### todevice ad-to_device
flex.autodiff -> to_device require_grad false loss(grads): 1.4316(0)
thread 'main' (3487295) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-autodiff-0.22.1/src/runtime/server.rs:74:17:
Cannot run backward: graph tape has already been consumed at NodeId(3). Burn does not support retain_graph. Combine losses before backward or recompute the forward pass. Use detach() only to intentionally sever gradients through the previous computation.
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit=101
### todevice plain-to_device

thread 'main' (3487339) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-tensor-0.22.1/src/tensor/api/autodiff.rs:75:9:
Tensor::backward requires a tracked autodiff tensor; call Tensor::autodiff().require_grad() on the source leaf before computing the output
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
flex -> to_device          require_grad false loss(grads):exit=101
### todevice lazy-to_device
lazy -> to_device          require_grad true  loss(grads): 1.1184(2) 0.8926(2) 0.7046(2) 0.5396(2) 0.3949(2)
exit=0
### todevice ad-fork
flex.autodiff -> fork      require_grad true  loss(grads): 0.8483(2) 0.6698(2) 0.5224(2) 0.3945(2) 0.2907(2)
exit=0
### todevice plain-fork

thread 'main' (3487496) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-tensor-0.22.1/src/tensor/api/autodiff.rs:75:9:
Tensor::backward requires a tracked autodiff tensor; call Tensor::autodiff().require_grad() on the source leaf before computing the output
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
flex -> fork               require_grad false loss(grads):exit=101
### todevice plain-to_device-train
flex -> to_device.train()  require_grad true  loss(grads): 0.8053(2) 0.5953(2) 0.4269(2) 0.2989(2) 0.2061(2)
exit=0
模型从哪来怎么搬结果
直接建在 metal.autodiff() 上不搬正常，2 个梯度
flex().autodiff() 上已初始化to_device第 0 步 0 个梯度，第 1 步 panic「graph tape has already been consumed」
flex()（无 autodiff）上已初始化to_device第 0 步就 panic「Tensor::backward requires a tracked autodiff tensor」
flex().autodiff() 上、没读过（懒参数）to_device正常
flex().autodiff() 上已初始化fork正常
flex() 上已初始化fork同第三行
flex() 上已初始化to_device 再 .train()正常

最隐蔽的是第二行：第 0 步 loss 照常算出来、backward 不报错，只是 GradientsParams 是空的，优化器什么都没改；要到第二次 backward 才 panic，而那条信息讲的是 retain_graph，和换设备看不出关系。如果只跑一步（比如一个「先试一下」的单步测试），看到的就是「训练没有效果」。

出处

Param<Tensor> 的 to_device 和 fork：

    fn to_device(mut self, device: &Device) -> Self {
        let reparameterization = self.reparameterization.take();
        let base = self.map_to_device(device, |tensor| tensor.to_device(device));
        // ...
    }

    fn fork(mut self, device: &Device) -> Self {
        let reparameterization = self.reparameterization.take();
        let base = self.map_to_device(device, |tensor| {
            let is_require_grad = tensor.is_require_grad();
            let mut tensor = tensor.to_device(device).detach();

            if is_require_grad {
                tensor = tensor.require_grad();
            }

            tensor
        });
        // ...
    }

三种结局对应三种张量：

  • autodiff → autodiff 用 to_device：在 autodiff 下 Tensor::to_device 是一个被记进图的算子，搬过去的参数是它的输出（GradInBackward），不是叶子，is_require_grad() 为假。backward 时梯度穿过它流回 CPU 上的原参数（Module 派生 的动脑筋答案里，原参数取到了形状 [4, 2] 的梯度），GradientsParams::from_grads 按新模型的参数去取，一个也取不到，于是是 0。这次 backward 同时消耗了 to_device 节点的步骤；优化器原样返回模型，下一步的前向再用这两个节点，backward 走到它们时缺步骤，就是 graph tape has already been consumed 那条 panic。循环里每步新建的只有 loss 和中间结果，反复使用的非叶子节点只有这两个参数。
  • 普通设备 → autodiff 用 to_device 或 fork：原参数不是 autodiff 张量，没有要梯度的祖先，搬过去也不要梯度；fork 只在原参数 is_require_grad() 时才补 require_grad()。整个前向没有被追踪的张量，Tensor::backward 开头的断言就失败了。
  • 懒参数：map_to_device 发现参数还没初始化，只改它的目标设备，到第一次读取时直接在新设备上初始化成叶子，两个方法都不走搬运闭包。

fork 先 detach() 把新设备上的张量和旧图断开，再 require_grad() 让它成为新的叶子，这才是训练要的。Learner 自己换设备时用的也是它（Learner::fork）。

修法

  • 训练用的模型换设备一律 fork(&device)：上表第五行。
  • 模型来自普通设备（推理时加载的权重、valid() 过的副本），先搬再 .train()：train() 把参数接回 autodiff 并按 is_active 打开梯度（见 Module 派生 的 train / valid 一节），上表最后一行。
  • 能不搬就不搬：在目标设备上直接建模型（第一行），或者建好之后先别读参数，让懒初始化落到目标设备上（第四行）。

to_device 本身没有错，它是推理时把模型搬到另一个设备的正确写法，也是可微的：需要让梯度穿过设备边界流回原参数时（比如一部分层在 CPU、一部分在 GPU），要的正是这个行为。

怎么确认没中招

  • 训练开始前打印一次 model.weight.is_require_grad()（或遍历所有参数）：搬完设备之后应当是 true。
  • 每步记录 GradientsParams::len()，第 0 步就等于参数个数；是 0 立刻停下。
  • 看到「graph tape has already been consumed」但代码里只有一个 backward，先查模型有没有被 to_device 过。

小结

  • 已初始化的参数 to_device 到 autodiff 设备，得到的是 to_device 算子的输出，不是叶子：第 0 步 0 个梯度，第 1 步「graph tape has already been consumed」。
  • 从普通设备搬过去，参数根本不要梯度，backward 直接断言失败。
  • 训练用 fork；从推理状态回来用 .train()；懒参数不受影响。

💬 关联讨论

发起讨论
加载中...