梯度检查点

主线位置：求导主线上的一条支线。前向时每个算子要把反向需要的输入交给检查点（见一个算子的反向），策略决定交出去的是张量本身还是「怎么把它重新算出来」；backward 时（见 backward 与 Gradients）反向步骤通过 Checkpointer 取回这些输入，需要时现场补算。梯度检查点（gradient checkpointing）就是用这点额外的计算换前向之后少占的内存。

flowchart TB
  ck["前向 checkpoint 输入"] --> prop{"输入节点的属性"}
  prop -->|"计算密集或未定"| comp["Computed 存克隆"]
  prop -->|访存密集| rec["Recompute 存重算方法"]
  comp --> build["反向前 build 统计次数"]
  rec --> build
  build --> ret["retrieve_node_output"]
  ret --> topo["拓扑排序到已算节点"]
  topo --> rf["逐个 retro_forward 补算"]
  rf --> getst["get_state 次数减一"]

两种策略

策略是 Autodiff<B, C> 的第二个类型参数，实现 CheckpointStrategy 的两个方法。0.22.1 只有两个实现。NoCheckpointing 是默认值：

impl CheckpointStrategy for NoCheckpointing {
    /// An operation marked as memory bound is actually compute bound.
    fn compute_property<R: RetroForward>(_retro_forward: R) -> ComputingProperty {
        ComputingProperty::ComputeBound
    }

    /// An operation marked as memory bound is actually compute bound.
    /// It's therefore useless to checkpoint the parents
    fn checkpoint_parents<'a, B2, A>(
        _parents: A,
        _builder: &mut CheckpointerBuilder,
    ) -> Result<(), CheckpointingError>
    where
        B2: Backend,
        A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
    {
        // Nothing to do here
        Ok(())
    }
}

BalancedCheckpointing 照算子自己的标记办：

impl CheckpointStrategy for BalancedCheckpointing {
    /// An operation marked as memory bound is memory bound.
    /// When memory bound, an operation needs to save its RetroForward
    fn compute_property<R: RetroForward>(retro_forward: R) -> ComputingProperty {
        ComputingProperty::MemoryBound {
            retro_forward: Arc::new(retro_forward),
        }
    }

    // ...
    fn checkpoint_parents<'a, B2, A>(
        parents: A,
        builder: &mut CheckpointerBuilder,
    ) -> Result<(), CheckpointingError>
    where
        B2: Backend,
        A: IntoIterator<Item = &'a AutodiffTensor<B2>>,
    {
        let mut can_checkpoint = true;

        for tensor in parents.into_iter() {
            if let crate::graph::Requirement::None = tensor.node.requirement {
                can_checkpoint = false;
            } else {
                builder.checkpoint(tensor, ActionType::Backup);
            }
        }

        if !can_checkpoint {
            *builder = CheckpointerBuilder::default();
            return Err(CheckpointingError::UntrackedParent);
        }

        Ok(())
    }
}

这两个方法就是上一页类型状态链里 .retro_forward(..) 和 .parents(..) 两步调用的东西：

compute_property 决定输出节点的 ComputingProperty。NoCheckpointing 把每个标了 memory_bound() 的算子都改回 ComputeBound，算子写的 RetroForward 被直接丢弃。
checkpoint_parents 只在 Balanced 下做事：一个输出将来可能要重算，重算就要用到它的输入，所以把每个输入也「备份」一份（ActionType::Backup）。只要有一个输入是常量（Requirement::None），就清空备份、返回错误，.parents(..) 收到错误就把这个算子退回 ComputeBound。源码给的理由是常量输入的需求事先不知道，没法稳妥地备份（CheckpointingError::UntrackedParent 的注释）。

哪些算子标了 memory_bound()：ops/tensor.rs 里是逐元素的算术和数学函数（add、sub、mul、div 及其标量版本、exp、log、sqrt、powf、三角和双曲函数、abs、round 等）和只改视图或搬数据的算子（reshape、permute、swap_dims、expand、slice、flip、select、repeat_dim 等）；ops/activation.rs 里是 relu、gelu、sigmoid、log_sigmoid。matmul、各种归约（sum、mean、max_dim……）、gather / scatter、cumsum、sort、cast、mask_where 是 compute_bound()；ops/module.rs 里的卷积、池化等一个 memory_bound() 也没有。

