Rust in Everything
基于 tenferro-runtime 0.8.0

第二次 backward 后梯度翻倍

主线位置:eager 主线的最后两步,loss.backward(&mut engine) → x.grad()。梯度槽、trace 和函数式 grad 的机制见 eager AD。

flowchart TB
  step["每一步 新的 loss"] --> bw["backward"]
  bw --> slot{"槽里已有梯度"}
  slot -->|没有| put["存入"]
  slot -->|有| add["相加"]
  add --> grow["n 步之和"]
  clear["clear_grad"] --> slot

实验 crate exp-tenferro8/pit2-exp,bin p18_accumulate(参数 clear 是修好的版本),CPU(CpuEngine::serial()),debug,M2 Max、macOS 27、rustc 1.99.0。

现象

一个最小的「训练循环」:w 是变量,x 是常数,每一步重新算 loss = sum(w * x) 再 backward。梯度应当每步都是 x:

let rt = EagerRuntime::new()?;
let mut engine = CpuEngine::serial();
let w = rt.variable_from(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0])?)?;
let x = rt.constant_from(Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0])?)?;
for step in 1..=3 {
  let (h, loss) = rt.with_eager_session(&mut engine, |s| {
    let h = s.mul(&w, &x)?;
    let loss = s.reduce_sum(&h, None)?;
    Ok::<_, tenferro_ad::Error>((h, loss))
  })?;
  loss.backward(&mut engine)?;
  println!("step {step}: w.grad = {:?}  h.grad = {:?}  x.grad = {:?}", g(&w)?, g(&h)?, g(&x)?);
}
step 1: w.grad = Some([10.0, 20.0, 30.0])  h.grad = None  x.grad = None
step 2: w.grad = Some([20.0, 40.0, 60.0])  h.grad = None  x.grad = None
step 3: w.grad = Some([30.0, 60.0, 90.0])  h.grad = None  x.grad = None

每一步的 loss 都是新建的,w.grad() 却是前几步之和。拿它更新参数,第 n 步的步长放大 n 倍。没有报错、没有警告。

同一个程序末尾,几个相邻的误会:

h.backward(): Err(grad requires a scalar output, got shape [3])
  NonScalarGrad { shape: [3] }
rt.grad(loss, h): Err(grad (GraphBuild): invalid argument wrt: grad is inactive for Input(User { id: 16 }); use grad_optional to observe the inactive state)
rt.grad_optional(loss, h): is_some = false
rt.grad(loss, w) = [10.0, 20.0, 30.0], w.grad still Some([30.0, 60.0, 90.0])
  • 中间结果 h 和常数 x 没有梯度槽,grad() 是 None,不是零。
  • h 的形状是 [3],backward 只收标量,报 NonScalarGrad;非标量要用 backward_with(&mut engine, &cotangent) 给出种子。
  • 函数式的 rt.grad(&mut engine, &loss, &h) 对中间结果求导报「grad is inactive」,消息里是 h 的内部键 Input(User { id: 16 }),看不出指的是哪个张量。
  • rt.grad(&mut engine, &loss, &w) 返回一个新张量,不碰槽位:槽里还是三步之和。

出处

EagerTensor::backward 先检查标量,再从全 1 的种子做一次反向,最后由 store_grads 写进槽位。槽里已有值时在 engine 上相加:

            let next = match current.as_ref() {
                Some(existing) => {
                    let existing_read = existing.tensor_read("EagerRuntime::store_grads")?;
                    let incoming_read = TensorRead::from_tensor(incoming);
                    let tensor = session
                        .add_read(existing_read, incoming_read)
                        .map_err(Error::from)?;
                    AdValueRecord::from_tensor(tensor, "EagerRuntime::store_grads")?
                }

这是有意的:backward 的文档示例就是在循环里调两次、断言梯度是 [4.0, 4.0, 4.0](一次是 2)。和 PyTorch 一样,累加让一个大 batch 可以拆成几次 backward 再一起更新。只有 variable_from 建的叶子有槽;中间结果和 constant_from 的常数没有,所以是 None。

函数式的 EagerRuntime::grad 拿到 None(wrt 不是活动的叶子)时,调 traced 一侧的 inactive_wrt_error 拼出这条 Validation 错误,grad_optional 则把同一个状态当成 Ok(None)。0.7.1 报的是 Internal("grad output is inactive ..."),提交 98319e01(「classify an inactive wrt as typed validation, not Internal」)把它改成了用户可以处理的校验错误。

修法

每一步 backward 之前清掉槽:

for step in 1..=3 {
  w.clear_grad()?;
  // ... 同上
}
step 1: w.grad = Some([10.0, 20.0, 30.0])  h.grad = None  x.grad = None
step 2: w.grad = Some([10.0, 20.0, 30.0])  h.grad = None  x.grad = None
step 3: w.grad = Some([10.0, 20.0, 30.0])  h.grad = None  x.grad = None
  • 一个变量用 clear_grad;整个运行时的所有槽用 EagerRuntime::clear_grads,相当于 PyTorch 的 optimizer.zero_grad()。
  • 不想管槽位,就别用 backward:let dw = rt.grad(&mut engine, &loss, &w)? 每次返回新的梯度张量(上面最后一行)。
  • 要中间结果的梯度,用 rt.grad_optional 先看它是不是活动的;需要真的梯度,就把它改成一个叶子再算后半段。

怎么确认没中招

  • 训练循环里每一次 backward 之前,搜一下有没有 clear_grad / clear_grads。
  • 第一步和第二步用同样的数据各跑一次,grad() 应当相等;第二步是第一步的两倍就是忘了清。
  • grad() 返回 None 时先看这个张量是不是 variable_from 建的(tracks_grad() 为真的不一定是叶子:中间结果 h 的 tracks_grad() 也是真)。

小结

  • eager 的 backward 把梯度加进叶子的槽位,不清就逐步累加;第 n 步的 grad() 是 n 步之和。
  • 每步之前 clear_grad()(或运行时的 clear_grads()),或者改用不碰槽位的 rt.grad。
  • 只有 variable_from 的叶子有梯度;非标量要 backward_with,对中间结果用 rt.grad 报 grad is inactive ... use grad_optional。

💬 关联讨论

发起讨论
加载中...