detach() 之后还在记图
主线位置：训练主线的 optimizer step 之后：每步从参数里取一份「不要梯度」的快照，更新 EMA 或目标网络。
版本：PyTorch 2.15.0（CPU，
torch.set_num_threads(1)）、burn 0.22.1（Flex 设备加 autodiff，release 构建）
flowchart TB w["参数 w 叶子 要梯度"] -->|"PyTorch detach"| c["常量"] w -->|"burn detach"| leaf["新叶子 仍要梯度"] leaf --> ema["ema 仍在记图"] ema --> chain["每步链长加一"] chain --> slow["backward 清理变慢"] ema --> loss["进 loss 第二步 panic"] w -->|without_autodiff| const2["常量"]
实验 crate pita-t 的 p10，PyTorch 侧 p10_ema.py。两边都是 Linear(512, 512)、输入 [64, 512]、SGD 学习率 1e-3、2000 步，每隔一段打印这一段的平均每步耗时和 RSS 增量（ps -o rss=）。burn 用的是模板的 Flex（关掉了默认 feature，没有 rayon），两边的绝对耗时不可比，只看各自随步数怎么变。测量时 uptime 的 1 分钟负载约 2.5。平台 M2 Max、macOS 27、rustc 1.99.0。
PyTorch 的写法
EMA（指数滑动平均）权重、BYOL / DQN 一类的目标网络，都在每步更新之后取一份参数的快照：
ema = model.weight.detach()
for step in range(1, steps + 1):
out = model(x)
if mode == "target":
loss = (out - x @ ema).pow(2).mean()
else:
loss = out.pow(2).mean()
opt.zero_grad()
loss.backward()
opt.step()
ema = ema * 0.99 + model.weight.detach() * 0.01
detach step 1: 9.65 ms/step rss + 10.6 MiB ema.requires_grad=False grad_fn=None
detach step 50: 0.32 ms/step rss + 13.3 MiB ema.requires_grad=False grad_fn=None
...
detach step 2000: 0.28 ms/step rss + 18.9 MiB ema.requires_grad=False grad_fn=None
target step 1: 9.90 ms/step rss + 11.0 MiB ema.requires_grad=False grad_fn=None
...
target step 2000: 0.31 ms/step rss + 22.6 MiB ema.requires_grad=False grad_fn=None
detach() 的结果不要梯度、没有 grad_fn，ema 是一个普通的常量张量。每步耗时不随步数变。
搬到 burn 的现象
let mut ema = model.weight.val().detach();
for step in 1..=steps {
let out = model.forward(x.clone());
let loss = if mode.starts_with("target") {
// pull the online output toward a slow-moving target network
(out - x.clone().matmul(ema.clone())).powi_scalar(2).mean()
} else {
out.powi_scalar(2).mean()
};
let grads = GradientsParams::from_grads(loss.backward(), &model);
model = optim.step(1e-3, model, grads);
ema = ema * 0.99 + model.weight.val().detach() * 0.01;
}
（实验程序里快照写成一个闭包，按模式选 detach() 或 without_autodiff()，下面「改法」里是另一种。）EMA 只用来保存、不进 loss：
detach step 1: 4.60 ms/step rss + 3.8 MiB ema.is_tracked=true is_require_grad=false
detach step 50: 1.03 ms/step rss + 7.9 MiB ema.is_tracked=true is_require_grad=false
detach step 100: 0.98 ms/step rss + 8.0 MiB ema.is_tracked=true is_require_grad=false
detach step 200: 1.11 ms/step rss + 8.3 MiB ema.is_tracked=true is_require_grad=false
detach step 500: 1.78 ms/step rss + 9.0 MiB ema.is_tracked=true is_require_grad=false
detach step 1000: 5.44 ms/step rss + 10.7 MiB ema.is_tracked=true is_require_grad=false
detach step 2000: 21.73 ms/step rss + 14.7 MiB ema.is_tracked=true is_require_grad=false
ema.is_tracked=true：它在记图。每步耗时从 1 ms 涨到 21.7 ms，越往后涨得越快；RSS 慢慢涨。不报错，训练结果也是对的，只是越跑越慢。
EMA 作为目标网络进了 loss：
target step 1: 2.03 ms/step rss + 5.0 MiB ema.is_tracked=true is_require_grad=false
thread 'main' (8846881) panicked at .../burn-autodiff-0.22.1/src/runtime/server.rs:74:17:
Cannot run backward: graph tape has already been consumed at NodeId(16). Burn does not support retain_graph. Combine losses before backward or recompute the forward pass. Use detach() only to intentionally sever gradients through the previous computation.
（本地路径前缀换成了 ...。）第二步就 panic。报错建议「Use detach()」，而问题恰恰出在 detach() 上。
原因
burn 的 Tensor::detach 文档写着「A leaf tensor also preserves its current require_grad setting」，autodiff 后端的 float_detach 第一行注释是「we still want to keep the require_grad (and distributed) setting」。参数是要梯度的叶子，detach() 之后得到一个新 id、同样要梯度的叶子。两边 detach 的差别见不求导的区域的「detach 不一样」。
于是 ema * 0.99 + snapshot * 0.01 每一步都是被追踪的运算，ema 挂着一条从第一步连到现在的图，而且 ema 一直被引用着，图里的节点一个也释放不了。
- 越跑越慢：每次 backward 之后，
cleanup_orphaned_entries复制一份全局的图表（GraphLocator里的graphs，每个节点一条），逐条检查所属的图还有没有用（maybe_useful()扫描图里的节点）。EMA 链每步多几个节点，这一步的工作量跟着变大。对照实验：同样的 EMA 更新，去掉 forward 和 backward（ema-only），2000 步每步一直是 0.07 ms，RSS 每步多约 3 KB；时间的增长来自 backward 之后的清理。 - 第二步 panic：EMA 进了 loss，
ema所在的图和 loss 的图合成了一张，第一步的 backward 把它消费掉；第二步的ema建在第一步的ema上，第二步的 backward 沿着它走回第一步，碰到已经消费掉的节点，就是图只能用一次那个 panic。
这几处在 origin/main（2026-10-28）上没有改动：git diff v0.22.1 origin/main 对 burn-autodiff 的 runtime/、graph/ 和 float_detach 都没有输出。detach 保留 require_grad 是写在文档里的设计；清理的开销随存活节点数增长，我认为是 0.22.1 的一个性能问题，不过只有在「图被一直引用着」这种用法下才显出来。
改法
要一个真正的常量，用 without_autodiff()（inner() 是同一个函数）。它把张量移到内层后端、丢掉图的引用；文档写着和带 autodiff 的张量一起运算时「it is treated as a constant for that operation」，所以目标网络的 x.matmul(ema) 照样能写：
let snapshot = |m: &Linear| m.weight.val().without_autodiff();
let mut ema = snapshot(&model);
// ...
ema = ema * 0.99 + snapshot(&model) * 0.01;
without step 1: 1.40 ms/step rss + 3.7 MiB ema.is_tracked=false is_require_grad=false
without step 50: 0.94 ms/step rss + 7.7 MiB ema.is_tracked=false is_require_grad=false
...
without step 2000: 0.93 ms/step rss + 9.8 MiB ema.is_tracked=false is_require_grad=false
target-without step 1: 1.69 ms/step rss + 3.8 MiB ema.is_tracked=false is_require_grad=false
target-without step 50: 1.34 ms/step rss + 7.8 MiB ema.is_tracked=false is_require_grad=false
...
target-without step 2000: 1.32 ms/step rss + 8.2 MiB ema.is_tracked=false is_require_grad=false
每步耗时不再随步数变，目标网络的版本跑满 2000 步。
只在日志里用一下快照（每步算 w.detach() 的范数、into_scalar() 读走）没有这个问题：同一个程序的 log 模式跑 2000 步，每步一直是 1.13 ms，RSS 停在 +8.0 MiB。结果没有被留住，那一小段图在下一次 backward 后就被清掉了。会出问题的是把 detach 的结果留下来、一步接一步往上累积的写法。
怎么检查
- 凡是 PyTorch 代码里的
.detach()，先问它的输入是参数（叶子）还是中间结果。中间结果两边一样；参数在 burn 里要写without_autodiff()。 - 长期保存的张量（EMA、目标网络、running stats、缓存），打印一次
is_tracked()，应当是false。 - 训练跑久了每步变慢、而模型没变，先怀疑有一条图被一直留着。
小结
- burn 的
detach()对要梯度的叶子保留require_grad，结果照样记图；PyTorch 的detach()得到常量。 - 用
p.detach()累积 EMA，每步耗时从 1 ms 涨到 21.7 ms（2000 步）；EMA 进 loss，第二步 panic「graph tape has already been consumed」。 - 要常量就
without_autodiff()，和 autodiff 张量一起运算时被当作常量。