Rust in Everything
基于 burn 0.22.1

nn 模块

主线位置：训练主线的 forward 那一步，也是它和算子主线的接口。Linear::forward 这类方法只做形状检查，然后调 burn-tensor/src/tensor/module.rs 里的函数式算子；后者经 Dispatch 落到后端的 ModuleOps，之后的路和一次 matmul 一样。模块的派生和 Param 见 Module 派生，XxxConfig 见 Config。

flowchart TB
  lin["Linear forward"] --> flin["module linear"]
  conv["Conv2d forward"] --> fconv["module conv2d"]
  ln["LayerNorm forward"] --> fln["module layer_norm"]
  flin --> dlin["Dispatch linear 默认实现"]
  dlin --> mm["Dispatch float_matmul 和 add"]
  fconv --> dconv["Dispatch conv2d"]
  fln --> dln["Dispatch layer_norm"]
  mm --> ad["Autodiff 或内层后端"]
  dconv --> ad
  dln --> ad
  ad --> be["Flex 或 Fusion CubeBackend"]

模块 = 参数 + 一次函数式调用

burn-nn（约 2.6 万行）里的层几乎都是同一个形状：一个 #[derive(Config)] 的配置、一个 #[derive(Module)] 的结构体、init 和 forward 两个方法。四个代表：

层参数（Param）常量 / 其他字段forward 调的函数式算子
Linearweight: [d_input, d_output]，bias: Option<[d_output]>无linear
Conv2dweight: [c_out, c_in / groups, k1, k2]，bias: Option<[c_out]>stride、kernel_size、dilation、groups；#[module(skip)] paddingconv2d
LayerNormgamma: [d_model]，beta: Option<[d_model]>epsilon: f64layer_norm
Dropout无张量prob: f64，training: Param<Flag>无，直接用张量 API

Linear 的全部前向（Linear::forward）：

    pub fn forward<const D: usize>(&self, input: Tensor<D>) -> Tensor<D> {
        let weight = self.weight.val();
        let [d_input, _] = weight.dims();
        assert_shape!(input, [.., d_input]);

        linear(input, weight, self.bias.as_ref().map(|b| b.val()))
    }

self.weight.val() 是第一次触发懒初始化的地方，之后每次只是克隆一个句柄；assert_shape! 在最后一维不对时 panic（它的写法见检查与报错）。模块里没有任何数学，数学都在函数式算子里。Conv2d::forward 多一步：按输入的高宽把 PaddingConfig2d（Valid、Same 或显式值）算成上下左右四个数，装进 ConvOptions 再调 conv2d。LayerNorm::forward 在输入是 f16 时先把输入和参数都转成 f32 再算、算完转回去（accumulation_dtype：平方和在 f16 里会溢出，bf16 指数位和 f32 一样，不转）。

📝NOTE
没有蠢问题

问： PyTorch 的 nn.Linear 权重是 [out_features, in_features]，算的是 x @ W^T；burn 的是 [d_input, d_output]，算 x @ W。从 PyTorch 导权重时不就全反了？

答： 是反的，函数式 linear 的文档专门有一段「Compatibility」说明它不转置：y = x @ weight，前向就是一次 matmul。从 PyTorch 导入时，由 burn-store 的适配器把 Linear 权重转置过来（见存储与加载）。另外还有一个 LinearConfig::with_layout(LinearLayout::Col)：LinearConfig::init 按 [d_output, d_input] 初始化，再给 Param 挂三个变换：init_mapper（初始化后转置）、save_mapper（保存前转回 [d_output, d_input]）、load_mapper（读入后转置）。模块里的张量始终是 [d_input, d_output]，变的只是存取时文件里的布局。

但 Col 不能当成「直接读 PyTorch 权重」的办法：加载时的形状校验用 Param::lazy_shape，它对没读过的 Col 参数报 [d_output, d_input]，对读过的报 [d_input, d_output]。于是同一个文件配 PytorchStore，懒参数报 Shape mismatch for 'fc.weight'，已物化的参数「加载成功」后前向时才在 assert_shape! 处 panic。复现和输出见存储与加载。

