Rust in Everything
基于 tenferro-ad 0.8.0

求导变换的代价

主线位置：变换主线的「grad / vjp / jvp」一步，量变换本身花的时间；TRG 主线的「分段求导得内能」为什么必须分段，答案在这一页。
版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native，单线程）
flowchart TB
  y0["y 用了两次"] --> y1["下一层的 y"]
  y1 --> y2["再下一层"]
  y2 --> key["结构键 ValueKey"]
  key --> tree["展开成树 2 的 n 次方"]
  tree --> eq["逐层比较 相等"]
  eq --> slow["变换和释放都变慢"]

JAX 用户很少单独量 grad 这个变换：jax.make_jaxpr(jax.grad(f)) 只是 Python 层的追踪，和图的大小成正比，几十毫秒的事。tenferro 的 ad.grad 也是对图做变换，产物还是图（traced AD），但在 0.8.0 里它的耗时可以随深度指数增长。导数程序本身的大小仍是线性的，执行也快；慢的是变换时的元数据簿记，以及释放导数图的时候。

实验代码：tenferro 一侧是 bin ad_blowup、trg_ad_steps（release，CpuEngine::serial()），JAX 一侧是 jax_ad_cost.py。M2 Max、macOS 27、rustc 1.99.0；负载均值 3.5–5.5（动脑筋那次是 45）。会跑很久的地方都加了上限：ad_blowup 在一次变换超过 1.5 s 后停止，trg_ad_steps 整个进程限 90 s（timeout）。

同一条链，两边的变换时间

三种链，输入都是 3×3 矩阵，深度 n，最后 sum：

  • mul：y ← (y·y)/2，每层把上一层的值用两次。
  • svd3：(u, s, vt) = svd(y)，y ← u·s + vt，一个多输出运算的三个输出都用上。
  • 另外三种只在 tenferro 一侧跑：svd1（只用 u，y ← u + y）、svd2（用 u 和 s）、qr2（y ← q + r）。
def chain(kind, n):
  def f(x):
    y = x
    for _ in range(n):
      if kind == "mul":
        y = (y * y) * 0.5
      else:  # svd3
        u, s, vt = jnp.linalg.svd(y)
        y = u * s[None, :] + vt
    return jnp.sum(y)
  return f

jp = jax.make_jaxpr(jax.grad(f))(x)       # 计时：追踪加变换
lo = jax.jit(jax.grad(f)).lower(x)        # 计时：降级
co = lo.compile()                          # 计时：XLA 编译

let x = TracedTensor::input_concrete_shape(DType::F64, &[3, 3])?;
let mut y = x.clone();
for _ in 0..n {
  y = match kind {
    "mul" => (&y * &y)?.scale_real(0.5)?,
    "svd3" => {
      let (u, s, vt) = y.svd()?;
      let s = s.reshape([3])?.broadcast_in_dim(&[3, 3], &[1])?;
      (&(&u.reshape([3, 3])? * &s)? + &vt.reshape([3, 3])?)?
    }
    // svd1、svd2、qr2 同理
  };
}
let l = y.reduce_sum(None)?;
let d = ad.grad(&l, &x)?;                  // 计时
// 编译 d、数算子；然后单独计时 drop(d)

# JAX
mul n=10 grad jaxpr eqns    62  make_jaxpr(grad)     3.08 ms  lower     5.56 ms  compile    14.50 ms
mul n=25 grad jaxpr eqns   152  make_jaxpr(grad)     6.25 ms  lower     8.19 ms  compile    17.85 ms
svd3 n=10 grad jaxpr eqns   102  make_jaxpr(grad)     7.24 ms  lower    11.83 ms  compile    78.02 ms
svd3 n=15 grad jaxpr eqns   152  make_jaxpr(grad)    10.60 ms  lower    17.52 ms  compile   101.78 ms
# tenferro
mul grad n=10 value ops   23  ad.grad       0.41 ms  ratio   1.34  derivative ops    51  drop derivative       0.09 ms
mul grad n=20 value ops   43  ad.grad      79.52 ms  ratio   2.10  derivative ops   101  drop derivative      77.61 ms
mul grad n=25 value ops   53  ad.grad    2328.13 ms  ratio   1.85  derivative ops   126  drop derivative    2239.12 ms
mul grad: stopped after n=25 (transform > 1500 ms cap)
svd3 grad n=10 value ops   71  ad.grad      12.40 ms  ratio   2.41  derivative ops   574  drop derivative      10.39 ms
svd3 grad n=15 value ops  106  ad.grad    2552.95 ms  ratio   2.99  derivative ops   859  drop derivative    2481.44 ms
svd3 grad: stopped after n=15 (transform > 1500 ms cap)

