推理时内存一直涨

主线位置：算子主线上的 Autodiff 那一环。设备是 .autodiff() 时，每个浮点算子都先在 Autodiff 层登记一个节点和一个反向步骤，再交给内层后端；这些记录只在 backward() 结束时清理。机制见一个算子的反向和 backward 与 Gradients 的「什么时候释放」。

flowchart TB
  op["autodiff 设备上的一个算子"] --> tracked{"有要梯度的输入"}
  tracked -->|是| save["登记步骤 保存张量"]
  tracked -->|否| untracked["登记空步骤 建一张新图"]
  save --> wait["等 backward"]
  untracked --> wait
  wait -->|"一直不 backward"| grow["内存只增不减"]
  wait -->|"backward"| freed["清理"]

实验两个 crate：exp-burn/pb-pit 的 bin mem（Metal，release，带一个统计堆上存活字节数的全局分配器）和 exp-burn/pb-adleak（只开 flex + autodiff，release）。M2 Max、macOS 27、rustc 1.99.0。

现象

有参数的模型

一个两层 MLP（Linear(512, 512) → Relu → Linear(512, 512)）在 metal.autodiff() 上建好，输入 [256, 512]，前向 100 次，每次读回 sum() 然后全部丢掉，不调 backward()。三种摆法：

fn infer(label: &str, model: &Mlp, input_device: &Device, gpu: &Device) {
  let x = Tensor::<2>::random([256, 512], Distribution::Default, input_device);
  let heap0 = HEAP.load(Relaxed);
  for i in 1..=100 {
    let y = model.forward(x.clone());
    let _: f32 = y.sum().into_scalar(); // read the result, drop everything
    if i % 50 == 0 {
      report(&format!("{label}, {i} forwards"), gpu, heap0);
    }
  }
}

// main
infer("A model + input on autodiff", &model, &ad, &gpu);
let valid = model.valid();
infer("B valid(), input on autodiff", &valid, &ad, &gpu);
infer("C valid(), input on plain gpu", &valid, &gpu, &gpu);

report 先 sync()，再打印内存池里在用的字节数和主机堆相对本组开始时的增量：

model built (2 x 1 MiB weights)          gpu in_use    0.0 MiB   host heap   +952.7 KiB
A model + input on autodiff, 50 forwards gpu in_use   52.5 MiB   host heap  +1637.0 KiB
A model + input on autodiff, 100 forwards gpu in_use  102.5 MiB   host heap  +1977.0 KiB
B valid(), input on autodiff, 50 forwards gpu in_use  103.0 MiB   host heap   +501.3 KiB
B valid(), input on autodiff, 100 forwards gpu in_use  103.0 MiB   host heap   +816.5 KiB
C valid(), input on plain gpu, 50 forwards gpu in_use  103.0 MiB   host heap   +155.7 KiB
C valid(), input on plain gpu, 100 forwards gpu in_use  103.0 MiB   host heap   +156.3 KiB

（参数是懒初始化的，建好时还没分配，所以第一行是 0。）

A：显存每次前向涨约 1 MiB，100 次涨到 102.5 MiB，主机堆也每次涨约 6.8 KiB（第 50 到 100 次之间 340 KiB）。参数要梯度，matmul 等算子为求导保存了输入。
B：模型 valid() 过了，显存不再涨（停在 A 留下的 103 MiB，A 的图没有被 backward，那些张量一直没释放）；但输入还在 autodiff 设备上，主机堆每次前向仍涨约 6.3 KiB。
C：输入也放到不带 autodiff 的设备上，两样都不涨。

没有任何 require_grad

把问题缩到最小（pb-adleak）：autodiff 设备上两个不要梯度的 16 元素张量，c = a.clone() + b.clone() 然后立刻 drop(c)，每 10 万次看一次进程 RSS：

