eager 还是 traced

主线位置：TRG 主线的「分段求导得内能」一步，用 eager 和 traced 各走一遍；先在变换主线的向量函数上量清两种模式各自的开销。

版本：JAX 0.12.0.1（CPU，x64）、PyTorch 2.15.0（CPU）、tenferro-rs 0.8.0（CPU，native）

flowchart TB
  q1{"形状和步数固定"} -->|是| q2{"要反复求值"}
  q1 -->|否| eg["eager"]
  q2 -->|是| tr["traced 准备一次"]
  q2 -->|否| eg
  eg --> q3{"图很深"}
  q3 -->|是| stg["分段 逐段 backward"]
  q3 -->|否| ok["直接 backward"]
  tr --> q4{"图很深"}
  q4 -->|是| stg2["分段 每段一个程序"]
  q4 -->|否| ok2["整图 grad"]

JAX 用户已经熟悉一种「两个模式」：不 jit 的 jax.grad 逐个算子执行，jit 之后整段编译。tenferro 的 eager 和 traced 不是这个关系。traced 对应 jit(grad(f))：先建图、做变换、编译、准备，之后每次只执行。eager 对应 torch.autograd：前向立即执行并记录，backward 对记录做一次语义 VJP。两边的 VJP 是同一个程序变换、同一组规则（设计取舍 第 5 节），所以结果逐位相同；代价落在不同的地方。

实验代码是 tenferro 的 bin eager_vs_traced、eager_scaling_n、eager_depth、trg_eager、trg_eager_staged、trg_staged、trg_traced（release），Python 侧 eager_jax.py、jit_fma_check.py、scaling_n.py、depth_torch.py、trg_torch.py。M2 Max、macOS 27、rustc 1.99.0；同机有别的任务在跑，负载均值 3–8。时间都是稳态的中位数或最小值（各处注明），engine 宽度写在表头：「1 线程」是 CpuEngine::serial() 和 torch.set_num_threads(1)；JAX 用默认的线程设置。

同一个函数：梯度逐位相同

函数还是 f(x) = Σ sin(x)·x²，n = 1000。eager 写法见 backward；traced 写法：

let xt = TracedTensor::input_concrete_shape(DType::F64, &[N])?;
let y = (&xt.sin()? * &(&xt * &xt)?)?.reduce_sum(None)?;
let g = y.grad(&xt)?;
let p = GraphCompiler::new().compile_with_input_specs(&g, &[(&xt, DType::F64, &[N][..])])?;
// prepare / run1：实验代码里包住 Runtime::prepare 和 Runtime::run 的两个小函数
let pr = prepare(&rt, &eng, &p, &[&xin])?;
let tg = run1(&rt, &mut eng, &pr, &[&xin])?;

JAX 一侧比较不 jit 和 jit：

g_eager = jax.grad(f)(x)
g_jit = jax.jit(jax.grad(f))(x)

# tenferro，1 线程、8 线程各比一次，输出相同
eager grad vs traced grad: max diff 0.000e0, bitwise equal: true
# JAX
no jit vs jit: max diff 4.441e-16, bitwise equal 777/1000

tenferro 的两种模式 1000 个梯度逐位相同。JAX 的不 jit 和 jit 反而有 223 个差在末位。差别来自 XLA 的融合：jit 之后编译出来的 HLO 只有一个 multiply_add_fusion 内核，最后那一步 sin(x)·2x + x²·cos(x) 的乘加可以合成一次 FMA，少一次舍入。jit_fma_check.py 用 Python 的 math.fma 重算一遍：

jit sin == no-jit sin: 1000/1000, cos: 1000/1000
no-jit grad == sin*(2x) + x*x*cos: 1000/1000
jit grad    == sin*(2x) + x*x*cos: 777/1000
jit grad    == fma(sin, 2x, x*x*cos): 993/1000

sin、cos 两边相同；不 jit 的梯度就是普通的乘加；jit 的梯度有 993 个等于 FMA 的结果，剩下 7 个说明融合内核里还有别的重排，本页没有再追。

TRG 上也一样。整张 4 步的图（χ = 16，k = [2, 4, 4, 16]），β = 0.30：

# eager，loss.backward
beta=0.30 lnZ/N=0.799095207857105 grad=0.844054073623855 u=-0.844054073623855 bits=3feb027db048a41a | ...
# traced，ad.grad 之后编译、准备、执行
  hex lnZ/N=0x1.992301d045eecp-1 grad=0x1.b027db048a41ap-1

0x1.b027db048a41ap-1 的位模式就是 3feb027db048a41a，两边逐位相同。所以从 traced 换到 eager（或反过来）调试时，数字的变化不会来自模式本身。

时间：eager 的反向有一段固定开销