函数式算子：burn-tensor 的 module.rs

burn-tensor/src/tensor/module.rs（1008 行）是一排自由函数：conv1d 到 conv_transpose3d、各种池化、embedding、interpolate、attention、batch_norm、layer_norm、linear……它们和 Tensor 的方法一样薄，只是对象是「一层网络」而不是「一个张量」。conv2d：

pub fn conv2d(
    x: Tensor<4>,
    weight: Tensor<4>,
    bias: Option<Tensor<1>>,
    options: impl Into<ConvOptions<2>>,
) -> Tensor<4> {
    let options = options.into();
    check!(TensorCheck::conv(
        "conv2d",
        x.dims(),
        weight.dims(),
        options.groups,
    ));

    Tensor::new(BridgeTensor::float(Dispatch::conv2d(
        x.primitive.into_float(),
        weight.primitive.into_float(),
        bias.map(|b| b.primitive.into_float()),
        options,
    )))
}

先用 TensorCheck::conv 检查通道数和分组，再把三个 BridgeTensor 拆成浮点原语，直接调 Dispatch::conv2d。linear 先处理两种特殊情况（1 维输入临时加一个批维；量化权重改走量化 matmul，不在每次前向把整块权重反量化），普通情况交给非泛型的 linear_impl，它调 Dispatch::linear。layer_norm 同样经 layer_norm_impl 调 Dispatch::layer_norm。*_impl 这种「泛型外壳 + 非泛型内核」的写法和它的编译时间理由见 BridgeTensor 与 *_impl。

落到 ModuleOps

Dispatch::conv2d 这些是 ModuleOps trait 的方法。和 FloatTensorOps 一样（见 Backend trait），它的方法分两种：后端必须实现的，和有默认实现、由别的算子拼出来的。三个算子的情况：

算子ModuleOps 默认实现Dispatch 转发Autodiff 覆盖Flex 覆盖CubeBackend / Fusion 覆盖
conv2d无（必须实现）是是，专用反向是是 / 是
layer_norm有：mean_dim、sub、mul、sqrt、div、reshape、add 拼成是否是，一个 CPU 内核否 / 否
linear有：linear::linear 用 matmul + add 拼成否，只有一段注释是，专用反向否否 / 否

Dispatch 的 ModuleOps 实现末尾是这样的（burn-dispatch/src/ops/module.rs）：

    // TODO: linear ops
    // fn linear(
    //         x: FloatTensor<Self>,
    //         weight: FloatTensor<Self>,
    //         bias: Option<FloatTensor<Self>>,
    //     ) -> FloatTensor<Self> {

    // }
}

Dispatch 没写 linear，于是 Dispatch::linear 是 trait 的默认实现，B = Dispatch：它在 Dispatch 这一层就把 linear 拆成了 Dispatch::float_matmul 和 Dispatch::float_add，每一个再各自分派。Autodiff 那个带专用反向的 linear 永远收不到调用。

用 lldb 验证（实验 nnpath，debug，Device::flex().autodiff() 上一个 Linear(4, 2) 前向一次）。断点一设在 Flex 的 gemm_call 上，断点二设在任何「Autodiff 的 ModuleOps 的 linear」上：