用户这一侧只有一个开关：Device::gradient_checkpointing 把 autodiff 设备的策略设成 Balanced，之后每个算子在分派时被 with_autodiff_backend 换成 Autodiff<B, BalancedCheckpointing>（见 Autodiff 装饰器）。对没开 autodiff 的设备调用会 panic：「Gradient checkpointing requires autodiff; use Device::autodiff().gradient_checkpointing()」。

前向：两种动作

前向时，prep.checkpoint(&x)（显式，ActionType::Explicit）和 checkpoint_parents（备份，ActionType::Backup）都调 CheckpointerBuilder::checkpoint，按被保存的那个节点的属性记一种动作：

        match &tensor.node.properties {
            ComputingProperty::ComputeBound | ComputingProperty::Ambiguous => {
                action_list.push(CheckpointingAction::Computed {
                    node_id: tensor.node.id,
                    state_content: Box::new(tensor.primitive.clone()),
                })
            }
            ComputingProperty::MemoryBound { retro_forward } => {
                action_list.push(CheckpointingAction::Recompute {
                    node_id: tensor.node.id,
                    retro_forward: retro_forward.clone(),
                })
            }
        }

Computed 存一份张量句柄的克隆：克隆不复制数据，但只要它活着，那块内存就不能回收。Recompute 只存一个 Arc<dyn RetroForward>，比如上一页的 RetroMul，里面只有两个 NodeId。叶子张量（AutodiffTensor::new 建的）属性是 Ambiguous，按 Computed 处理；它本来就在内存里，克隆不额外占地方。

这些动作挂在算子自己的 CheckpointerBuilder 上，随步骤一起登记进图（AutodiffServer 的 actions_builder）。所以检查点省下的是「图里挂着的张量」：前向结束时，用户已经不再持有的中间结果，NoCheckpointing 下被克隆拉住，Balanced 下如果是访存密集算子的输出，就没人拉住，随即释放。

backward 前：统计每个状态要用几次

build_tape 收集完要执行的步骤后，把它们的 CheckpointerBuilder 合并，调 build 建出 Checkpointer：

find_stop_nodes：所有 Computed 动作的节点，重算时追溯到这里为止。
build_n_required_map：只看显式动作。Computed 的节点计数加 1；Recompute 的节点计数加 1，并且沿着父节点往上递归，直到碰到停止节点，沿途每个节点都加 1（重算它要用到它们）。
insert_checkpoints：对计数表里的每个节点，先在显式动作里找，找不到再去备份动作里找，变成 State::Computed { state_content, n_required } 或 State::Recompute { n_required }。两处都找不到就 panic「Node {:?} is needed but never checkpointed」。

计数表里没有出现的备份动作在 build 结束时随 self 一起丢掉，它们拉住的张量也就释放了。ActionType::Backup 的注释说的就是这件事：「If no explicit action ever asks for the child's output, then the backup output will go out of scope when the checkpointer is built.」

backward 中：取一个状态

反向步骤调 retrieve_node_output：

    pub fn retrieve_node_output<T>(&mut self, node_id: NodeId) -> T
    where
        T: Clone + Send + 'static,
    {
        let sorted = self.topological_sort(node_id);
        let num_nodes = sorted.len();
        log_autodiff(AutodiffLogLevel::Basic, move || {
            format!("retrieve_node_output {node_id:?}: {num_nodes} node(s) to compute")
        });

        sorted.into_iter().for_each(|node| {
            log_autodiff(AutodiffLogLevel::Full, move || {
                format!("execute_retro_forward {node:?}")
            });
            self.retro_forwards
                .execute_retro_forward(node, &mut self.backward_states)
        });

        self.backward_states.get_state::<T>(&node_id)
    }

topological_sort 从要取的节点往上走：Recompute 的节点先访问它所有父节点再把自己放进结果，Computed 的节点直接放进结果、不再往上，所以得到的序列里父节点总在子节点前面。然后逐个执行：execute_retro_forward 只对仍是 Recompute 的节点调它的 RetroForward::forward（之后从表里删掉，每个重算方法最多用一次），已经算好的跳过。retro_binary! 生成的 forward 是这样的：

            fn forward(&self, states: &mut BackwardStates, out_node: NodeId) {
                let lhs = states.get_state::<B::FloatTensorPrimitive>(&self.lhs_id);
                let rhs = states.get_state::<B::FloatTensorPrimitive>(&self.rhs_id);
                let out = $ops(lhs, rhs);
                states.save(out_node, out)
            }