同一个函数，n 从 10 到 10⁶，稳态中位数，单位 µs。tenferro eager 一栏是前向（记录）加 backward，traced 一栏是准备好的梯度程序跑一次（程序里自带前向）：

拆开看 eager（1 线程）：

       n  eager fwd  eager bwd  eager sum     traced
      10       11.4       58.3       69.7        8.1
    1000       15.2       60.1       75.3       18.8
  100000      447.8      691.0     1138.8      724.3
 1000000     5081.0     7009.5    12090.5     7845.3

• 小数组上 eager 慢在反向的固定开销。 n = 10 和 n = 1000 的 backward 都是约 60 µs，几乎和数据量无关；同一个梯度的 traced 程序整个只要 8 到 19 µs。PyTorch 的反向在 n = 10 是 17.8 µs。前向（记录）也比不记录贵一倍：n = 1000 时常数输入的前向 7.6 µs，变量输入 14.7 µs（eager_vs_traced）。
• 大数组上开销摊薄。 n = 10⁶ 单线程时 eager 的反向 7.0 ms，和整个 traced 梯度程序（7.8 ms）差不多，总时间是 traced 的 1.5 倍、PyTorch 的 1.1 倍。8 线程时 traced 快到 1.1 ms，eager 只到 5.9 ms，这一行 eager 没有像 traced 那样吃到并行。
• JAX 不 jit 的 grad 不是 eager 的对照物。 它从 10 到 1000 都是 1.1 ms，比 tenferro eager 慢 15 倍；jit 之后在小数组上最快。

进程里的第一次：eager 第一次前向 0.22 ms、第一次 backward 0.53 ms（变换、编译、准备导数程序，之后命中缓存）；traced 建图加变换 0.08 ms、编译 0.02 ms、准备 0.10 ms、第一次执行 0.04 ms。这两组都比 JAX 的第一次小：不 jit 的 grad 第一次 116 ms，jit(grad) 第一次 12 ms（eager_jax.py）。

图越深，每次反向越贵

固定开销不是常数。把函数换成 sum(sin(sin(…sin(x))))，n = 10，算术可以忽略，只看图的深度（1 线程，中位数，µs）：

eager 的前向、traced 的梯度程序、PyTorch 的反向都随层数线性增长；eager 的 backward 从 100 层到 400 层涨了 13 倍，400 层时是 traced 梯度程序的 39 倍、PyTorch 反向的 31 倍。这些 backward 全部命中了准备好的导数程序缓存。

缓存只省掉了变换、编译和准备。每次 backward 仍要把整条 trace 整理成源程序（缓存键要用它的语义指纹），执行之后再为每个梯度建一份导数 trace，以便函数式 grad 的结果还能再求导。semantic_eager_vjp_many 在执行完准备好的程序之后：

            let semantic_trace = derivative_trace_from_frozen_program(
                &source,
                derivative_program.frozen_program(),
                derivative_output_index,
                &[(seed_input_index, Arc::clone(&cotangent_tensor))],
                &[&output_trace, wrt_trace, &cotangent_trace],
                None,
                "semantic_eager_vjp",
            )?;

而 backward 拿到梯度就在 backward_from_seed 里把它拆成纯数值，刚建好的 trace 随即丢掉：

                let tensor = match grad.into_value() {
                    Ok(tensor) => tensor,
                    // ...
                };
                cotangents.insert(key, tensor);

用 macOS 的 sample 对 400 层的循环采样 3 秒：backward 下 1708 个样本里，semantic_eager_vjp_many 占 1328 个，其中 derivative_trace_from_frozen_program 686 个；into_value 丢掉 trace 又占 252 个。栈顶最多的是符号维度（RawSymDim）的 clone 和 drop、malloc / free、ValueKey 的相等比较，都是元数据的簿记。tenferro 卷的设计取舍第 5 节在一个浅图的训练循环上看到了同样的分布；这里补上的是它随深度超线性增长。

TRG：整张图的 eager 反向，每一次都付

TRG 的每一步有两次 SVD，SVD 的三个输出又被重复使用。traced 的 ad.grad 在这种图上随步数指数增长（求导变换的代价）。eager 的 backward 是同一个变换，而且按上一节，每次都要重做一部分。整张图的 eager TRG，β = 0.30，1 线程，四次里后三次的最小值，单位 ms；PyTorch 同样 1 线程：

4 步那一行逐次打印缓存：

  rep 0: backward 7502.34 ms, cache CacheStats { entries: 1, retained_bytes: 28827, hits: 0, misses: 1, evictions: 0, clears: 0 }
  rep 1: backward 7487.65 ms, cache CacheStats { entries: 1, retained_bytes: 28827, hits: 1, misses: 1, evictions: 0, clears: 0 }
  rep 2: backward 7456.97 ms, cache CacheStats { entries: 1, retained_bytes: 28827, hits: 2, misses: 1, evictions: 0, clears: 0 }
  rep 3: backward 7480.52 ms, cache CacheStats { entries: 1, retained_bytes: 28827, hits: 3, misses: 1, evictions: 0, clears: 0 }