(lldb) breakpoint set -r "burn_flex::ops::matmul::gemm_call"
Breakpoint 1: 3 locations.
(lldb) breakpoint set -r "Autodiff.*ModuleOps.*>::linear$"
Breakpoint 2: no locations (pending).
WARNING:  Unable to resolve breakpoint to any actual locations.
(lldb) run linear
...
(lldb) bt 12
...
     frame #4: 0x0000000102042a2c nnpath`<burn_flex::backend::Flex as burn_backend::backend::ops::tensor::FloatTensorOps<burn_flex::backend::Flex>>::float_matmul(lhs=<unavailable>, rhs=FlexTensor @ 0x000000016fde3450) at float.rs:185:9
     frame #5: 0x00000001007ca864 nnpath`<burn_autodiff::backend::Autodiff<burn_flex::backend::Flex> as burn_backend::backend::ops::tensor::FloatTensorOps<burn_autodiff::backend::Autodiff<burn_flex::backend::Flex>>>::float_matmul(lhs=AutodiffTensor<burn_flex::backend::Flex> @ 0x000000016fdf0868, rhs=AutodiffTensor<burn_flex::backend::Flex> @ 0x000000016fdf08f0) at tensor.rs:670:21
     frame #6: 0x00000001003d4340 nnpath`<burn_dispatch::backend::Dispatch as burn_backend::backend::ops::tensor::FloatTensorOps<burn_dispatch::backend::Dispatch>>::float_matmul(lhs=DispatchTensor @ 0x000000016fdf7c38, rhs=DispatchTensor @ 0x000000016fdf8178) at tensor.rs:105:9
     frame #7: 0x000000010004b5e4 nnpath`burn_backend::backend::ops::modules::linear::linear::<burn_dispatch::backend::Dispatch>(x=DispatchTensor @ 0x000000016fdf8888, weight=DispatchTensor @ 0x000000016fdf89d0, bias=Option<burn_dispatch::tensor::DispatchTensor> @ 0x000000016fdf8b18) at linear.rs:47:26
     frame #8: 0x000000010004ed44 nnpath`<burn_dispatch::backend::Dispatch as burn_backend::backend::ops::modules::base::ModuleOps<burn_dispatch::backend::Dispatch>>::linear(x=<unavailable>, weight=DispatchTensor @ 0x000000016fdf8f78, bias=Option<burn_dispatch::tensor::DispatchTensor> @ 0x000000016fdf9278) at base.rs:325:9
     frame #9: 0x0000000100047c0c nnpath`burn_tensor::tensor::module::linear_impl(input=<unavailable>, weight=BridgeTensor @ 0x000000016fdfb6c0, bias=Option<burn_tensor::bridge::kind::BridgeTensor> @ 0x000000016fdfb840) at module.rs:787:25
     frame #10: 0x000000010001b07c nnpath`burn_tensor::tensor::module::linear::<2>(input=Tensor<2, burn_tensor::bridge::kind::Float> @ 0x000000016fdfc400, weight=Tensor<2, burn_tensor::bridge::kind::Float> @ 0x000000016fdfc580, bias=Option<burn_tensor::tensor::api::base::Tensor<1, burn_tensor::bridge::kind::Float>> @ 0x000000016fdfc700) at module.rs:775:17
     frame #11: 0x000000010002e010 nnpath`<burn_nn::modules::linear::Linear>::forward::<2>(self=0x000000016fdfcbf0, input=Tensor<2, burn_tensor::bridge::kind::Float> @ 0x000000016fdfce40) at linear.rs:138:9

从下往上读：Linear::forward → module::linear → linear_impl → Dispatch 的 ModuleOps::linear（base.rs:325，即默认实现那一行）→ linear::linear::<Dispatch> → Dispatch::float_matmul → Autodiff::float_matmul → Flex::float_matmul。断点二「no locations」不只是没命中：这个函数根本没有被单态化进二进制。对整个二进制 nm -C，名字以 ModuleOps<..>>::linear 结尾的符号只有 1 个，<Dispatch as ModuleOps<Dispatch>>::linear；conv2d 有 18 个、layer_norm 有 16 个（每种后端组合一个）。

