checkpoint 之后栈溢出
主线位置：TRG 主线「分段求导得内能」：想用检查点代替分段的那条岔路。
版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native，单线程）
flowchart TB jx["每步 jax.checkpoint"] --> port["每步 t.checkpoint"] port --> s4["4 步 梯度 15 ms"] port --> chain["链表每步乘 24"] chain --> s5["5 步 释放时溢出"] chain --> s6["6 步 建图时溢出"] port --> stage["改成分段"] stage --> ok["20 步正常"]
实验代码：tenferro 一侧是 bin p26_checkpoint（release，CpuEngine::serial()），JAX 一侧是 p26_checkpoint.py。χ = 16，截断数 [2, 4, 4, 16, …]，β = 0.3。M2 Max、macOS 27、rustc 1.99.0。会栈溢出的那几次放在子进程里跑：父进程对 1 到 6 步各起一个子进程，只报告子进程的退出状态，自己不受影响。
现象
JAX 里给每一步套 jax.checkpoint 是常见的省内存写法，步数多少都能跑：
st = jax.checkpoint(trg_scan.step, static_argnums=(1,))
for n, k in enumerate(ks):
t, ln = st(t, k)
steps= 4 grad=0.844054073623855 grad with jax.checkpoint per step=0.844054073623855
steps= 5 grad=0.744793433247367 grad with jax.checkpoint per step=0.744441835948536
steps= 6 grad=0.712432002551333 grad with jax.checkpoint per step=0.712397003365131
steps=20 grad=0.704498679529762 grad with jax.checkpoint per step=0.704389022843801
（从 5 步起加了 jax.checkpoint 的梯度和不加的不同，这是 JAX 一侧的另一个问题，见 checkpoint。）
tenferro 也有 checkpoint，照样每步之后调一次：
let lnz = trg_with(&beta, &ks, &mut |t| {
t.checkpoint(&mut compiler, &rt, &mut eng)?; // 每一步之后
Ok(())
})?;
let g = ad.grad(&lnz, &beta)?;
steps=4 ks=[2, 4, 4, 16]
after step 1: chain nodes 1 (distinct graphs 1)
after step 2: chain nodes 25 (distinct graphs 2)
after step 3: chain nodes 601 (distinct graphs 3)
after step 4: chain nodes 14425 (distinct graphs 4)
build with checkpoints 5.0 ms
ad.grad + run 15.3 ms, grad = 0.844054073623855
drop 0.8 ms
child exit status: exit status: 0
steps=5 ks=[2, 4, 4, 16, 16]
...
after step 5: chain nodes 346201 (distinct graphs 5)
build with checkpoints 75.5 ms
ad.grad + run 75.6 ms, grad = 0.744793433247367
thread 'main' (10173389) has overflowed its stack
fatal runtime error: stack overflow, aborting
child exit status: signal: 6 (SIGABRT)
steps=6 ks=[2, 4, 4, 16, 16, 16]
...
after step 5: chain nodes 346201 (distinct graphs 5)
thread 'main' (10173390) has overflowed its stack
fatal runtime error: stack overflow, aborting
child exit status: signal: 6 (SIGABRT)
- 4 步：
ad.grad加执行 15 ms，不加检查点时光变换就要 2.1 s（求导变换的代价）。梯度和 JAX 不加检查点的值逐位相同。看上去找到了出路。 - 5 步：梯度算出来了，也和 JAX 相同，然后在释放（
drop）时栈溢出，进程被 abort，没有Result可以接住。 - 6 步：第 6 步的
checkpoint自己就溢出了。 chain nodes每步乘约 24（1、25、601、14425、346201），不同的图只有步数那么多。
出处
checkpoint 把当前的图求值成一个新叶子，旧图记成一个 CheckpointNode 挂在链表上，求导时再接回去。二元运算把两个操作数的链表合并，merge_chains 不去重：同一个 T 在一步里被用了好几次，它的链表就被复制好几份，节点数按步数指数增长。链表的 Arc 一个接一个递归释放，几十万个节点就压爆 8 MiB 的主线程栈。逐行的源码、最小复现（标量反复平方，n 轮 2ⁿ − 1 个节点）和 1 GiB 栈的实验见 checkpoint。这是上游的问题，0.8.0 复现。
两边的 checkpoint 也不是一回事：jax.checkpoint 是反向时重算前向、省内存；tenferro 的 checkpoint 是现在就求值、图里不能有占位符（β 只能是数据，每个 β 都要重建），它省的是求导变换的时间。
修法
不用 checkpoint，按步分段：每一步单独建图、单独 vjp，步与步之间手工传余切（分段）。动手页的程序 ising_trg：
tenferro 0.8.0 eager staged chi=16 steps=5 width=1 h=1e-5 ks=[2, 4, 4, 16, 16]
beta=0.3000 lnZ/N=0.792026414104853 u=-0.744793433247366 c=0.371554629451 | value_and_grad 40.1 ms
tenferro 0.8.0 eager staged chi=16 steps=20 width=1 h=1e-5 ks=[2, 4, 4, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16]
beta=0.3000 lnZ/N=0.790559063800966 u=-0.704498679529762 c=0.286288378175 | value_and_grad 555.8 ms
5 步的梯度和上面检查点的结果差 1 ulp（…366 对 …367），20 步和 JAX 不加检查点的梯度打印出来相同。每一步是一张独立的小图，求导变换只看见一步的深度，也就没有要合并的链表。
怎么确认没中招
- 用了
checkpoint的程序，用tenferro_runtime::ad_support::checkpoint_chain(&t)取出链表头、collect_graphs().len()数一次节点（这是给tenferro-ad用的内部接口）：节点数比检查点次数多出几个数量级，就离溢出不远了。 - 栈溢出是 abort，不是 panic，
catch_unwind接不住。要试探上限，像本页这样放在子进程里跑，看退出状态。
小结
- 每步
checkpoint能让 4 步 TRG 的整图梯度从秒级降到毫秒级，但检查点链表每步长 24 倍。 - 5 步时梯度算完、释放时栈溢出，6 步在建图时溢出；进程直接 abort。
- 长链用分段代替
checkpoint；jax.checkpoint对应的是分段，不是 tenferro 的checkpoint。