（ratio 是这一层的变换时间除以上一层的。）

  • JAX 是线性的。 导数 jaxpr 的方程数每层加 6（mul）或 10（svd3），追踪加变换从 3 ms 到 6 ms、从 7 ms 到 11 ms，XLA 编译也只缓慢增长。
  • tenferro 的导数程序也是线性的（mul 每层 5 个算子，svd3 每层 57 个），但 ad.grad 的时间在 mul 上每层翻倍，在 svd3 上每层乘 3。25 层的 y·y 求导 2.3 s，126 个算子的程序。
  • 释放导数图和变换一样贵。 drop(d) 2.2 s 和 2.5 s，几乎等于变换本身。所以只计 ad.grad 的时间会少算一半。

五种链都量到深处（每层的比值）：

链每层重复使用n = 10 的比值最后几层的比值最深一层
muly 用 2 次1.101.94–2.05（n = 18–25）n = 25，2316 ms
svd1y 用 2 次（svd 和加法）1.091.43–1.72（n = 14–16）n = 16，10.5 ms
qr2qr(y) 的 2 个输出1.181.61–1.82（n = 14–16）n = 16，14.7 ms
svd2svd(y) 的 2 个输出1.251.67–1.90（n = 14–16）n = 16，30.6 ms
svd3svd(y) 的 3 个输出2.272.97–3.01（n = 13–15）n = 15，2469 ms

（svd1、qr2、svd2 只跑到 n = 16，比值还在往 2 爬。jvp 也一样：svd3 的 ad.jvp 在 n = 15 是 2627 ms，比值 2.97。）

渐近的比值就是每层把上一层用了几次。试验阶段把这个现象归到「多输出运算（SVD）的输出被重复使用」，范围说窄了：mul 里没有多输出运算，只是 y 用了两次，同样指数增长。多输出运算让它来得更早、底数更大，因为三个输出各算一次使用。

时间去了哪里

用 macOS 的 sample 对 mul 链采样 2 秒（n 在 23 到 25 之间）：主线程 1383 个样本里，915 个在释放导数图（TracedTensor 的 drop → GlobalMetadataScope::drop），468 个在 ad.grad（grad_with_optional_rules → derivative_traces_from_frozen_program → register_scoped_global_metadata_batch）。两条路的栈顶都是同一个函数一层套一层地调用自己：computegraph 的 ValueKey::eq。

ValueKey 是一个值的结构身份：输入，或者「某个运算的第几个输出」，运算又按结构记下它的全部输入键（OperationKey）。比较两个键的 PartialEq：

impl<Op: GraphOperation> PartialEq for ValueKey<Op> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Input(lhs), Self::Input(rhs)) => lhs == rhs,
            (
                Self::Derived {
                    operation: lhs_op,
                    output_slot: lhs_slot,
                },
                Self::Derived {
                    operation: rhs_op,
                    output_slot: rhs_slot,
                },
            ) => {
                lhs_slot == rhs_slot
                    && (Arc::ptr_eq(lhs_op, rhs_op) || lhs_op.as_ref() == rhs_op.as_ref())
            }
            _ => false,
        }
    }
}

两个键指向同一个 Arc 时一步就判完；不是同一个 Arc 时，lhs_op.as_ref() == rhs_op.as_ref() 逐个比较输入键，每个输入键又递归下去，没有记忆。图是 DAG，但这样比较是把 DAG 展开成树来走：y·y 的键有两个输入，都指向上一层的 y，n 层就是 2ⁿ 个叶子；svd3 每层三个输出各指一次 svd(y)，是 3ⁿ。

