Rust in Everything
基于 burn 0.22.1

Autodiff 装饰器

主线位置：求导主线的起点，也是算子主线上可选的一环。算子主线里 Dispatch 按设备 match 之后，如果设备是 .autodiff() 的，先进 Autodiff<B, C> 记一个节点，再调内层后端算；本页讲这层装饰器长什么样、张量上多带了什么、运行期的标记怎么变成类型参数。match 本身怎么生成见 Dispatch，算子里登记了什么见一个算子的反向。

flowchart TB
  devn["Device flex autodiff"] --> dd["DispatchDevice Autodiff"]
  dd --> tens["张量带 Enabled 策略"]
  tens --> merge{"合并各输入 context"}
  merge -->|Disabled| plain["Flex 直接算"]
  merge -->|"Enabled Disabled"| nock["Autodiff Flex NoCheckpointing"]
  merge -->|"Enabled Balanced"| bal["Autodiff Flex Balanced"]
  nock --> op["Autodiff 的 FloatTensorOps"]
  bal --> op
  op --> fwd["调 Flex 前向"]
  op --> step["登记 Step 到图"]

本页的实验 crate 是 ad-exp（从模板 base 复制，burn =0.22.1，feature 与模板相同），在 Device::flex() 上跑，debug 构建，平台 M2 Max、macOS 27、rustc 1.99.0。

Autodiff<B, C>：只换浮点张量

Autodiff 是一个只有两个 PhantomData 的空结构体，第二个类型参数默认是 NoCheckpointing。它在 BackendTypes 里只改了一样东西：

pub struct Autodiff<B, C = NoCheckpointing> {
    _b: PhantomData<B>,
    _checkpoint_strategy: PhantomData<C>,
}

impl<B: BackendTypes, C: CheckpointStrategy> BackendTypes for Autodiff<B, C> {
    type Device = B::Device;

    type FloatTensorPrimitive = AutodiffTensor<B>;

    type IntTensorPrimitive = B::IntTensorPrimitive;

    type BoolTensorPrimitive = B::BoolTensorPrimitive;

    type QuantizedTensorPrimitive = B::QuantizedTensorPrimitive;

    // A replayed graph would skip re-recording the autodiff tape, so capture
    // is not supported under autodiff.
    type GraphPrimitive = burn_backend::GraphUnsupported;
}

设备类型照搬内层，整数、布尔、量化张量也照搬：只有浮点张量能求导，其他种类的张量过这一层时原样转手。Backend 的实现几乎全是转发（sync、seed、memory_*、supports_dtype 都直接调 B::），只有 ad_enabled 返回 true，name 在内层名字外面包一层 autodiff<...>。AutodiffBackend 这个 trait 多出 backward、grad、grad_remove、inner / from_inner 几个方法，都是一行，转给 AutodiffTensor 自己的方法。真正的工作分在两处：每个浮点算子的 FloatTensorOps 实现（ops/tensor.rs，下一页），和 backward() 背后的图（runtime/，第三页）。

后端 trait 的整体结构（BackendTypes、按种类拆开的 ops trait、AutodiffBackend）见 Backend trait。

AutodiffTensor：原语、节点、引用计数

AutodiffTensor<B> 有三个字段：

#[derive(Debug, Clone)]
pub struct AutodiffTensor<B: BackendTypes> {
    pub(crate) primitive: B::FloatTensorPrimitive,
    pub(crate) node: NodeRef,
    pub(crate) rc: NodeRefCount,
}
  • primitive 是内层后端的浮点张量，前向的数值全在这里，Autodiff 不复制数据。
  • node 是 Arc<Node>，图的元数据。
  • rc 是 NodeRefCount，也就是 Arc<NodeId>。它不携带信息（id 在 node 里也有），存在的意义是它的强引用计数：图里登记步骤时会存一份 rc，第三页的内存管理靠 Arc::strong_count 判断「用户手里还有没有指向这个节点的张量」。#[derive(Clone)] 让张量的 clone 只多一份 Arc。