BackwardStates::get_state 每取一次把 n_required 减 1，减到 0 就把状态从表里移走、把张量本身交出去，不再保留；save 把重算结果存成 Computed，沿用原来的计数。反向进行到哪里，用完的状态就释放到哪里。

📝NOTE

冷知识：补算前的排序曾经是平方级的

PR #5679（「perf(autodiff): avoid redundant traversal in checkpoint topological sort」，首次进入 v0.22.0-pre.4，正式版 0.22.0）之前，topological_sort 是递归返回 Vec 的写法：对 Recompute 节点的每个父节点各排一遍序，再逐个用 sorted.contains(&ps) 去重后拼起来。同一个祖先被几条分支共享时，它的整棵子树会在每条分支上重新展开一次，去重又是线性查找。提交说明写的是用 HashSet 记录访问过的节点，「expand each ancestor once instead of per-branch」；现在函数上方的注释把后果写得更具体：不这样做，祖先链变长时一次取值会退化成平方级甚至更糟，例子是被多条路径反复访问的循环链。提交里附带的性能复现测试在合并前又删掉了。

实验一：哪些节点被补算

autodiff 有自己的日志开关，但和 fusion 不同，没有 BURN_AUTODIFF_LOG 这样的环境变量（BurnConfig::override_from_env 只认 BURN_FUSION_LOG、BURN_FUSION_MAX_EXPLORATIONS、BURN_REMOTE_LOG）。要从当前目录往上找到的 burn.toml 里配置：

[autodiff.logger]
level = "full"
stdout = true

0.22.1 里全部的 autodiff 日志就是上面 retrieve_node_output 里的两行：basic 级打「取哪个节点、要经过几个节点」，full 级再打每个经过的节点。AutodiffLogLevel 的文档说 basic 会记「backward graph size and the checkpoint strategy applied per forward pass」，medium 会记每个被检查点或重算的张量，源码里没有这些调用。

网络是 ad-exp 的 ckpt：每层 h = silu(h.matmul(w))，silu(x) 在 burn 里就是 x * sigmoid(x)，所以每层是一个计算密集的 matmul 加两个访存密集的算子。权重 require_grad，输入 x 不要梯度：

  let weights: Vec<Tensor<2>> = (0..layers)
    .map(|_| Tensor::<2>::random([n, n], Distribution::Normal(0.0, scale), &device).require_grad())
    .collect();
  let x = Tensor::<2>::random([n, n], Distribution::Default, &device);
  // ...
    let mut h = x.clone();
    for w in &weights {
      h = silu(h.matmul(w.clone()));
    }
    let loss = h.sum();
    // ...
    let grads = loss.backward();

在 Device::flex() 上取两层、n = 64，节点 id 依次是：w0 = 0、w1 = 1、x = 2，第一层 matmul / sigmoid / mul = 3 / 4 / 5，第二层 = 6 / 7 / 8，sum = 9。NoCheckpointing 的日志（debug，在放了 burn.toml 的目录里运行）：

device = Device<Autodiff { device: Flex(Cpu), checkpointing: Disabled }>, n = 64, layers = 2
...
retrieve_node_output NodeId { value: 6 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 6 }
retrieve_node_output NodeId { value: 7 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 7 }
retrieve_node_output NodeId { value: 6 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 6 }
retrieve_node_output NodeId { value: 5 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 5 }
retrieve_node_output NodeId { value: 1 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 1 }
retrieve_node_output NodeId { value: 3 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 3 }
retrieve_node_output NodeId { value: 4 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 4 }
retrieve_node_output NodeId { value: 3 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 3 }
retrieve_node_output NodeId { value: 2 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 2 }

