Rust in Everything
基于 tenferro-ad 0.8.0

分段

主线位置：TRG 主线的「分段求导得内能」：20 步不建成一张图，每种形状的一步一个程序，前向逐步跑，反向从最后一步往回串。
版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native，单线程）
flowchart TB
  b["beta"] --> i0["初值程序"]
  i0 --> t0["T0"]
  t0 --> s1["步程序 2 2"]
  s1 --> t1["T1"]
  t1 --> sx["步程序 按形状复用"]
  sx --> tn["T20"]
  tn --> fin["末尾 ln tr"]
  fin --> ctn["ct20"]
  ctn --> bk["步的反向程序"]
  bk --> ct0["ct0"]
  ct0 --> ib["初值的反向"]
  ib --> u["d lnZ/N d beta"]

JAX 里 20 步的 TRG 写成 Python 循环，jit(grad(f)) 把它整个展开、整个求导，第一次调用 1.2 s，之后每次 223 ms（计时）。同样的写法在 tenferro 里走不通：整图的 ad.grad 4 步就要 2.1 s，5 步 90 s 内做不完（求导变换的代价）。tenferro 也没有 scan、while 这类控制流算子（记图与 jit 的「控制流」）。本页先看 JAX 的 lax.scan 做了什么，再看 tenferro 怎么用一组按形状复用的小程序把同一件事手工做出来。

实验代码：JAX 一侧 trg_scan.py，tenferro 一侧 bin trg_staged（release），eager 的分段见 trg_eager_staged。M2 Max、macOS 27、rustc 1.99.0；负载均值 4.7–7.1。χ = 16，截断数 [2, 4, 4, 16, 16, ...]，JAX 用和 tenferro 默认规则相同的 custom_jvp SVD。

JAX：lax.scan 只能扫相同形状的步

TRG 的前 4 步形状各不相同（(d, k) 是 (2, 2)、(2, 4)、(4, 4)、(4, 16)），之后 16 步都是 (16, 16)。lax.scan 的循环体只追踪、编译一次，要求每次迭代的形状相同，所以前 4 步留在 Python 循环里，后 16 步交给 scan；每步 ln‖T‖ 的权重 2^(N−n) 作为扫描的输入：

def body(t, wn):
  t, ln = step(t, 16)
  return t, wn * ln

t = initial_tensor(beta)
for n in range(4):                         # 形状各不相同的前 4 步
  t, ln = step(t, ks[n])
  acc = acc + 2.0 ** (steps - n) * ln
t, terms = lax.scan(body, t, w)            # 后 16 步，w = [2^(20-4), ..., 2^1]
for i in range(terms.shape[0]):            # 按展开时的顺序累加
  acc = acc + terms[i]

# 负载均值 7.1，jit(grad)，β = 0.3
unrolled (trg_jax.make_trg)  grad=0.704498679529762 |d|=0.0e+00  trace+lower   398.0 ms  compile   520.0 ms  first call   374.0 ms  steady   225.8 ms  temp 197.17 MiB
scan                         grad=0.704498679529763 |d|=6.7e-16  trace+lower   133.0 ms  compile   325.9 ms  first call   261.3 ms  steady   231.8 ms  temp 240.99 MiB  (head 4 unrolled)
  • scan 把追踪加降级从 398 ms 降到 133 ms、编译从 520 ms 降到 326 ms：循环体只处理一次。
  • 稳态几乎不变（226 对 232 ms），计算量没有变。临时缓冲区反而大了一点（197 对 241 MiB），scan 要把每次迭代的残差堆成一个数组留给反向。
  • 梯度差 6.7e-16，一个 ulp：两种写法的求和顺序已经对齐，剩下的是 XLA 对循环体和展开代码的不同融合。
📝NOTE
没有蠢问题

问： 我的 custom_jvp SVD 在展开的循环里好好的，一放进 lax.scan 就报错，为什么？