Node 的字段：

字段类型含义
parentsVec<Parent>需要梯度的父节点 id，Parent.is_leaf 记父节点是不是叶子
orderusize深度：叶子是 0，算子输出是父节点最大深度加 1
idNodeId全局唯一的 u64
requirementRequirement要不要梯度，见下一节
propertiesComputingPropertyComputeBound / MemoryBound { retro_forward } / Ambiguous，给检查点用
clientAutodiffClientImpl登记步骤、发起 backward 的入口，0.22.1 里是一个零大小的 GraphMutexClient
distributed_paramsOption<DistributedParams>多设备训练时标记参数，本卷不展开

NodeId::new 是一个全局 static COUNTER: AtomicU64 的 fetch_add，所以 id 在进程内单调递增，和设备、后端都无关。后几页实验里打印出来的 NodeId(5)、NodeId { value: 6 } 就是这个计数器的值。

AutodiffTensor 对外只给出 node()、id()、primitive()、into_primitive()、into_parts() 这几个方法，字段是 pub(crate)。node() 返回的不是 NodeRef，而是同时持有 node 和 rc 的 NodeGuard，原因放在下一页的冷知识里。

📝NOTE
没有蠢问题

问： Autodiff<B, C> 有检查点策略 C，AutodiffTensor<B> 却没有。同一个张量换了策略，类型不就对不上了吗？

答： 策略只在前向登记时起作用：算子用 C::compute_property、C::checkpoint_parents 决定存输出还是存重算方法（见梯度检查点）。登记完以后，图里的步骤是 Box<dyn Step>，backward()、grad() 都不需要知道 C；Gradients 也是一个按 u64 存类型擦除张量的容器。把 C 留在张量类型外面，Dispatch 就能在 BackendTensor::Autodiff(FloatTensor<Autodiff<B>>) 里只存一种张量（这里的 Autodiff<B> 取了默认的 C，但 FloatTensor<Autodiff<B, C>> 对任何 C 都是同一个 AutodiffTensor<B>），策略作为运行期的值另外存，每个算子再临时选 C。

Requirement：叶子、中间结果、常量

Requirement 有三个值：

pub enum Requirement {
    /// Operations that require gradients.
    Grad,
    /// Operations that require gradients only for backprop.
    GradInBackward,
    /// Operations that don't need gradients, therefore not to be included in the graph.
    None,
}
  • Grad：调过 require_grad() 的叶子。backward 之后它的梯度留在 Gradients 里给用户取。
  • GradInBackward：至少有一个输入要梯度的算子输出。backward 途中要用到它的梯度，用完就删（第三页）。
  • None：常量。新建的张量（from_data、zeros、random……都走 AutodiffTensor::new）一开始都是 None，输入全是 None 的算子输出也是 None。

算子输出的要求由 Requirement::from_nodes 推出来，规则只有一条（infer）：所有输入都是 None 才是 None，否则是 GradInBackward。所以算子的输出永远不会是 Grad。张量 API 上的三个查询对应这几个值：

方法看什么为真的条件
is_autodiff()设备张量带着 autodiff 标记，和要不要梯度无关
is_tracked()requirement不是 None
is_require_grad()requirement是 Grad

require_grad() 的实现（AutodiffTensor::require_grad）不改原来的节点（Arc<Node> 是共享的，改不了），而是用同一个 id 新建一个 Grad 节点，再登记一个什么也不做的根步骤 RootStep：

    pub fn require_grad(mut self) -> Self {
        match self.node.requirement {
            Requirement::Grad => self,
            Requirement::GradInBackward => {
                panic!("Can't convert a non leaf tensor into a tracked tensor")
            }
            Requirement::None => {
                self.node = Node::new(
                    vec![],
                    0,
                    self.node.id,
                    Requirement::Grad,
                    self.node.properties.clone(),
                    self.node.client.clone(),
                    self.node.distributed_params.clone(),
                )
                .into();
                self.register_root()
            }
        }
    }