conv2d 的同一个实验（断点设在 Flex 的 ModuleOps::conv2d）走的是另一条路，Autodiff::conv2d 在中间：

     frame #1: 0x0000000100777e7c nnpath`<burn_autodiff::backend::Autodiff<burn_flex::backend::Flex> as burn_backend::backend::ops::modules::base::ModuleOps<burn_autodiff::backend::Autodiff<burn_flex::backend::Flex>>>::conv2d(...) at module.rs:647:25
     frame #2: 0x00000001003968e0 nnpath`<burn_dispatch::backend::Dispatch as burn_backend::backend::ops::modules::base::ModuleOps<burn_dispatch::backend::Dispatch>>::conv2d(...) at module.rs:63:9
     frame #3: 0x000000010001aa24 nnpath`burn_tensor::tensor::module::conv2d::<burn_std::ops::ConvOptions<2>>(...) at module.rs:192:37
     frame #4: 0x000000010002f008 nnpath`<burn_nn::modules::conv::conv2d::Conv2d>::forward(...) at conv2d.rs:181:9

（参数列表用 ... 省略。）layer_norm 的路线取决于设备。断点同时设在 Flex 的 ModuleOps::layer_norm 和 Flex::float_mean_dim 上，第一个命中的：

== Device::flex()
     frame #0: ... <burn_flex::backend::Flex as burn_backend::backend::ops::modules::base::ModuleOps<burn_flex::backend::Flex>>::layer_norm(...) at module.rs:752:9
     frame #1: ... <burn_dispatch::backend::Dispatch as burn_backend::backend::ops::modules::base::ModuleOps<burn_dispatch::backend::Dispatch>>::layer_norm(...) at module.rs:575:9
== Device::flex().autodiff()
     frame #0: ... <burn_flex::backend::Flex as burn_backend::backend::ops::tensor::FloatTensorOps<burn_flex::backend::Flex>>::float_mean_dim(...) at float.rs:821:9
     frame #1: ... <burn_autodiff::backend::Autodiff<burn_flex::backend::Flex> as burn_backend::backend::ops::tensor::FloatTensorOps<burn_autodiff::backend::Autodiff<burn_flex::backend::Flex>>>::float_mean_dim(...) at tensor.rs:2676:53
     frame #2: ... <burn_autodiff::backend::Autodiff<burn_flex::backend::Flex> as burn_backend::backend::ops::modules::base::ModuleOps<burn_autodiff::backend::Autodiff<burn_flex::backend::Flex>>>::layer_norm(...) at base.rs:994:20
     frame #3: ... <burn_dispatch::backend::Dispatch as burn_backend::backend::ops::modules::base::ModuleOps<burn_dispatch::backend::Dispatch>>::layer_norm(...) at module.rs:575:9

（地址和参数用 ... 省略。）普通 Flex 上，Dispatch 把 layer_norm 交给 Flex 自己的实现（crate::ops::activation::layer_norm，一个 CPU 内核）；autodiff 设备上，Autodiff 没有覆盖 layer_norm，用的是 ModuleOps 的默认实现（base.rs:994 是其中第一行 B::float_mean_dim），每一步都是一个被记录、各自有反向的 Autodiff 算子，Flex 的融合实现用不上。GPU 上 CubeBackend 和 Fusion 都没有覆盖 layer_norm，默认实现拆出来的这串逐元素和归约算子交给融合去合并（见融合搜索）。

设备linear 实际执行反向图里的节点
flex()Dispatch 默认实现 → Flex 的 matmul + add无
flex().autodiff()同上，每个经过 Autodiffmatmul、偏置的 reshape、add 各一个通用步骤
0.21 的 Autodiff<Flex> 当 BAutodiff::linear → Flex 默认实现一个 LinearWithBias 步骤

最后一行是从源码推出来的：0.21 用户直接把 Autodiff<B> 当后端类型参数，B::linear 就是 Autodiff::linear。0.22 一切都先经过 Dispatch，Dispatch 不转发，路就断了。结果（数值）是一样的，梯度也对，只是专用反向里「哪个输入没被追踪就不存另一个」这类省内存的处理用不上。