答： 报错是 Custom JVP rule lorentz_svd_jvp for function lorentz_svd must produce a pair ... the JVP rule output's first element had container/pytree structure: (float64[256,256], float64[256], float64[256,256]) while the custom_jvp-decorated function lorentz_svd had output container/pytree structure: SVDResult(U=float64[256,256], S=float64[256], Vh=float64[256,256])。jnp.linalg.svd 返回的是具名元组 SVDResult，规则里返回的是普通元组，两种 pytree 结构不同。展开的循环里同一个规则没有报错，scan 的路径做了这个检查。函数本体改成 return tuple(jnp.linalg.svd(a, full_matrices=False)) 就好了；本页的数字都是改过之后的。

tenferro：每种形状一个小程序

lax.scan 背后的想法是「相同的步只处理一次」。tenferro 里把它做成显式的：按 (d, k) 把步分组，每组建三个 traced 程序，输入全是占位符，所以 20 步、所有 β 共用：

for &k in ks {
  if !steps.contains_key(&(d, k)) {
    let t = TracedTensor::input_concrete_shape(DType::F64, &[d, d, d, d])?;
    let ct = TracedTensor::input_concrete_shape(DType::F64, &[k, k, k, k])?;
    let c = TracedTensor::input_concrete_shape(DType::F64, &[])?;
    let (tn, ln) = trg_step(&t, d, k, None)?;
    let l = (&(&ct * &tn)?.reduce_sum(None)? + &(&c * &ln)?)?;
    let g = ad.grad(&l, &t)?;
    steps.insert((d, k), StepProgs {
      fwd: mk(rt, eng, &tn, &[(&t, vec![d, d, d, d])])?,
      lognorm: mk(rt, eng, &ln, &[(&t, vec![d, d, d, d])])?,
      bwd: mk(rt, eng, &g, &[(&t, vec![d, d, d, d]), (&ct, vec![k, k, k, k]), (&c, vec![])])?,
    });
  }
  d = k;
}

（mk 是实验代码里的小函数：compile_with_input_specs 加 prepare，prepare 用一组形状对的零张量。初值和末尾的 ln tr T 各另有一对程序。）

反向程序 bwd 是这一步的 VJP。整条链上，记 lnZ/N = s·(Σₙ 2^(N−n) ln‖Tₙ‖ + ln tr T_N)，s = 2^(−N)，ctₙ 是 lnZ/N 对 Tₙ 的导数。Tₙ 只经过第 n 步影响结果：直接通过 ln‖Tₙ‖，间接通过 T_{n+1}。所以

ctₙ = ∂/∂Tₙ [ ⟨T_{n+1}(Tₙ), ct_{n+1}⟩ + cₙ · ln‖Tₙ‖ ]，cₙ = s·2^(N−n)，

把 ct_{n+1} 和 cₙ 当常数。这正是 bwd 里那个标量 l 的梯度：ct 和 c 是占位符，运行时从后往前喂进去。和 JAX 的对照是 jax.vjp(step, t) 返回的函数用余切 (ct, c) 调一次；lax.scan 的反向在内部做的就是这件事。

求值和梯度：

let mut ts = vec![run(rt, eng, &s.init_fwd, &[&bt])?];
for (i, &k) in s.ks.iter().enumerate() {
  let p = &s.steps[&(d, k)];
  let ln = run(rt, eng, &p.lognorm, &[&ts[i]])?.as_slice::<f64>()?[0];
  acc += 2f64.powi((n - i) as i32) * ln;
  ts.push(run(rt, eng, &p.fwd, &[&ts[i]])?);
  d = k;
}
// ...
let mut ct = /* s 乘 fin_bwd 的输出：d ln tr T_N / d T_N */;
for i in (0..n).rev() {
  let c = scalar(2f64.powi((n - i) as i32) * scale)?;
  ct = run(rt, eng, &s.steps[&(dims[i], dims[i + 1])].bwd, &[&ts[i], &ct, &c])?;
}
let gb = run(rt, eng, &s.init_bwd, &[&bt, &ct])?.as_slice::<f64>()?[0];

tenferro 0.8.0 staged chi=16 steps=20 schedule=rank width=1 ks=[2, 4, 4, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16]
build+grad+compile+prepare all programs 27.4 ms (5 step shapes, 1330 ops in total)
  step (d, k) = (2, 2): fwd 40 ops, bwd 204 ops
  ...
  step (d, k) = (16, 16): fwd 40 ops, bwd 204 ops