GradInBackward 那一支说明：中间结果不能改成要梯度。burn 没有 PyTorch 的 retain_grad，想要一个中间量的梯度，只能把它 detach() 成新的叶子再 require_grad()，那样梯度也就不再往前传了。实验里的几种 panic，原文如下（ad-exp 的 backend，debug）：

thread 'main' (3144670) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-autodiff-0.22.1/src/tensor.rs:190:17:
Can't convert a non leaf tensor into a tracked tensor
...
thread 'main' (3144670) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-tensor-0.22.1/src/tensor/api/float.rs:1185:13:
Tensor::require_grad requires autodiff; call Tensor::autodiff first
...
thread 'main' (3144670) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-tensor-0.22.1/src/tensor/api/autodiff.rs:75:9:
Tensor::backward requires a tracked autodiff tensor; call Tensor::autodiff().require_grad() on the source leaf before computing the output

第一条来自 burn-autodiff 本身（(w * 2.0).require_grad()）；第二条是 burn-tensor 在分派之前用 assert! 拦下的（在 Device::flex() 上建的张量调 require_grad()）；第三条是对一个全是常量算出来的张量调 backward()，同样在 burn-tensor 里拦下。

从 Device::autodiff() 到 Autodiff<Flex, C>

0.22 的用户代码里不出现 Autodiff 这个类型。Device::autodiff 把设备包成 DispatchDevice::Autodiff，里面是 AutodiffDevice：

pub struct AutodiffDevice {
    pub(crate) inner: Box<DispatchDevice>,
    pub(crate) checkpointing: GradientCheckpointingStrategy,
}

checkpointing 是 GradientCheckpointingStrategy，Balanced 或默认的 Disabled，Device::gradient_checkpointing 把它改成 Balanced。设备相等只比较 inner，autodiff 标记和策略都不参与；Backend::name 也不带策略：

flex     Device<Flex(Cpu)>  name=dispatch<flex>
ad       Device<Autodiff { device: Flex(Cpu), checkpointing: Disabled }>  name=dispatch<autodiff<flex>>
ad_ckpt  Device<Autodiff { device: Flex(Cpu), checkpointing: Balanced }>  name=dispatch<autodiff<flex>>
flex == ad: true, ad == ad_ckpt: true

在这样的设备上建的张量，Dispatch 层的表示是 DispatchTensor 的两个字段：kind 指向具体后端的张量，autodiff 是 DispatchAutodiffContext，Disabled 或 Enabled(策略)。浮点张量和整数、布尔张量的样子不一样：

张量kindautodiff
Device::flex() 上的浮点Flex(BackendTensor::Float(FlexTensor))Disabled
.autodiff() 设备上的浮点Autodiff(Box<Flex(BackendTensor::Autodiff(AutodiffTensor<Flex>))>)Enabled(策略)
.autodiff() 设备上的整数 / 布尔Flex(BackendTensor::Int(..)) / Bool(..)Enabled(策略)

整数和布尔不会被求导，所以不套 AutodiffTensor，只靠 autodiff 字段记住自己属于哪种设备，device() 据此报告 Autodiff { .. }。