命中缓存的第 2 到 4 次和第一次一样慢。在第三次 backward 期间采样 3 秒，1977 个样本有 1894 个落在 backward_from_seed → into_value → 丢掉导数 trace 的元数据 → ValueKey 的相等比较上。traced 在同一张 4 步的图上把这笔钱付一次：ad.grad 2097 ms，梯度程序（863 个运算）编译、准备 9.3 ms，之后每次执行 1.2 ms。PyTorch 每步的反向只和运算数成正比，4 步 1.2 ms。

所以整图的 eager 在 TRG 上比整图的 traced 更不可用：traced 至少能「慢一次、快很多次」，eager 每个 β 都要从头付。

分段：每步一个 eager 图

办法和 traced 一样：不让 backward 看见整条链。每一步的输入做成一个新的叶子（用上一步的数值 variable_from），这一步自己是一张小图；反向从最后一步往回走，每一步用一个标量替身 L_n = ⟨T_{n+1}, ct_{n+1}⟩ + c_n·ln‖T_n‖ 调一次 backward，从叶子的槽里取出 ct_n，交给前一步。这就是 traced 分段（分段）的同一个替身，只是换成 eager 写：

// forward: stage 0 = initial tensor, stage i+1 = RG step i on a fresh leaf
let t_init = rt.with_eager_session(eng, |s| initial(s, &beta, &a, &bb))?;
for (i, &k) in ks.iter().enumerate() {
  let leaf = rt.variable_from(cur.to_tensor(eng)?)?;
  let (tn, ln) = rt.with_eager_session(eng, |s| step(s, &leaf, d, k))?;
  // ...
}
// reverse: one backward per stage, newest first
lt.backward(eng)?;
let mut ct = grad_tensor(rt, eng, &last)?;
for i in (0..n).rev() {
  let (tn, ln) = &outs[i];
  let c = 2f64.powi((n - i) as i32) * scale;
  let surrogate = rt.with_eager_session(eng, |s| {
    let p = s.mul(tn, &ct)?;
    let p = s.reduce_sum(&p, None)?;
    let l = s.scale_real(ln, c)?;
    s.add(&p, &l)
  })?;
  surrogate.backward(eng)?;
  ct = grad_tensor(rt, eng, &leaves[i])?;
}

4 步时它和整图的 eager 逐位相同（grad 的位模式都是 3feb027db048a41a），反向从 6981 ms 降到 4.29 ms。20 步，χ = 16，1 线程，四次里后三次的最小值：

tenferro 0.8.0 eager staged chi=16 steps=20 width=1 ks=[2, 4, 4, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16]
first call: forward 377.40 ms backward 331.61 ms
beta=0.30 lnZ/N=0.790559063800966 u=-0.704498679529753 err=+3.913e-7 hex grad=3fe68b40d094f9a0 | min forward 348.91 ms backward 304.52 ms
beta=0.40 lnZ/N=0.879363060306853 u=-1.106045407611192 err=+3.380e-5 hex grad=3ff1b25cab594c06 | min forward 355.53 ms backward 303.75 ms
beta=0.44 lnZ/N=0.928723962577072 u=-1.402785274913334 err=-5.583e-4 hex grad=3ff671cef8f1029d | min forward 359.20 ms backward 298.20 ms
beta=0.50 lnZ/N=1.025792858455948 u=-1.745583468150079 err=-1.889e-5 hex grad=3ffbede8ee42491c | min forward 342.79 ms backward 295.31 ms
beta=0.60 lnZ/N=1.210133028164417 u=-1.909086939792358 err=-7.621e-7 hex grad=3ffe8b9ebf3a0ff9 | min forward 345.57 ms backward 306.24 ms
cache: CacheStats { entries: 7, retained_bytes: 42848, hits: 433, misses: 7, evictions: 0, clears: 0 }

• err 是内能 u 和 Onsager 精确解的差，五个 β 都和 traced 分段的结果到所列位数相同（β = 0.44 靠近临界点，误差来自 χ = 16 的截断）。和 traced 分段的梯度比，β = 0.30 差 10 个 ulp（1.1e-15），两个程序的结构不同，舍入顺序不同。
• 缓存只有 7 条：形状 (d, k) 有 5 种步，加上初值和末尾的 log trace，20 步、5 个 β 全部复用。
• 同一个 β = 0.30，紧接着各跑一次：eager 分段前向 339.6 ms 加反向 298.0 ms，共 638 ms；traced 分段的「求值加梯度」945 ms。eager 快，因为它的反向直接用前向留下的残差，而 traced 每步的反向程序（204 个运算，前向程序 40 个）要把这一步的前向、包括两次 SVD 重算一遍。
• PyTorch 一张图直接反向，20 步前向 204.8 ms、反向 77.0 ms，但梯度是 nan（简并奇异值，见 简并 SVD）；tenferro 默认的 SVD 规则给出上面这些有限、正确的数。