谁在比较：tenferro 把每个值的元数据（dtype、形状）登记在一张全局表里，键就是 ValueKey（register_scoped_global_metadata_batch）。哈希只用缓存的指纹，是常数时间；指纹相同的键要再用 eq 确认。变换时 derivative_traces_from_frozen_program 为导数图登记元数据，导数图里引用了原图的值；释放时 GlobalMetadataScope 逐个注销，又要按键查表。两次查表都撞上已经登记过、结构相同的键，于是走进完整的递归比较。走到这里的两个键不共用 Arc（否则第一步就返回了），所以导数图里的键是重新构造出来的、结构相等的副本。这一点是从栈和代码推出来的，没有在 tenferro 里加打印核实。

📝NOTE
冷知识：哈希先变成了常数时间，相等比较没有跟上

computegraph 的 PR #3（「Optimize global graph key hashing」，2026-05-19 合并）之前，键类型（当时叫 GlobalValKey / GlobalOpKey）直接 #[derive(Clone, Debug, PartialEq, Eq, Hash)]，运算键内联在值键里：哈希、比较、克隆都递归走整棵树。这个 PR 把运算键放进 Arc，在构造时算好一个指纹，Hash 只哈希指纹；文档注释写「maps keyed by recursively derived values can avoid repeatedly re-hashing the whole input tree. Equality still checks the full structure after the fingerprint prefilter」。同时加的测试 derived_keys_with_distinct_op_arcs_are_structurally_equal 断言两个不同 Arc、结构相同的键相等，这正是上面慢的那条路：指纹让查表找到桶，确认相等仍要走整棵树。computegraph-rs 没有打 tag，第一个发布到 crates.io 的 0.1.0（594cdce）包含这个提交，tenferro 0.8.0 用的就是它。

TRG 整张图

同样的事落在 TRG 上。χ = 16，截断数 [2, 4, 4, 16, ...]，β 是标量占位符，整张图建好后 ad.grad：

let beta = TracedTensor::input_concrete_shape(DType::F64, &[])?;
let lnz = trg(&beta, &ks, None)?;
let g = ad.grad(&lnz, &beta)?;            // 计时
// 编译、准备、执行 g；最后单独计时 drop(g)

eq, t_tr, t_lo, t_co, co = cost(make_trg(ks, lorentz_svd), b)   # 同上的三段计时

# tenferro
steps=1 ks=[2] build 3.44 ms | ad.grad       4.78 ms | drop grad graph       0.23 ms | compile+prepare 4.93 ms | run 0.199 ms | ops value 71 grad 266 | grad=1.667309214024311
steps=2 ks=[2, 4] build 0.18 ms | ad.grad       5.56 ms | drop grad graph       3.91 ms | compile+prepare 5.40 ms | run 0.343 ms | ops value 116 grad 465 | grad=1.277612005080206
steps=3 ks=[2, 4, 4] build 0.20 ms | ad.grad      91.08 ms | drop grad graph      88.82 ms | compile+prepare 6.94 ms | run 0.524 ms | ops value 160 grad 665 | grad=1.023802944795145
steps=4 ks=[2, 4, 4, 16] build 0.31 ms | ad.grad    2125.22 ms | drop grad graph    2099.40 ms | compile+prepare 9.32 ms | run 1.230 ms | ops value 204 grad 863 | grad=0.844054073623855
exit=124
# JAX
trg steps= 1 grad jaxpr eqns   266  make_jaxpr(grad)    37.83 ms  lower    23.28 ms  compile    73.53 ms  grad=1.667309214024311
trg steps= 4 grad jaxpr eqns   837  make_jaxpr(grad)    49.05 ms  lower    63.09 ms  compile   187.54 ms  grad=0.844054073623855
trg steps= 6 grad jaxpr eqns  1239  make_jaxpr(grad)    64.64 ms  lower    93.85 ms  compile   261.77 ms  grad=0.712432002551333
trg steps=20 grad jaxpr eqns  4053  make_jaxpr(grad)   212.64 ms  lower   302.40 ms  compile   481.34 ms  grad=0.704498679529762
  • tenferro 3 步到 4 步，变换乘了 23；5 步在 90 s 的上限内没有做完（exit=124 是 timeout 终止）。试验阶段 6 步跑了 3 分钟没有返回。
  • 释放导数图又是一份同样的时间，4 步一共约 4.2 s。
  • 导数程序线性增长（每步约 200 个算子，1 步时和 JAX 的 266 个方程正好相同），4 步的执行只要 1.2 ms。变换做完以后一切都便宜。
  • JAX 20 步的追踪加变换 213 ms，编译 481 ms，和步数成正比。两边的梯度 1、2、4 步逐位相同，3 步差 1 ulp（...145 对 ...146）。