fn run(name: &str, a: Tensor<1>, b: Tensor<1>, backward_every_step: bool) {
  let device = a.device();
  print!("{name:<22}");
  for step in 0..=4 {
    if step > 0 {
      for _ in 0..100_000 {
        let c = a.clone() + b.clone();
        drop(c);
      }
      if backward_every_step {
        tiny_backward(&device); // an unrelated 1-element graph, one backward
      }
    }
    print!(" {:>3}k:{:6.1}MB", step * 100, rss_mb());
  }
  println!();
}

flex                     0k:   2.2MB 100k:   2.5MB 200k:   2.5MB 300k:   2.5MB 400k:   2.5MB
flex.autodiff            0k:   2.2MB 100k: 119.6MB 200k: 236.3MB 300k: 353.0MB 400k: 469.7MB
autodiff, a tracked      0k:   2.3MB 100k:  68.4MB 200k: 134.3MB 300k: 229.0MB 400k: 266.1MB
autodiff + backward      0k:   2.2MB 100k: 124.1MB 200k: 124.2MB 300k: 124.3MB 400k: 124.3MB
without_autodiff         0k:   2.2MB 100k:   2.5MB 200k:   2.5MB 300k:   2.5MB 400k:   2.5MB

普通 Flex 不涨；同样的计算放在 Device::flex().autodiff() 上、没有任何张量要梯度，每 10 万次涨约 117 MB，也就是每个算子约 1.2 KB，没有上限。每 10 万次之后对一个毫不相干的小图做一次 backward，内存就停在第一轮的高度（分配器把释放的内存留着复用，RSS 不会回落）。

用计数分配器数得更准（pb-adleak 的 bin count，每组 1 万次）：

flex                 after 10000 dropped adds: +0 bytes in +0 blocks (0 B/op, 0 blocks/op)
                     after an unrelated backward: +1888 bytes, +5 blocks vs before the loop
flex.autodiff        after 10000 dropped adds: +10791872 bytes in +83332 blocks (1079 B/op, 8 blocks/op)
                     after an unrelated backward: -952 bytes, -8 blocks vs before the loop
autodiff, a tracked  after 10000 dropped adds: +5801260 bytes in +51666 blocks (580 B/op, 5 blocks/op)
                     after an unrelated backward: +5801260 bytes, +51666 blocks vs before the loop
                     after a backward through a: -1180 bytes, -12 blocks vs before the loop

未追踪的加法：每次 1079 字节、8 次分配，一个不相干的 backward 之后全部归还。
a 要梯度时：每次 580 字节、5 次分配；不相干的 backward 收不回，必须对经过 a 的图 backward 一次才归还。

出处

要梯度的情形

算子把求导要用的张量存进图里的检查点动作，节点只在 backward 之后被清理，丢掉输出张量没有用。这部分 backward 与 Gradients 的动脑筋已经量过（Metal 上每次前向 2 MiB，一次 backward 降回去）。

不要梯度的情形

所有输入都不要梯度时，算子走 OpsKind::UnTracked，finish 不存状态，但照样登记一个步骤：

    pub fn finish(self, output: FloatTensor<B>) -> AutodiffTensor<B> {
        let output = AutodiffTensor::from_parents(
            output,
            &self.nodes.each_ref().map(|guard| guard.node.clone()),
            self.requirement,
            self.compute_property,
        );
        // ...
        // We register the ops in the graph even if untracked, otherwise memory bound operations
        // that have an untracked parent would not be able to retrieve it
        register_step(
            self.nodes,
            output,
            UntrackedOpsStep::new(ops),
            self.checkpointer_builder,
        )
    }

from_parents 只把要梯度的输入记成父节点，所以这个步骤的父节点列表是空的。登记时 GraphLocator::analyse 看到没有父节点，就以这个节点为起点 new_graph：每个未追踪的算子都在全局的 GraphLocator 里新建一张图（一个 Arc<Graph>，里面一个带锁的 AutodiffServer，它的步骤表、检查点动作表、节点表、前沿集合各存一条），再加上步骤本身的 Box 和节点。这就是上面每次 8 次分配、约 1 KB 的来源。a 要梯度时，加法的父节点是 a，所有步骤登记进 a 那一张图，不新建图，所以每次少一些（580 字节）。

