checkpoint
主线位置：TRG 主线的「分段求导得内能」，看两边名叫 checkpoint 的东西能不能代替分段。
版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native，单线程）
flowchart TB jc["jax.checkpoint"] --> rm["反向时重算前向"] rm --> mem["省内存 多算一遍"] tc["tenferro checkpoint"] --> ev["现在就求值"] ev --> leaf["换成带数据的叶子"] ev --> keep["旧图挂在链表上"] keep --> grad["求导时接回旧图"] keep --> grow["链表合并不去重"]
两个函数同名，做的事情不同。jax.checkpoint（又叫 remat）决定反向时哪些中间值重算、哪些存下来，是内存和时间的取舍，数学上不改变导数。tenferro 的 checkpoint 是「现在就把这个张量算出来，换成一个带数据的叶子，旧图留着给求导回放」。本页分别量两者，再说 tenferro 里和 jax.checkpoint 真正对应的是什么。
实验代码：JAX 一侧 trg_scan.py、ckpt_diag.py、ckpt_policy.py、ckpt_eps.py、ckpt_window.py、ckpt_window_builtin.py；tenferro 一侧 bin trg_ckpt_chain、ckpt_chain、ckpt_placeholder（release）。M2 Max、macOS 27、rustc 1.99.0；负载均值 4.7–9.4。TRG 都是 χ = 16、截断数 [2, 4, 4, 16, ...]、β = 0.3。可能跑飞的地方加了上限：ckpt_chain 只跑到 18 轮，tenferro 的 TRG 加检查点只跑到 6 步，每次都用 timeout（120 s 或 300 s）包住。
jax.checkpoint：重算换内存
给 TRG 的每一步套 jax.checkpoint，反向时只用每步的输入，步里面的中间值重算：
st = jax.checkpoint(step, static_argnums=(1,)) # step(t, k) -> (t_next, ln_norm)
for n, k in enumerate(ks):
t, ln = st(t, k)
acc = acc + 2.0 ** (steps - n) * ln
# 20 步，jit(grad)，负载均值 7.1
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)
unrolled + checkpoint grad=0.704389022843801 |d|=1.1e-04 trace+lower 141.3 ms compile 556.0 ms first call 467.0 ms steady 334.0 ms temp 35.27 MiB
scan + checkpoint grad=0.704427618195070 |d|=7.1e-05 trace+lower 57.4 ms compile 275.9 ms first call 371.7 ms steady 340.0 ms temp 35.83 MiB (head 4 unrolled)
（temp 是 compiled.memory_analysis().temp_size_in_bytes，XLA 为这个程序分配的临时缓冲。）
内存和时间符合预期：临时缓冲从 197 MiB 降到 35 MiB，稳态从 226 ms 涨到 334 ms，多出的正好约是一遍前向（求值 161 ms 里的大部分）。
梯度不符合预期：差了 1.1e-4。u 和精确解的误差本来是 3.9e-7，现在被这个差淹没了。看是哪一步开始的：
steps= 4 grad=0.844054073624 grad(checkpoint)=0.844054073624 again=0.844054073624 fd=0.844054073951 |grad-fd|=3.3e-10 |ckpt-fd|=3.3e-10
steps= 5 grad=0.744793433247 grad(checkpoint)=0.744441835949 again=0.744441835949 fd=0.744793433716 |grad-fd|=4.7e-10 |ckpt-fd|=3.5e-04
steps= 6 grad=0.712432002551 grad(checkpoint)=0.712397003365 again=0.712397003365 fd=0.712432002875 |grad-fd|=3.2e-10 |ckpt-fd|=3.5e-05
steps=20 grad=0.704498679530 grad(checkpoint)=0.704389022844 again=0.704389022844 fd=0.704498679716 |grad-fd|=1.9e-10 |ckpt-fd|=1.1e-04
fd 是求值函数的中心差分（h = 1e-5）。1 到 4 步两者相同；从第 5 步起，不加 checkpoint 的梯度和差分一致到 1e-10，加了的偏 1e-4 量级，而且每次运行结果一样（again）。第 5 步是第一次分解 256×256、秩只有 16 的矩阵的地方，也是内置 jnp.linalg.svd 的梯度开始变成 NaN 的地方（简并奇异值）。
换几种存储策略（policy），20 步：
fd of value: 0.704498679716
no checkpoint grad=0.704498679530 |grad-fd|=1.9e-10 temp 197.17 MiB steady 222.7 ms
checkpoint, default policy grad=0.704389022844 |grad-fd|=1.1e-04 temp 35.27 MiB steady 332.2 ms
checkpoint, everything_saveable grad=0.704498679530 |grad-fd|=1.9e-10 temp 230.21 MiB steady 235.3 ms
checkpoint, save only svd outputs grad=0.704382232148 |grad-fd|=1.2e-04 temp 65.84 MiB steady 388.6 ms
全存（everything_saveable）就回到正确的值；只存 SVD 的三个输出（checkpoint_name 加 save_only_these_names）还是错。偏差也不随 SVD 规则里的展宽 ε 变化（ε 从 1e-12 到 1e-3 都在 1.0e-4 到 1.1e-4，ckpt_eps.py），不是展宽那种「巨大的数抵消不净」。拿精确的中间张量做短窗口：单独一步套不套 checkpoint，VJP 完全相同（差 0）；连续两步各套一个，就差到和梯度本身同一量级，换成 JAX 内置的 SVD 规则也一样（steps 3..4 (ks [4, 16]): max|g|=2.369e+00 max|g - g_ckpt|=1.797e+00）。到这里原因还没有查清。可以确定的是：在这个含秩亏 SVD 的链上，jax.checkpoint 改变了 JAX 算出的梯度，而差分站在不加 checkpoint 的那一边。用 jax.checkpoint 省内存之前，先拿差分核对一次梯度。
tenferro 的 checkpoint：现在求值，旧图留着
let beta = TracedTensor::from_vec_col_major(vec![], vec![0.3_f64])?; // 必须带数据
let mut compiler = GraphCompiler::new();
let lnz = trg_with(&beta, &ks, None, &mut |t| {
t.checkpoint(&mut compiler, &rt, &mut eng)?; // 每一步之后
Ok(())
})?;
let g = ad.grad(&lnz, &beta)?;
checkpoint 编译、准备、执行 t 到目前为止的图，把结果作为数据挂上，再由 checkpoint_tensor 把 t 的图换成一个新的输入叶子，旧图和「新叶子就是旧图的那个输出」这层对应记进一个 CheckpointNode，挂在 t 的检查点链表上。之后的运算从这个叶子接着建图，求值时不再重算前面；求导时编译器把链表里的旧图和对应关系接回去，梯度能一直传到 beta。
因为要现在就执行，图里不能有占位符：
checkpoint on placeholder graph: Err: placeholder User { id: 0 } has no runtime input binding
所以 β 只能是数据，每个 β 都要重新建图、重新检查点，失去了分段那种「准备一次、所有 β 复用」。它省下的是求导变换的时间：
traced TRG + checkpoint after every step, chi=16 steps=4 ks=[2, 4, 4, 16] beta=0.30 stack=None MiB
after step 1: checkpoint 1.12 ms, chain nodes 1 (distinct graphs 1)
after step 2: checkpoint 0.47 ms, chain nodes 25 (distinct graphs 2)
after step 3: checkpoint 0.52 ms, chain nodes 601 (distinct graphs 3)
after step 4: checkpoint 1.93 ms, chain nodes 14425 (distinct graphs 4)
lnZ/N graph built with checkpoints: 5.5 ms, chain nodes on lnZ/N 15679 (distinct graphs 4)
ad.grad 4.2 ms | grad compile+prepare+run 12.0 ms (871 ops) | grad=0.844054073623855
drop 0.9 ms | total 22.7 ms
4 步的 ad.grad 4.2 ms，不加检查点是 2125 ms（求导变换的代价），梯度逐位相同（0.844054073623855）。检查点之后的值从一个新的输入键开始，结构键不再一路递归到 beta，比较就便宜了；这是从机制推出来的，没有单独采样核实。
但每一行的 chain nodes 在涨：1、25、601、14425，每步约乘 24，而不同的图只有 4 个。继续往下：
traced TRG + checkpoint after every step, chi=16 steps=5 ks=[2, 4, 4, 16, 16] beta=0.30 stack=None MiB
...
after step 5: checkpoint 54.58 ms, chain nodes 346201 (distinct graphs 5)
lnZ/N graph built with checkpoints: 81.5 ms, chain nodes on lnZ/N 376305 (distinct graphs 5)
ad.grad 19.7 ms | grad compile+prepare+run 68.3 ms (1073 ops) | grad=0.744793433247367
thread 'main' (10083079) has overflowed its stack
fatal runtime error: stack overflow, aborting
exit=134
traced TRG + checkpoint after every step, chi=16 steps=6 ks=[2, 4, 4, 16, 16, 16] beta=0.30 stack=Some("1024") MiB
...
after step 6: checkpoint 730.30 ms, chain nodes 8308825 (distinct graphs 6)
lnZ/N graph built with checkpoints: 1344.3 ms, chain nodes on lnZ/N 9031331 (distinct graphs 6)
ad.grad 393.9 ms | grad compile+prepare+run 103.3 ms (1272 ops) | grad=0.712432002551333
drop 222.3 ms | total 2188.4 ms
- 5 步时梯度算出来了（和 JAX 5 步的
0.744793433247367相同），然后在释放时栈溢出：链表的Arc是一个接一个递归释放的。用 lldb 截下崩溃时的栈，顶上 40 层全是<alloc::sync::Arc<tenferro_runtime::checkpoint::CheckpointNode>>::drop_slow。 - 放到 1 GiB 栈的线程上，6 步能跑完，链表 900 万个节点，建图加检查点 1.3 s，释放 0.22 s，梯度正确（和 JAX 6 步的
0.712432002551333相同）。7 步按每步 24 倍估计是两亿多个节点，没有跑。 - 试验阶段记的是「4 步能救、6 步栈溢出」。实际上 5 步就在默认的主线程栈上溢出了，只是溢出发生在梯度算完之后的释放阶段。
最小复现，一个标量反复平方、每次检查点：
let x = TracedTensor::from_vec_col_major(vec![], vec![1.01_f64])?;
let mut y = x.clone();
for n in 1..=max_n {
y = (&y * &y)?;
y.checkpoint(&mut compiler, &rt, &mut eng)?;
// 数链表节点数，计时 ad.grad(&y, &x)
}
n= 1 chain nodes 1 (distinct graphs 1) | mul+checkpoint 3.85 ms | ad.grad 0.32 ms, grad chain nodes 0, drop 0.00 ms
n= 2 chain nodes 3 (distinct graphs 2) | mul+checkpoint 0.05 ms | ad.grad 0.09 ms, grad chain nodes 0, drop 0.00 ms
n= 3 chain nodes 7 (distinct graphs 3) | mul+checkpoint 0.04 ms | ad.grad 0.09 ms, grad chain nodes 0, drop 0.00 ms
...
n=17 chain nodes 131071 (distinct graphs 17) | mul+checkpoint 7.10 ms | ad.grad 5.79 ms, grad chain nodes 0, drop 0.01 ms
n=18 chain nodes 262143 (distinct graphs 18) | mul+checkpoint 14.12 ms | ad.grad 11.55 ms, grad chain nodes 0, drop 0.01 ms
thread 'main' (10081053) has overflowed its stack
fatal runtime error: stack overflow, aborting
exit=134
（节点数用 tenferro_runtime::ad_support::checkpoint_chain 取出链表头、collect_graphs() 数出来；这是给 tenferro-ad 用的内部接口。）
n 轮之后链表有 2ⁿ − 1 个节点，不同的图只有 n 个。y * y 的两个操作数带着同一条链表，二元运算把两边的链表合并，merge_chains：
(Some(lhs_head), Some(rhs_head)) => {
let mut nodes: Vec<&CheckpointNode> = Vec::new();
let mut current: Option<&CheckpointNode> = Some(&lhs_head);
while let Some(node) = current {
nodes.push(node);
current = node.prev.as_deref();
}
let mut prev: Option<Arc<CheckpointNode>> = Some(rhs_head);
for node in nodes.into_iter().rev() {
prev = Some(Arc::new(CheckpointNode {
graph: node.graph.clone(),
// ...
prev,
}));
}
prev
}
左边链表的每个节点复制一份接到右边链表的前面，不检查两边是不是同一条链、有没有共同的尾巴。同一条链和自己合并，长度就翻倍。TRG 的一步里 T 参与了很多次二元运算（求范数、归一化、两种矩阵化之后的收缩），所以每步乘 24。运算一多，Arc 链表的递归释放就先撞上栈。
PR #712（「fix: merge both checkpoint chains in apply_binary for symmetric grad」，2026-04-13 合并，首次发布于 v0.1.0）之前，二元运算用 .or() 只保留左操作数的检查点链表。提交说明写「Previously apply_binary used .or() which only preserved the lhs checkpoint chain, silently dropping the rhs chain. This caused grad() to panic with 'vjp output is inactive' when differentiating w.r.t. the rhs operand after both operands were independently checkpointed」。于是加了 merge_chains，「links both chains into a single linked list」，einsum 的路径也改成把所有输入的链表折叠起来。PR 新增的测试全是「两个操作数各自独立检查点」（grad_both_independently_checkpointed_add_wrt_rhs 等），两条链表互不相交；两个操作数共用同一条链表的情况（y * y，或者 TRG 里同一个 T 用多次）没有测，合并就不去重地把它复制了一遍。
要修的话，merge_chains 至少要在两条链表是同一条时（Arc::ptr_eq(&lhs_head, &rhs_head)）直接返回一边；有公共尾巴的情况还要按 alias_key 去重。本页没有改上游代码去验证这个修法，只把复现留在实验代码里。
y * y 换成 y + x（x 是没有检查点的原始叶子），n = 10 时又有几个？按上面 merge_chains 的代码推。
tenferro 里谁对应 jax.checkpoint
jax.checkpoint 的效果是「反向只用每步的输入，步内的前向重算」。tenferro 里有这个效果的是 traced 分段（分段）：只存 20 个 Tₙ，每步的反向程序带着这一步的前向，求值加梯度 933 ms，是求值的 2.8 倍；JAX 每步 checkpoint 之后是 334 ms，是求值的 2.1 倍。存残差、不重算的对应物是 eager 分段，以及 JAX 的不加 checkpoint 或 everything_saveable。
| 想要的效果 | JAX | tenferro |
|---|---|---|
| 反向重算每步、只存步的输入 | 每步 jax.checkpoint | traced 分段 |
| 存残差、不重算 | 不加 checkpoint，或 scan | eager 分段 |
| 把一段图先算出来、之后从数据接着建 | 没有直接对应（最接近的是不 jit 的逐步执行） | checkpoint，β 必须是数据，步数多了链表爆炸 |
动脑筋答案
前一问：n=10 chain nodes 1023 (distinct graphs 10)，2¹⁰ − 1 = 1023。每一轮先把同一条链表和自己合并（长度翻倍），检查点再在前面加一个节点：Lₙ = 2·Lₙ₋₁ + 1。
后一问：ckpt_chain 12 add（输出里的 mul+checkpoint 是固定的标签，这次的运算是加法）：
n= 1 chain nodes 1 (distinct graphs 1) | mul+checkpoint 3.80 ms | ad.grad 0.19 ms, grad chain nodes 0, drop 0.00 ms
...
n=10 chain nodes 10 (distinct graphs 10) | mul+checkpoint 0.03 ms | ad.grad 0.11 ms, grad chain nodes 0, drop 0.00 ms
n=12 chain nodes 12 (distinct graphs 12) | mul+checkpoint 0.03 ms | ad.grad 0.12 ms, grad chain nodes 0, drop 0.00 ms
- 10 个，每轮加一个。
x没有链表，merge_chains(Some(..), None)走(lhs, None) => lhs，原样返回左边那条，不复制。 - 所以爆炸的条件是「两个操作数带着同一条（或有公共尾巴的）链表」。TRG 里
T和自己相乘、和自己的变形相乘，正是这种情况。
小结
jax.checkpoint用重算换内存：TRG 20 步的临时缓冲 197 → 35 MiB，稳态 226 → 334 ms。- 在这条含秩亏 SVD 的链上，
jax.checkpoint（默认策略、或只存 SVD 输出）让梯度偏 1e-4，从第 5 步开始，可复现，原因没有查清；everything_saveable正确。省内存之前先用差分核对。 - tenferro 的
checkpoint是「现在求值、换成叶子、旧图留给求导」，图里不能有占位符；它让 4 步 TRG 的ad.grad从 2.1 s 降到 4 ms，梯度不变。 - 但它的检查点链表在二元运算合并时不去重：
y·y每轮翻倍，TRG 每步乘 24；5 步在 8 MiB 栈上释放时溢出，6 步要 1 GiB 栈和 900 万个节点。 - tenferro 里对应
jax.checkpoint的是 traced 分段；对应「不 remat」的是 eager 分段。