顺序就是反向执行的顺序：第二层的 mul（8）取 m2（6）和 s2（7），sigmoid（7）取它的输入 6，matmul（6）取 h1（5）和 w1（1），第一层同理，最后第一层的 matmul（3）只取 x（2）：x 不要梯度，w0 的梯度要用 x，而 x 的梯度不用算，所以 w0 从来没被保存。所有状态都是前向存好的，每次都是「1 node(s)」，execute_retro_forward 那行只是日志，对 Computed 的节点什么也不做。8 号节点（第二层输出）没有出现，sum 的反向只要形状。

❗IMPORTANT

动脑筋：同一个两层网络换成 Device::flex().autodiff().gradient_checkpointing()，日志里 node(s) to compute 大于 1 的有几行？一共补算了几个节点、是哪几个？

ad-exp 的 ckpt balanced flex 64 2 1，debug，平台 M2 Max、macOS 27。

实验二：内存和时间

把网络放大到 8 层、n = 1024（每个张量 1024 × 1024 × 4 字节 = 4 MiB），在 Metal 上读前向结束后内存池里在用的字节数，并计时前向和反向（计时前后 device.sync()）。为了不让 autotune-checks 这个调试 feature 干扰计时，这组实验放在单独的 ad-time crate 里，feature 只开 flex、metal、autodiff、fusion、autotune，release 构建；每种策略跑 31 次，第 0 次包含内核编译和 autotune，取后 30 次的中位数。两种策略交替各跑两轮，结果一致，下面是第一轮：

device = Device<Autodiff { device: Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })), checkpointing: Disabled }>, n = 1024, layers = 8
params + input          in_use =    36.0 MiB
rep 0: after forward in_use =   128.0 MiB, forward 332.22ms, backward    1.08s, sum(grad w0) = -26673.1992
rep 1: after forward in_use =   128.0 MiB, forward   6.37ms, backward  11.73ms, sum(grad w0) = -26673.1973
rep 2: after forward in_use =   128.0 MiB, forward   4.98ms, backward   9.23ms, sum(grad w0) = -26673.2031
reps 1..30: median forward 4.41ms, median backward 8.33ms
device = Device<Autodiff { device: Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })), checkpointing: Balanced }>, n = 1024, layers = 8
params + input          in_use =    36.0 MiB
rep 0: after forward in_use =    68.0 MiB, forward 180.55ms, backward 288.32ms, sum(grad w0) = -26673.1992
rep 1: after forward in_use =    68.0 MiB, forward   6.46ms, backward  11.63ms, sum(grad w0) = -26673.1992
rep 2: after forward in_use =    68.0 MiB, forward   5.05ms, backward   9.56ms, sum(grad w0) = -26673.1953
reps 1..30: median forward 4.39ms, median backward 8.58ms

同一个网络在 Device::flex() 上（release，每种策略 11 次，取后 10 次的中位数，两轮）：

reps 1..10: median forward 80.31ms, median backward 114.05ms
reps 1..10: median forward 82.59ms, median backward 161.43ms
reps 1..10: median forward 79.50ms, median backward 112.49ms
reps 1..10: median forward 78.81ms, median backward 159.02ms

（第 1、3 行是 NoCheckpointing，第 2、4 行是 Balanced。）

内存：参数和输入是 9 个张量、36 MiB。前向之后，NoCheckpointing 多出 92 MiB，也就是 23 个张量：每层的 matmul 输出 m 和 sigmoid 输出 s（各 8 个），加上前 7 层的 silu 输出 h（被下一层 matmul 保存；最后一层的只给 sum，不保存）。Balanced 只多出 32 MiB，8 个张量，正好是每层的 m：s 和 h 是访存密集算子的输出，只记了重算方法。
时间：Metal 上两种策略的反向中位数都在 8.3–8.6 ms，看不出补算的代价；Flex 上 Balanced 的反向从约 113 ms 变成约 160 ms，多了四成。两种策略反向时执行的梯度算子相同，多出来的是补算：每次反向补算 8 个 sigmoid 和 7 个 mul，每个 100 万元素（个数见动脑筋答案）。前向两边一样，策略不改变前向算什么，只改变留下什么。
梯度：两种策略的 sum(grad w0) 都落在 -26673.195 到 -26673.206 之间（四轮所有输出行），同一种策略的不同轮次之间也是这样浮动，是 GPU 上浮点累加顺序的差别。
这里只测了前向结束时的占用，没有测反向过程中的峰值；GPU 时间只对这台机器（M2 Max、macOS 27、rustc 1.99.0）、这个形状有效。Flex 没有内存池，memory_pool_usage() 返回 None。
第 0 次的时间里有内核编译和 autotune。cubecl 把编译和调优结果写进 target/environment/default.db（相对于从当前目录往上找到的最外层 Cargo.toml 所在目录），所以同一个目录里第二个进程的第 0 次快得多，表里 Balanced 的第 0 次就是这样；autotune 的细节见 autotune。