📝NOTE
冷知识：ModuleOps::linear 只活过了一个版本
#4747「Add linear op to ModuleOps for fused matmul+bias」（首次发布于 v0.21.0-pre.4）给 ModuleOps 加了 linear 和三个反向方法，提交说明：「allowing backends to fuse the matmul and bias addition into a single operation」，并让 module::linear 改调 B::linear。同一个 PR 的评审修改里有两条和本页直接相关：autodiff 的实现「checkpoint x only when weight is tracked and vice versa, mirroring the matmul pattern. Saves memory for frozen layers」；以及 burn-fusion 删掉了自己的覆盖，「so the default decomposition (matmul + add / sum) runs at the Fusion<B> level, letting burn-cubecl-fusion see the primitive IR nodes it can actually fuse」。不到三周后，#4717（v0.22.0-pre.1）把所有张量改走 Dispatch，在 Dispatch 的 ModuleOps 里留下了 // TODO: linear ops 那段注释，一直到 0.22.1 没有补上。

初始化器

Initializer 是一个 #[derive(Config)] 枚举：Constant、Ones、Zeros、Uniform、Normal、KaimingUniform、KaimingNormal、XavierUniform、XavierNormal、Orthogonal。它作为配置字段存在 LinearConfig、Conv2dConfig 里，所以会随配置一起存进 JSON（见 Config）。

init_with(shape, fan_in, fan_out, device) 返回 Param::uninitialized：闭包里才调 init_tensor 生成张量，外面包一层 device.memory_persistent_allocations（文档：「Sets the current allocation mode to persistent」，参数会一直活着，分配器按长期内存对待），设备是 autodiff 时再 require_grad()。所以「初始化器」决定的是第一次 val() 时抽什么随机数，抽的时机见 Module 派生的懒初始化一节。

Kaiming 和 Xavier 的边界：

  • KaimingUniform { gain, fan_out_only }：a = sqrt(3) * gain / sqrt(fan)，在 [-a, a] 上均匀抽样；fan 是 fan_in（fan_out_only 时用 fan_out）。
  • XavierUniform { gain }：a = sqrt(3) * gain * sqrt(2 / (fan_in + fan_out))。
  • 没传 fan（用 init 而不是 init_with）时，kaiming_std 会 panic：「Can't use Kaiming initialization without specifying fan. Use init_with method.」

Linear 和 Conv2d 的默认值是 KaimingUniform { gain: 1/sqrt(3), fan_out_only: false }，代进去 a = 1/sqrt(fan_in)，和 Linear 文档注释里写的 U(-k, k), k = sqrt(1 / d_input) 一致。Conv2d 的 fan_in = c_in / groups * k1 * k2。偏置用同一个初始化器、同一组 fan。LayerNorm 不看配置里的初始化器：gamma 用 Ones，beta 用 Zeros。实验 init（debug，Flex，seed(1)）：

Linear(256, 64).weight             min  -0.06250  max   0.06250  bound 0.06250
Linear(256, 64).bias               min  -0.05866  max   0.06228  bound 0.06250
Conv2d weight shape [8, 3, 3, 3]
Conv2d([3, 8], [3, 3]).weight      min  -0.19154  max   0.19203  bound 0.19245
Linear XavierUniform{gain: 1}      min  -0.13692  max   0.13693  bound 0.13693
LayerNorm(4) gamma [1.0, 1.0, 1.0, 1.0] beta [0.0, 0.0, 0.0, 0.0]

bound 一列是按上面的公式手算的：1/sqrt(256) = 0.0625、1/sqrt(27) = 0.19245、sqrt(6/320) = 0.13693。16 384 个权重抽样，最小最大值贴着边界；卷积只有 216 个，离边界还有一点。Conv2d 的权重形状是 [c_out, c_in, k1, k2]，和 Linear 不同，它和 PyTorch 的卷积布局一致。

train / valid 怎么影响 Dropout