每个算子开头把所有输入的 context 用 merge 合并：Disabled 和任何值合并得到对方，两个 Enabled 必须策略相同，否则 assert_eq! 失败。合并的结果决定这个算子走具体后端，还是走 Autodiff<B, C>、C 取哪一个。#[backend_dispatch] 为 Dispatch::float_mul 生成的代码（在实验目录里对 burn-dispatch 0.22.1 只开 flex、autodiff 两个 feature 跑 cargo expand，只留结构）：

            fn float_mul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
                let __burn_selected_context = lhs.autodiff;
                let mut __ad_ctx = crate::DispatchAutodiffContext::Disabled;
                __ad_ctx = __ad_ctx.merge(lhs.autodiff);
                __ad_ctx = __ad_ctx.merge(rhs.autodiff);
                match lhs.kind {
                    crate::DispatchTensorKind::Flex(__burn_selected) => {
                        match __ad_ctx {
                            crate::DispatchAutodiffContext::Disabled => {
                                type B = crate::backends::Flex;
                                // ... 取出 rhs 的 Flex 原语，B::float_mul，包回 Flex(Float)
                            }
                            crate::DispatchAutodiffContext::Enabled(__strategy) => {
                                match __strategy {
                                    crate::GradientCheckpointingStrategy::Balanced => {
                                        type B = crate::backends::Autodiff<
                                            crate::backends::Flex,
                                            burn_autodiff::checkpoint::strategy::BalancedCheckpointing,
                                        >;
                                        let lhs = <B as burn_backend::AutodiffBackend>::from_inner(
                                            __burn_selected.float(),
                                        );
                                        // ... rhs 同理，B::float_mul，包成 Autodiff(Flex(Autodiff(..)))
                                    }
                                    crate::GradientCheckpointingStrategy::Disabled => {
                                        // ... 同上，C = NoCheckpointing
                                    }
                                }
                            }
                        }
                    }
                    crate::DispatchTensorKind::Autodiff(__inner) => {
                        match *__inner {
                            crate::DispatchTensorKind::Flex(__burn_selected) => {
                                match __ad_ctx {
                                    crate::DispatchAutodiffContext::Disabled => {
                                        // panic: "an autodiff float primitive must have an enabled autodiff context"
                                    }
                                    crate::DispatchAutodiffContext::Enabled(__strategy) => {
                                        // ... 和上面一样按 __strategy 选 C，lhs 取 __burn_selected.autodiff()
                                    }
                                }
                            }
                            // ...
                        }
                    }
                    // ...
                }
            }

几件事从这里能直接读出来：

  • 运行期的策略在这里变成类型参数：每个分支里 type B = Autodiff<Flex, BalancedCheckpointing> 或 Autodiff<Flex, NoCheckpointing>，然后调 B::float_mul。两种策略的 Autodiff 都被单态化，编进同一个二进制。做这件事的是过程宏里的 routing::with_autodiff_backend；burn-dispatch 里同名的声明宏 with_autodiff_backend! 生成同样的 match，用在跨设备搬运、量化和分布式这几处手写的分派里。
  • 一个没有 autodiff 标记的浮点张量（kind 是 Flex(Float)）和 autodiff 张量一起参与运算时，合并出来是 Enabled，它被 from_inner 包成一个新的 AutodiffTensor。from_inner 就是 AutodiffTensor::new，Requirement::None，所以它在这次运算里是常量。
  • 输出包成 DispatchTensorKind::Autodiff(..)，autodiff 字段是合并后的 context，所以结果仍在 autodiff 设备上。
📝NOTE
冷知识：Dispatch 一开始只有不带检查点的 autodiff

burn-dispatch 是 0.21 开发中由 PR #4508（「[Feat] Global backend Dispatch」，首次进入 v0.21.0-pre.2）加进来的。刚出现时，autodiff 分支写死成 type B = Autodiff<$Backend<f32>>，第二个参数取默认的 NoCheckpointing，AutodiffDevice 只是 Box<DispatchDevice> 的一层包装，没有办法通过 Dispatch 用上 BalancedCheckpointing。PR #4629（「Dispatch autodiff checkpointing strategy support」，首次进入 v0.21.0-pre.3，正式版 0.21.0）给 AutodiffDevice 加上 checkpointing 字段，给 DispatchTensor 加上同名字段，并新增了 with_autodiff_backend!，注释就是今天那一句「Helper to map the runtime strategy to the compile-time Autodiff generic.」。当时的枚举叫 CheckpointingStrategy，两个值是 Balanced 和 None。0.22 开发中，PR #5400（「fix!: address fusion and autodiff edge cases」，首次进入 v0.22.0-pre.3）把它改名为 GradientCheckpointingStrategy，None 改叫 Disabled；PR #5520（「refactor(dispatch): unify routing and strengthen autodiff contract」，首次进入 v0.22.0-pre.4）再把张量上的字段换成现在的 DispatchAutodiffContext。