beta=0.30 lnZ/N=0.790559063800966 exact=0.790559070951263 err=-7.150e-9
  u=-0.704498679529754 exact=-0.704499070832445 err=+3.913e-7 | grad=0.704498679529754 fd=0.704498679671150 |grad-fd|=1.414e-10
  run ms: value first 365.68 steady 336.906; value+grad first 937.81 steady 933.257
βlnZ/N 误差u 误差|grad − 差分|求值 / 求值加梯度（ms）
0.30−7.150e-9+3.913e-71.4e-10337 / 933
0.40−7.605e-7+3.380e-57.7e-10333 / 899
0.44−4.578e-6−5.583e-41.837e-4346 / 926
0.50+4.576e-8−1.889e-56.9e-10332 / 898
0.60+6.399e-7−7.621e-71.3e-10331 / 905
  • 每种步的程序只有 40 个前向算子、204 个反向算子，ad.grad 每次只看一步，指数增长没有机会出现。全部 5 种步加初值和末尾，1330 个算子，27 ms 建好。
  • 和 JAX 的梯度（0.704498679529762）差 8e-15；五个 β 上的最大差见 怎么读这一卷 的「线程数也会改末位」（lnZ/N 5.6e-16，梯度 2.4e-12）。β = 0.44 的误差和差分偏差来自 χ = 16 的截断，靠近临界点，两边一样。
  • 梯度比求值贵 2.7 倍：bwd 是 ad.grad 对「从 t 出发的整个一步」求导得到的程序，里面带着这一步的前向，两次 SVD 要重算一遍。

分段的代价：反向重算前向

JAX 的 scan 把每步的残差（SVD 的结果等）存下来给反向用，反向不重算。tenferro 的分段只存每步的输入 Tₙ（20 个，16⁴ 个 f64 一个，共约 10 MiB），反向从 Tₙ 重算这一步。这和 JAX 里给每一步套 jax.checkpoint 是同一种取舍：省内存，多算一次前向。下一页 checkpoint 量了 JAX 这样做的时间和内存。

不重算的分段在 tenferro 里也有，是 eager 写法：每一步是一个 eager 图，反向直接用前向留下的残差，同一个 β 上前向 345 ms 加反向 308 ms，比 traced 分段的 933 ms 快（代码和缓存行为见 eager 还是 traced）：

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 405.60 ms backward 348.74 ms
beta=0.30 lnZ/N=0.790559063800966 u=-0.704498679529753 err=+3.913e-7 hex grad=3fe68b40d094f9a0 | min forward 345.22 ms backward 308.35 ms
...
cache: CacheStats { entries: 7, retained_bytes: 42848, hits: 433, misses: 7, evictions: 0, clears: 0 }

两种分段怎么选：

traced 分段eager 分段
对应的 JAX 写法scan 加每步 jax.checkpointscan（存残差）
准备一次 27 ms，之后所有 β 复用第一次 backward 时按形状准备，缓存 7 条
20 步一个 β（1 线程）求值 337 ms，求值加梯度 933 ms前向 345 ms 加反向 308 ms
内存只存 20 个 Tₙ存每步的残差
只要求值跑前向程序即可，不付记录的开销前向照样记录残差

要大量只求值（扫 β 画 lnZ 曲线）选 traced；要梯度、内存够，选 eager。两种写法的 u 到所列位数相同。

小结

  • JAX 的 lax.scan 只扫形状相同的步，TRG 的前 4 步要展开；它把追踪和编译从 0.92 s 降到 0.46 s，稳态不变，残差照存。
  • tenferro 没有 scan：按 (d, k) 每种步建前向、ln‖T‖、反向三个程序，占位符让 20 步、所有 β 共用，全部 27 ms 准备好。
  • 反向用替身 ⟨T_{n+1}, ct_{n+1}⟩ + cₙ·ln‖Tₙ‖ 串起来，就是逐步的 jax.vjp；五个 β 上梯度和 JAX 最多差 2.4e-12（单线程）。
  • traced 分段的反向程序重算这一步的前向，代价和每步 jax.checkpoint 一样；eager 分段用残差，梯度更快。
  • custom_jvp 放进 lax.scan 时，规则的输出 pytree 要和函数本体完全一致（SVDResult 对普通元组会报错）。

💬 关联讨论

发起讨论
加载中...