📝NOTE

没有蠢问题

问：教科书里的梯度检查点是每隔几层存一个激活、反向时把整段重跑一遍，能把内存从 O(n) 压到 O(√n)。burn 的 Balanced 每层都存了 matmul 的输出，这算哪门子检查点？

答：它不是分段的那种。Balanced 是逐算子决定的：计算密集的算子（matmul、卷积、归约）输出总是存，访存密集的（逐元素、改视图）只记怎么重算，重算最多往上追到最近的一个已存节点。好处是不用用户划分段落，也不会重跑 matmul 这类昂贵的算子；代价是省下的只有逐元素算子的输出。在 silu 这种「一个 matmul 后面跟几个逐元素算子」的网络里，这部分恰好是大头，上面的实验省了约三分之二。0.22.1 里 CheckpointStrategy 只有这两个实现，名字叫 Balanced 也说明了它的定位。

动脑筋答案

device = Device<Autodiff { device: Flex(Cpu), checkpointing: Balanced }>, n = 64, layers = 2
...
retrieve_node_output NodeId { value: 6 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 6 }
retrieve_node_output NodeId { value: 7 }: 2 node(s) to compute
execute_retro_forward NodeId { value: 6 }
execute_retro_forward NodeId { value: 7 }
retrieve_node_output NodeId { value: 6 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 6 }
retrieve_node_output NodeId { value: 5 }: 3 node(s) to compute
execute_retro_forward NodeId { value: 3 }
execute_retro_forward NodeId { value: 4 }
execute_retro_forward NodeId { value: 5 }
retrieve_node_output NodeId { value: 1 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 1 }
retrieve_node_output NodeId { value: 3 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 3 }
retrieve_node_output NodeId { value: 4 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 4 }
retrieve_node_output NodeId { value: 3 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 3 }
retrieve_node_output NodeId { value: 2 }: 1 node(s) to compute
execute_retro_forward NodeId { value: 2 }

两行大于 1，补算了三个节点：

取 7（第二层的 sigmoid 输出 s2）：排序得到 [6, 7]，6 是存好的 m2，只补算 7。
取 5（第一层的 silu 输出 h1，第二层 matmul 要用它算 w1 的梯度）：排序得到 [3, 4, 5]，3 是存好的 m1，补算 4（s1）和 5（h1）。
随后第一层的 mul 取 4 时只有「1 node(s)」：4 在上一步被补算出来后，因为还有人要用（n_required 没减到 0），以 Computed 的形式留在表里，不用再算一次。
取的顺序、每次取的节点和 NoCheckpointing 完全一样，区别只在每次要经过几个节点。推到 8 层：每层的 sigmoid 输出都要补算，silu 输出除了最后一层都要补算，一共 15 个逐元素算子，和实验二里 Flex 上多出的反向时间对得上。

小结

检查点策略是 Autodiff 的类型参数，Device::gradient_checkpointing() 选 BalancedCheckpointing，默认是 NoCheckpointing；两者的区别只在 compute_property 和 checkpoint_parents 两个方法里。
NoCheckpointing 把所有算子当计算密集，反向要用的输入都以张量克隆的形式挂在图上；Balanced 对逐元素和改视图的算子只记 RetroForward，有常量输入的算子退回计算密集。
反向前 CheckpointerBuilder::build 统计每个状态要用几次，用不到的备份随即释放；反向中 retrieve_node_output 拓扑排序到最近的已存节点，按序补算，get_state 用一次减一次、减到 0 就释放。
8 层、1024 宽的 silu 网络：前向后保存的状态 92 MiB 对 32 MiB；Metal 上反向时间看不出差别，Flex 上反向慢了约四成。
autodiff 日志只能用 burn.toml 打开，0.22.1 里只记录检查点取值这两行。