实验里 c 建在 Device::flex() 上，w 在 .autodiff() 设备上 require_grad，c * w 的结果和梯度：

c*w device: Device<Autodiff { device: Flex(Cpu), checkpointing: Disabled }>
w.grad = Some([2.0, 3.0, 4.0, 5.0])

loss = (a * w + c * w).sum()，w 的梯度是 a + c，a = [[1, 2], [3, 4]]、c 全是 1，结果对。两种策略混用则在 merge 里 panic：

thread 'main' (3144670) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-dispatch-0.22.1/src/tensor.rs:230:17:
assertion `left == right` failed: Gradient checkpointing strategy mismatch: Balanced vs Disabled. Tensors in the same operation must share a strategy.
  left: Balanced
 right: Disabled

backward() 不经过策略选择。Dispatch 的 AutodiffBackend::backward 只看 kind，取出 AutodiffTensor<Flex> 直接调它的 backward()；张量 API 那边，Tensor::backward 返回的 Gradients 是对 burn_autodiff::grads::Gradients 的一层不透明包装。

❗IMPORTANT
动脑筋：ad-exp 的 backend 里，a 是 .autodiff() 设备上的普通张量，w 调过 require_grad()，c 建在不带 autodiff 的 Device::flex() 上。a、w、c、a*a、a*w、c*w 六个张量的 is_autodiff / is_tracked / is_require_grad 各是什么？平台 M2 Max、macOS 27，burn 0.22.1，debug 构建。

动脑筋答案

a          is_autodiff=true  is_tracked=false is_require_grad=false strategy=Some(Disabled)
w          is_autodiff=true  is_tracked=true  is_require_grad=true  strategy=Some(Disabled)
c(flex)    is_autodiff=false is_tracked=false is_require_grad=false strategy=None
a*a        is_autodiff=true  is_tracked=false is_require_grad=false strategy=Some(Disabled)
a*w        is_autodiff=true  is_tracked=true  is_require_grad=false strategy=Some(Disabled)
c*w        is_autodiff=true  is_tracked=true  is_require_grad=false strategy=Some(Disabled)
  • a：在 autodiff 设备上，但新建的张量是 Requirement::None，不追踪。
  • w：require_grad() 之后是 Grad，三项都是真。
  • c：普通 Flex 张量，DispatchAutodiffContext::Disabled，gradient_checkpointing_strategy() 返回 None。
  • a*a：两个输入都是 None，输出也是 None。它仍在 autodiff 设备上，也仍然在图里登记了一个空步骤（下一页讲为什么）。
  • a*w：有一个输入要梯度，输出是 GradInBackward，所以「被追踪」但「不保留梯度」。
  • c*w：c 被当成常量包进来，结果和 a*w 一样；结果的 autodiff 标记来自 w。

小结

  • Autodiff<B, C> 是空结构体，只把浮点张量换成 AutodiffTensor<B>，设备、整数、布尔、量化张量都照搬内层后端，Backend 方法几乎全是转发。
  • AutodiffTensor = 内层原语 + Arc<Node> + Arc<NodeId>；NodeId 来自全局原子计数器，rc 的强引用计数是后面释放节点的依据。
  • Requirement：require_grad() 的叶子是 Grad，有输入要梯度的算子输出是 GradInBackward，其余是 None；中间结果不能 require_grad()。
  • Device::autodiff() 只给设备打标记，张量带着 DispatchAutodiffContext；生成的分派代码合并各输入的 context，在运行期选出 Autodiff<Flex, NoCheckpointing> 或 Autodiff<Flex, BalancedCheckpointing>。
  • 不带标记的浮点张量和 autodiff 张量混算时被当成常量；两种检查点策略混算会 panic；backward() 不看策略。

💬 关联讨论

发起讨论
加载中...