每次 backward 都一样慢

主线位置：变换主线的「grad」一步，eager 写法反复调用：同一个深的计算图，每次前向加 backward。TRG 主线整图求导时就是这个形状。

版本：PyTorch 2.15.0（CPU，单线程）、tenferro-rs 0.8.0（CPU，native，CpuEngine::serial()）

flowchart TB
  f["前向记 400 层"] --> b["backward"]
  b --> hit["命中准备好的程序"]
  hit --> ex["执行"]
  ex --> tr["为每个梯度建导数 trace"]
  tr --> drop["拆成数值后丢掉"]
  drop --> f
  t["traced 程序"] --> once["编译准备一次"]
  once --> run["每次只执行"]

实验代码：PyTorch 脚本一个，tenferro 程序 p15_deep_eager（参数是深度）。函数 loss = Σ sin(sin(…sin(x)))，n = 10，算术可以忽略，只看图的深度。每种做 50 次，打印前 5 次和 50 次的中位数，墙钟时间。跑的时候机器上还有别的任务（uptime 负载 2.3–2.8），数字是量级，不是基准。平台 M2 Max、macOS 27、rustc 1.99.0。

现象

PyTorch 的习惯是每步重建图、backward，除了第一次有初始化开销，之后每次都快：

for call in range(50):
  x = xs.clone().requires_grad_()
  h = torch.sin(x)
  for _ in range(399):
    h = torch.sin(h)
  loss = h.sum()
  loss.backward()

torch call 0: forward 1.437 ms, backward 13.731 ms
torch call 1: forward 0.322 ms, backward 0.824 ms
torch call 2: forward 0.320 ms, backward 0.864 ms
...
torch median of 50: forward 0.295 ms, backward 0.787 ms

tenferro 的 eager 写法同样每次重建：

let x = ert.variable_from(Tensor::from_vec_col_major(vec![N], xs.clone())?)?;
let l = ert.with_eager_session(&mut eng, |s| {
  let mut h = s.sin(&x)?;
  for _ in 1..depth {
    h = s.sin(&h)?;
  }
  s.reduce_sum(&h, None)
})?;
l.backward(&mut eng)?;

depth 400, n 10, serial engine
eager call 0: forward 5.507 ms, backward 60.215 ms
eager call 1: forward 0.885 ms, backward 22.540 ms
eager call 2: forward 0.827 ms, backward 22.485 ms
...
eager median of 50: forward 0.802 ms, backward 22.163 ms

第一次之后确实快了（准备好的导数程序进了缓存），但之后每次仍是 22 ms，是 PyTorch 的 28 倍，而且不随调用次数下降。换深度看：

depth 100, n 10, serial engine
eager median of 50: forward 0.210 ms, backward 1.787 ms
...
depth 200, n 10, serial engine
eager median of 50: forward 0.404 ms, backward 5.880 ms

深度 100 → 200 → 400，backward 1.8 → 5.9 → 22.2 ms，每翻一倍约乘 3.3 到 3.8，比线性快得多；前向是线性的。

出处

缓存只省掉了变换、编译和准备。每次 backward 仍要把整条 trace 整理成源程序，执行之后再为每个梯度建一份导数 trace，好让函数式 grad 的结果还能再求导：semantic_eager_vjp_many 里执行完准备好的程序之后调 derivative_trace_from_frozen_program。而 backward 在 backward_from_seed 里拿到梯度就把它拆成纯数值，刚建好的导数 trace 随即丢掉。这些整理工作每一次调用都要做，量出来的总时间随图的大小超线性增长。逐段代码和 TRG 整图上的测量（4 步每次约 7 秒）见 eager 还是 traced 的「图越深，每次反向越贵」和「TRG：整张图的 eager 反向，每一次都付」两节。

这不是「求导变换随深度指数增长」那个问题：那里是值被重复使用时变换本身爆炸（见求导变换的代价），这里每层只用一次上一层的值，变换不爆炸，贵的是每次调用都重复的整理工作。

修法

同一个函数要反复求导时，用 traced：记图、ad.grad、编译、准备各做一次，之后每次只执行。

let xt = TracedTensor::input_concrete_shape(DType::F64, &[N])?;
let mut h = xt.sin()?;
for _ in 1..depth {
  h = h.sin()?;
}
let g = ad.grad(&h.reduce_sum(None)?, &xt)?;
let p = GraphCompiler::new().compile_with_input_specs(&g, &[(&xt, DType::F64, &[N][..])])?;
let pr = prepare(&rt, &eng, &p, &[&xin])?;
run1(&rt, &mut eng, &pr, &[&xin])?;   // 每次只做这一行

build + grad + compile + prepare: 15.549 ms
traced call 0: 1.931 ms
traced call 1: 0.589 ms
...
traced median of 50: 0.591 ms

一次性的 15.5 ms 之后，每次 0.59 ms（前向加反向），和 PyTorch 的前向加反向（约 1.1 ms）同一量级，约为 eager 前向加反向（23 ms）的 1/39。必须用 eager 时（例如要按值分支），把深图切成短段，每段一个 eager 图、段间手工串反向：TRG 20 步这样做每个 β 0.64 s，见 eager 还是 traced 的「分段」一节和分段。

怎么确认没中招

量一下第 2 到第 5 次的 backward：它们和第 1 次之后的时间一样、又远大于前向，就是这个问题。
把深度减半再量：时间降到 1/3 以下，说明是图大小的超线性代价，不是算术。
训练循环、优化循环、对很多 β 求导这类「同一个图反复求导」的场景，默认用 traced。

小结

tenferro eager 的 backward 第一次之后命中缓存，但每次仍要整理 trace、建再丢掉导数 trace，代价随深度超线性：400 层每次 22 ms，PyTorch 0.79 ms。
这和值被重复使用时的指数爆炸是两回事。
反复求导的深图用 traced，编译准备一次，之后每次 0.59 ms；必须 eager 时按段切开。

每次 backward 都一样慢

现象

出处

修法

怎么确认没中招

小结

💬 关联讨论

登录