Dropout::forward：

    pub fn forward<const D: usize>(&self, input: Tensor<D>) -> Tensor<D> {
        // Both, and for different reasons. The device says a backward is possible at all; the
        // flag says this layer takes part in one, which a subtree frozen in place on the training
        // device does not.
        if !self.training.is_enabled() || !input.device().is_autodiff() || self.prob == 0.0 {
            return input;
        }

        let prob_keep = 1.0 - self.prob;
        let random = input.random_like(Distribution::Bernoulli(prob_keep));
        let x = input * random;

        x * (1.0 / prob_keep)
    }

三个条件任何一个不满足，输入原样返回：

  1. training 这个 Flag 是开的。valid() 和 freeze() 会把它关掉（见 Module 派生的 train / valid 一节）。
  2. 输入张量的设备是 autodiff 设备。 这一条和模块状态无关，看的是这次前向的输入。
  3. 概率不是 0。

保留下来的元素乘 1 / (1 - prob)，期望不变（inverted dropout）。随机掩码用 random_like(Bernoulli) 在输入所在的后端上生成，GPU 上由 cubek 的随机数内核生成。DropoutConfig::init 会拒绝 [0, 1) 以外和非有限的概率：#5860（v0.22.0）把原来的「[0, 1] 之间」收紧成这样，概率 1 会让 1 / (1 - prob) 除以零。

❗IMPORTANT
动脑筋：下面六行各丢掉了多少个元素？

实验 dropout（debug）：一个 DropoutConfig::new(0.5).init()，对 Tensor::<2>::ones([1, 10_000], device) 前向一次，数输出里 0 的个数和最大值。六种情况：普通 Device::flex()；flex().autodiff()；flex().autodiff() 但 dropout 先 valid()；flex().autodiff() 但 dropout 先 freeze()；普通 Metal；metal(DefaultDevice).autodiff()。

BatchNorm 用同样的两个条件（输入在 autodiff 设备上、training 开着）在两条路之间选：训练路径用本批的统计量，并更新 RunningState 里的运行均值和方差；否则走推理路径，用运行值。

动脑筋答案

flex                               zeros     0  max 1
flex.autodiff()                    zeros  4964  max 2
flex.autodiff(), dropout.valid()   zeros     0  max 1
flex.autodiff(), dropout.freeze()  zeros     0  max 1
metal                              zeros     0  max 1
metal.autodiff()                   zeros  4938  max 2
  • 只有两个 autodiff 设备上真的丢了，约一半（4964、4938，Metal 那一行每次运行都不同，程序只给 Flex 设了种子），保留的元素放大成 2。
  • 普通 Flex 和普通 Metal 上一个都没丢：模块是新建的，training 是开的，但输入不在 autodiff 设备上，第二个条件就返回了。在不带 autodiff 的设备上做推理，不调 valid() 也不会有 dropout。
  • 反过来，在训练设备上评估（同一个 .autodiff() 设备上跑验证集）而忘了 valid()，dropout 是开着的，指标会偏低；而且这样前向会一直往计算图里记节点（见backward 与 Gradients）。

小结

  • burn-nn 的层 = 一个 Config + 一个 Module 结构体；forward 只做形状检查，然后调 burn-tensor/src/tensor/module.rs 的函数式算子，后者直接调 Dispatch 的 ModuleOps 方法。
  • ModuleOps 里 conv2d 必须实现，linear、layer_norm 有用基础算子拼出来的默认实现。0.22.1 的 Dispatch 不转发 linear，在自己这一层就拆成 matmul + add，Autodiff 的专用 linear 反向没有被编进二进制；layer_norm 只在普通 Flex 上走融合实现，autodiff 设备上拆成一串被记录的算子。
  • 初始化器是 #[derive(Config)] 枚举，init_with 返回懒参数；默认的 KaimingUniform { gain: 1/sqrt(3) } 给出 U(-1/sqrt(fan_in), 1/sqrt(fan_in))。Linear 权重是 [d_input, d_output]，和 PyTorch 相反。
  • Dropout 生效要三个条件：Flag 开着、输入在 autodiff 设备上、概率非 0；所以普通设备上推理自动关闭，训练设备上评估必须 valid()。

💬 关联讨论

发起讨论
加载中...