TRG 每一步里 T 被用了好几次（求范数、两种矩阵化），两次 SVD 的三个输出也都用上，所以底数比 svd3 还大。eager 的 backward 用的是同一个变换，整图 4 步每次约 7 秒，而且每个 β 都要重付（eager 还是 traced）。

❗IMPORTANT
动脑筋：把 mul 换成 y ← (y·y·y)/2，同一个 y 每层用三次。n = 12 时 ad.grad 是 23 ms，n = 16 时是多少？导数程序有几个算子？

同一个 ad_blowup，release，单线程。

怎么避开

指数的底来自深度，所以办法是不让一次变换看见很深的图：

  • 分段。 每种形状的一步建一个小程序，单独 ad.grad，步与步之间手工串反向。TRG 共 5 种步，全部程序的建图、变换、编译、准备一共 27 ms，20 步、所有 β 复用（分段）。这是本卷推荐的写法。
  • checkpoint。 把中间结果求值成叶子，下游的结构键从这个叶子重新开始，4 步的 ad.grad 从 2125 ms 降到 4 ms。但它的检查点链表自己按步数指数变长，5 步就在 8 MiB 的主线程栈上溢出（checkpoint）。
  • 检查一下自己的链。 不只是 SVD：任何「值用两次」的长链都会撞上。迭代 20 次以上、每次都复用上一轮结果的图，先用十几层量一次变换时间的比值，比值趋近 2 或 3 就别整图求导。

动脑筋答案

cube grad n=12 value ops   39  ad.grad      23.07 ms  ratio   2.91  derivative ops   109  drop derivative      22.89 ms
cube grad n=13 value ops   42  ad.grad      69.38 ms  ratio   3.01  derivative ops   118  drop derivative      68.33 ms
cube grad n=14 value ops   45  ad.grad     308.11 ms  ratio   4.44  derivative ops   127  drop derivative     361.42 ms
cube grad n=15 value ops   48  ad.grad    1054.91 ms  ratio   3.42  derivative ops   136  drop derivative    1116.39 ms
cube grad n=16 value ops   51  ad.grad    3218.53 ms  ratio   3.05  derivative ops   145  drop derivative    3413.15 ms
cube grad: stopped after n=16 (transform > 1500 ms cap)
  • 用三次，底数就是 3：n = 13 和 16 的比值都是 3.0。按 3⁴ = 81 估计 n = 16 约 1.9 s，实测 3.2 s；多出来的部分在 n = 14 那一步（比值 4.44），这次运行时机器的负载均值到了 45（另一个进程在跑满 CPU），单次的时间被拉长了。
  • 导数程序 145 个算子，每层多 9 个，线性。
  • 释放又是 3.4 s。和 svd3 比：同样每层乘 3，cube 里没有任何多输出运算。

小结

  • JAX 的 grad 变换和图的大小成正比；20 步 TRG 的追踪加变换 0.2 s。
  • tenferro 0.8.0 的 ad.grad / ad.jvp 在值被重复使用的链上随深度指数增长，底数是每层的使用次数：y·y 是 2，SVD 三个输出都用是 3，TRG 一步约 20；导数程序的大小仍然线性。
  • 释放导数图同样贵，计时时要算上。
  • 时间耗在 computegraph 的 ValueKey::eq：不共用 Arc 的结构相等键被展开成树来逐个比较，查全局元数据表时触发。
  • 整图 TRG 3 步 0.09 s、4 步 2.1 s，5 步 90 s 内没做完；要分段，每段的变换只看一步。

💬 关联讨论

发起讨论
加载中...