AutodiffTensor 没有 Drop 实现，丢掉 c 不会通知图。释放只有一个入口，GraphMutexClient::backward：反向执行完，先清理这张图里用不到的节点，再调 cleanup_orphaned_entries 扫一遍所有的图，对「没有任何节点还被引用」的图（maybe_useful() 为假）调 free_unused_roots。未追踪加法那一张张小图都满足条件，所以任何一次 backward 都能把它们扫掉；a 那张图里 a 还被引用着，maybe_useful() 为真，扫描跳过它，只有经过它的 backward 才会清理里面的死节点。

一个只做推理、从不 backward 的进程里，这两条路都不会走到。

这算不算上游的问题

在默认的不做检查点的策略（Device::autodiff() 的 Disabled）下，这些未追踪步骤没有用处：它们的 step 是空函数，又因为 from_parents 只登记要梯度的父节点，任何 backward 的遍历都走不到它们；注释里给的理由（内存密集算子补算时要找未追踪的父节点）只和检查点有关。代价却是每个算子约 1 KB、只有 backward 能回收。我认为这是 burn 0.22.1 的一个上游问题，最小复现就是上面的 pb-adleak：

[dependencies]
burn = { version = "=0.22.1", default-features = false, features = ["std", "flex", "autodiff"] }

let ad = Device::flex().autodiff();
let a = Tensor::<1>::from_floats([1.0f32; 16], &ad);
let b = Tensor::<1>::from_floats([2.0f32; 16], &ad);
for _ in 0..1_000_000 {
  drop(a.clone() + b.clone()); // RSS grows ~1.2 KB per iteration, never freed without a backward
}

📝NOTE

冷知识：未追踪节点的清理一直挂在 backward 上

#3957「Fix autodiff graph cleanup」（首次进入 v0.20.0-pre.1）给客户端加了 register_untracked：未追踪的步骤照常登记（源码注释「might be needed by tracked children, required for checkpointing」），同时记进一个 untracked_nodes 集合，「mark this node for cleanup after backward」，在每次 backward 结束时删掉。三周后 #4039「Cleanup autodiff unused roots」（首次进入 v0.20.0-pre.3）删掉了 register_untracked 和那个集合，换成现在的做法：backward 之后扫描所有图，free_unused_roots 清掉没人引用的根。两个版本的清理时机都是「backward 结束时」，没有 backward 的进程里未追踪的记录从来没有被回收过。两个提交都只有标题，没有说明。

修法

推理不要放在 autodiff 设备上。模型 valid()，输入也建在不带 autodiff 的设备上（上面的 C 组）；已经在 autodiff 设备上的张量用 without_autodiff() 摘掉（pb-adleak 的最后一行，平的）：

let valid = model.valid();            // 参数去掉 autodiff，也关掉 Dropout
let x = x.without_autodiff();          // 或者一开始就建在 Device::metal(..) 上
let y = valid.forward(x);

Device 的相等比较不看 autodiff（设备），所以「训练设备」和「推理设备」写成同一个变量时很难察觉：推理代码里单独建一个 Device::metal(..) / Device::flex()，不要复用训练用的 .autodiff() 那个。

训练循环里每步都有 backward，这些记录每步都会被清掉，不会累积；会累积的是训练循环之外的部分：验证、生成样本、导出前的冒烟测试。

怎么确认没中招

推理入口处断言 !x.device().is_autodiff()，模型参数 !model.weight.val().is_autodiff()。
长时间运行的推理进程，看一眼 RSS 和 device.memory_pool_usage() 的走势：每批都涨就是图在累积。
用一个计数的全局分配器（本页 count 的写法）测「一次前向之后堆上多了多少」，应当接近 0。

小结

autodiff 设备上的每个浮点算子都会登记记录，只有 backward() 结束时清理；只做前向的进程里，它们一直累积。
参数要梯度时累积的是保存的张量（显存）；完全不要梯度时，每个算子仍在主机上留约 1 KB（一张新图），这是 0.22.1 的上游问题。
推理用 valid() 的模型和不带 autodiff 的设备，或对输入 without_autodiff()。