形状变了：缓存按形状分

eager 的准备好的导数程序按「engine 目标、语义指纹、具体的输入形状」缓存（PreparedDerivativeCacheKey）。和 jit 一样，换一个形状就要重新准备，换回来又命中。traced 的 Prepared 则绑定在一个形状上，换形状的处理见 记图与 jit。

按值分支：eager 能读数

eager 的每个结果在前向时就有数值，回调里可以直接读出来、用 Rust 的 if 分支：

let l = ert.with_eager_session(&mut eng, |s| {
  let m = s.reduce_max(&xv, None)?;
  let big = m.value()?.as_slice::<f64>()?[0] > 1.0;
  if big {
    let sq = s.mul(&xv, &xv)?;
    s.reduce_sum(&sq, None)
  } else {
    let sn = s.sin(&xv)?;
    s.reduce_sum(&sn, None)
  }
})?;
l.backward(&mut eng)?;

eager: max(x) = 2 read inside the session, branch big=true
eager: grad[:3] = [-4.0, -3.991991991991992, -3.983983983983984]

这和 PyTorch、不 jit 的 JAX 一样。同一个 Python 函数放进 jax.jit 就报错：

jit: TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
  ... This concrete value was not available in Python because it depends on the value of the argument x.

traced 也一样：中间结果在建图时没有数（max(x).attached_value().is_some() = false），只能用 where_select 两边都算再挑，得到同一个梯度（逐位相同，梯度程序 13 个运算），细节见 记图与 jit 的「控制流」一节。分支取决于数据、而且每次可能走不同的路时，eager 是唯一不用改写的选择。

怎么选

• 形状、步数固定，同一个梯度要算很多次（扫参数、优化循环）：traced。准备一次，之后每次只付执行；同一个梯度和 eager 逐位相同，可以先用 eager 写对，再换。
• 分支依赖数值、形状每次不同、只算一两次：eager。浅图上一次 backward 几十 µs 的固定开销可以接受。
• 图很深或有重复使用的多输出运算（TRG 的 SVD）：两种模式都要分段。eager 分段的反向直接用残差，这次实测比 traced 分段快；traced 分段每种形状只建一个程序，准备一次之后所有 β 复用，没有每次反向的簿记。
• 在 eager 里不要把很长的链交给一次 backward：开销按每次反向付，不会被缓存掉。

动脑筋答案

== shape change, eager
n=1001: backward    197.0 us  entries=2 hits=500 misses=2
n=1001: backward     64.6 us  entries=2 hits=501 misses=2
n=1000: backward     60.9 us  entries=2 hits=502 misses=2

• n = 1001 的第一次未命中（misses 从 1 到 2），要做变换、编译、准备，197 µs，是三次里最慢的；缓存多了一条（entries = 2）。
• 第二次 n = 1001 命中新的那条，64.6 µs。
• 回到 n = 1000 命中旧的那条，60.9 µs。缓存键里的输入形状是具体的，两条并存，互不淘汰（evictions 仍为 0）。
• 变换缓存（ad_transform_cache）同样按形状分：另打印一次，n = 1001 之后它是 entries: 2, hits: 0, misses: 2，换形状时变换也重做了。这个未命中的 197 µs 比进程里第一次 backward（528 µs）小，差在哪里本页没有再追。JAX 的 jit 换形状要重新追踪并交给 XLA 编译，记图与 jit 量过那个时间。

小结

• eager 和 traced 的梯度逐位相同（向量函数和 4 步 TRG 都验证过）；JAX 的 jit 和不 jit 差在末位，因为 XLA 把乘加融合成了 FMA。
• eager 每次 backward 都要整理 trace、建导数 trace 再丢掉：浅图约 60 µs 起步，400 层的链 24 ms，是 traced 梯度程序的 39 倍；准备好的程序命中缓存也省不掉。
• 大数组上开销摊薄，n = 10⁶ 单线程 eager 是 traced 的 1.5 倍；不 jit 的 JAX 每次约 1.1 ms，不是 eager 的对照物，jit 才是。
• 整图的 eager TRG 4 步每次 backward 约 7 秒，缓存无济于事；按步分段、每步一个 backward，20 步 0.64 s，u 和 traced 分段到所列位数相同，比 traced 分段快，因为反向直接用残差。
• eager 的导数程序缓存按具体形状分，换形状重新准备、换回来命中；结果能在前向时读数，按值分支不用改写。