Rust in Everything
基于 burn 0.22.1

Module 派生

主线位置：训练主线的第一步。#[derive(Module)] 为模型生成 visit / map 这些遍历，后面 GradientsParams::from_grads（一个 ModuleVisitor）和 ModuleOptimizer::step（一个 ModuleMapper）都靠它找到每个参数；Param 上的 ParamId 是梯度和优化器状态的键。一步训练怎么串起来见一个训练步，本页讲这些遍历和参数本身。

flowchart TB
  derive["derive Module"] --> impl["impl Module for Mlp"]
  impl --> field{"逐个字段"}
  field -->|子模块| rec["递归调字段的方法"]
  field -->|常量或 skip| copy["原样搬过去"]
  rec --> pt["Param Tensor D"]
  rec --> pf["Param Flag"]
  pt --> vf["visit_float 或 map_float"]
  pf --> vfl["visit_flag 或 map_flag"]
  vf --> id["按 ParamId 认参数"]

Module trait：没有后端参数

Module 在 burn-core/src/module/base.rs：

#[diagnostic::on_unimplemented(
    note = "If `{Self}` is used as configuration or runtime state in a `#[derive(Module)]` type, mark that field with `#[module(skip)]`."
)]
pub trait Module: Clone + Send + core::fmt::Debug {

0.21 是 pub trait Module<B: Backend>，访问者是 ModuleVisitor<B>，参数是 Param<Tensor<B, D>>。#4717 去掉张量的后端参数时，Module 的也一起去掉了：后端由参数张量所在的 Device 决定（见 Device），一个 Linear 放在 Flex 上和放在 Metal 上是同一个类型。

trait 里有八个方法没有默认实现，由派生宏生成；其余都是在这八个之上写好的：

方法签名要点作用
visit&self, &mut V: ModuleVisitor只读遍历每个参数和控制标志
mapself, &mut M: ModuleMapper → Self按值消耗，返回换过参数的新模块
collect_devices&self, Devices → Devices收集参数所在的设备，不触发懒初始化
to_device / forkself, &Device → Self搬到另一个设备，后者保证参数在新设备上仍是叶子
train / validself / &self → Self训练态和验证态之间切换
materializeself → Self把 LoRA 这类重参数化（reparameterization）并回权重

materialize 这个名字容易误会：它和懒初始化无关，处理的是 apply_lora 挂上去的重参数化，没挂的参数原样返回，连懒参数都不会被初始化。

写好的方法里，num_params、set_require_grad、no_grad、freeze 这些都由文件开头一个内部的 module! 宏现场定义一个访问者或映射器，再调 visit / map。比如 num_params 是一个把每个浮点参数的 lazy_shape().num_elements() 加起来的访问者。into_record、save_file 也在这个 trait 上，记录的格式和加载见存储与加载。

#[derive(Module)] 生成什么

入口 module_derive 转给 derive_impl：结构体和枚举各有一个代码生成器，联合体直接报错。生成器先判断有没有「模块字段」：有就走 generate_module_standard，一个都没有（比如 struct Relu;，或者只有 f64 字段）就走 generate_module_stateless，后者整个 impl Module 只有一行 burn::empty!(module);，八个方法全是「什么也不做」或 clone。

实验 crate module-exp 的 expand（依赖同 base 模板），一个模块字段加一个常量字段的结构体，再加一个枚举：

#[derive(Module, Debug)]
struct Head {
  proj: Linear,
  scale: f64,
}

#[derive(Module, Debug)]
enum Block {
  Small(Linear),
  Large(Head),
}

cargo expand --bin expand 的结果（Head 部分，省略了 collect_devices、fork、materialize，它们和 to_device / train 同形）：

impl burn::module::Module for Head {
    fn num_params(&self) -> usize {
        let mut num_params = 0;
        num_params += burn::module::Module::num_params(&self.proj);
        num_params
    }
    fn visit<Visitor: burn::module::ModuleVisitor>(&self, visitor: &mut Visitor) {
        visitor.enter_module("proj", "Struct:Head");
        burn::module::Module::visit(&self.proj, visitor);
        visitor.exit_module("proj", "Struct:Head");
    }
    fn map<Mapper: burn::module::ModuleMapper>(self, mapper: &mut Mapper) -> Self {
        mapper.enter_module("proj", "Struct:Head");
        let proj = burn::module::Module::map(self.proj, mapper);
        mapper.exit_module("proj", "Struct:Head");
        let scale = self.scale;
        Self { proj, scale }
    }
    // ...
    fn to_device(self, device: &burn::tensor::Device) -> Self {
        let proj = burn::module::Module::to_device(self.proj, device);
        let scale = self.scale;
        Self { proj, scale }
    }
    // ...
    fn valid(&self) -> Self {
        let proj = burn::module::Module::valid(&self.proj);
        let scale = self.scale.clone();
        Self { proj, scale }
    }
    fn train(self) -> Self {
        let Self { proj, scale } = self;
        let proj = burn::module::Module::train(proj);
        let scale = scale;
        Self { proj, scale }
    }
    // ...
}
impl core::fmt::Display for Head {
    // ...
}
impl burn::module::ModuleDisplayDefault for Head {
    // ...
}
impl Clone for Head {
    fn clone(&self) -> Self {
        let proj = self.proj.clone();
        let scale = self.scale.clone();
        Self { proj, scale }
    }
}
impl burn::module::ModuleDisplay for Head {}

几点：

  • 每个方法都是「对模块字段递归调同名方法，其余字段原样搬过去」。没有任何运行期反射，模块树的形状在编译期就写死在这些函数里。
  • visit / map 进出每个字段时调 enter_module(字段名, "Struct:Head") / exit_module，访问者靠它拼出 layers.0.weight 这样的路径（下文实验）。常量字段 scale 连 enter_module 都没有。
  • num_params 被派生宏覆盖成逐字段相加，到 Param 那一层才落回 trait 的默认实现（一次访问者遍历）。
  • 宏还生成了 Clone、Display（打印 PyTorch 风格的模块树）、ModuleDisplayDefault、ModuleDisplay。Debug 不生成，trait 要求它，所以要自己写 #[derive(Module, Debug)]。

枚举的展开是同样的东西套一层 match，路径里用变体名：

    fn visit<Visitor: burn::module::ModuleVisitor>(&self, visitor: &mut Visitor) {
        match self {
            Self::Small(module) => {
                visitor.enter_module("Small", "Enum:Block");
                burn::module::Module::visit(module, visitor);
                visitor.exit_module("Small", "Enum:Block");
            }
            Self::Large(module) => {
                visitor.enter_module("Large", "Enum:Block");
                burn::module::Module::visit(module, visitor);
                visitor.exit_module("Large", "Enum:Block");
            }
        }
    }

parse_variants 只接受「恰好一个字段的元组变体」：结构体变体、单元变体、多字段元组变体各有一条编译错误。

📝NOTE
没有蠢问题

问： 模型要能 clone()，我在 #[derive(Module, Debug)] 里再加一个 Clone，为什么编译不过？

答： 因为 Clone 已经由 Module 的派生宏生成了（上面展开里的 impl Clone for Head），再派生一次就是两个实现。clone_err（debug）：

error[E0119]: conflicting implementations of trait `std::clone::Clone` for type `Model`
 --> src/bin/clone_err.rs:4:18
  |
4 | #[derive(Module, Clone, Debug)]
  |          ------  ^^^^^ conflicting implementation for `Model`
  |          |
  |          first implementation here

去掉 Clone 就好，.clone() 照样能用。生成的 clone 逐字段调 clone()，到参数那一层只是复制一个 Arc 和张量句柄，不拷数据，见下文 Param。

字段：子模块、常量和 #[module(skip)]

派生宏只看语法，不看类型实现了什么 trait。parse_module_field_type 给每个字段分一类：

字段判定生成的代码里
常量is_primitive_type：bool、char、整数、浮点、String，以及只由它们组成的元组、数组、引用，或者装在 is_constant_container 列出的容器（Option、Vec、Box、Arc、HashMap……）里不遍历、不搬设备，valid 里 clone，显示时用 Debug
#[module(skip)]写了这个属性和常量一样；只要求 Clone + Debug + Send
子模块其余所有字段，包括泛型参数 M所有方法都递归调用，要求实现 Module 和 ModuleDisplay

判定是按类型路径的最后一段名字做的：Vec<f32> 是常量，Vec<Linear> 是子模块，Block<f32> 这种自定义类型即使参数是 f32 也算子模块（源码里有一个测试专门防止容器「藏住」自定义模块）。Param 字段不能 skip，带 Module 约束的泛型也不能 skip，各有一条编译错误。

Conv2d 是一个现成的例子：stride: [usize; 2]、groups: usize 被认成常量，padding: PaddingConfig2d 是自定义枚举，只能写 #[module(skip)]（见 nn 模块）。

burn-core 给常见类型写好了 Module：Option<T>、Vec<T>、[T; N]、2 到 10 元的元组（param/primitive.rs），String、bool 和数值类型用 empty! 宏实现成空模块。裸的 Tensor<D, K> 字段也有 Module 实现，但它的 visit 和 map 什么也不做：会跟着 to_device 搬、valid 时去掉 autodiff，却不会被优化器看到，也不进记录（源码注释：TODO: tensor record should persist）。要训练的张量得包成 Param。

0.21 的派生宏靠「字段类型里有没有 B」区分 Linear<B> 和配置字段，没有 B 的自定义类型自动忽略。0.22 没有 B 了，规则反过来：认不出来的都当子模块。一个自定义配置枚举忘了 skip（skip_err，debug，4 个错误，节选第一个）：

#[derive(Clone, Debug)]
enum Pooling {
  Mean,
  Max,
}

#[derive(Module, Debug)]
struct Model {
  linear: Linear,
  pooling: Pooling,
}

error[E0277]: the trait bound `Pooling: burn::module::Module` is not satisfied
  --> src/bin/skip_err.rs:10:10
   |
10 | #[derive(Module, Debug)]
   |          ^^^^^^ unsatisfied trait bound
   |
help: the trait `burn::module::Module` is not implemented for `Pooling`
  --> src/bin/skip_err.rs:5:1
   |
 5 | enum Pooling {
   | ^^^^^^^^^^^^
   = note: If `Pooling` is used as configuration or runtime state in a `#[derive(Module)]` type, mark that field with `#[module(skip)]`.
   = help: the following other types implement trait `burn::module::Module`:
             (L0, L1)
             (L0, L1, L2)
...
           and 124 others

其余三个错误是同一个字段对 Module 和 ModuleDisplay 的不满足，每个都带同一句 note。

📝NOTE
冷知识：那句 note 是 0.22.1 才加的

note: 那一行来自 Module 和 ModuleDisplay trait 上的 #[diagnostic::on_unimplemented]。它是 #5997「fix(module): clarify derive field handling and suggest #[module(skip)]」加的，首次发布于 v0.22.1，0.22.0 里没有：那时同样的代码只报「Pooling: Module is not satisfied」，不告诉你该加什么。同一个 PR 在迁移指南里补了「Module fields」一节，原话是：「Without backend generics, the derive can no longer distinguish Linear<B> from configuration by the presence of B. Fields now default to modules, except recognized constants such as usize, f32, bool and String.」并把 char、u128、i128 加进了常量名单。#[module(skip)] 本身更早，#4618（v0.21.0-pre.3）引入它并弃用 Ignored<T> 包装。

Param：id、懒初始化和训练状态

参数是 Param<T>：

pub struct Param<T: ParameterValue> {
    /// The unique ID of this parameter. This is used by eg. optimizers to associate a gradient with a specific parameter.
    pub id: ParamId,
    /// Shared lazy initialization state across all clones of this parameter.
    /// The `Arc` exists solely to coordinate lazy initialization. It is not a general
    /// shared-ownership mechanism. Any mutation forks into a new `LazyInitState`.
    pub(crate) state: Arc<LazyInitState<T>>,
    pub(crate) param_mapper: ParamMapper<T>,
    // ...
    pub(crate) is_active: bool,
    /// Optional transformation that materializes the effective value from the stored base.
    pub(crate) reparameterization: Option<Box<dyn DynReparameterization>>,
}

T 是一个封闭的集合（ParameterValue 是 sealed trait）：三种张量 Tensor<D>、Tensor<D, Int>、Tensor<D, Bool>，以及控制标志 Flag。Param<Tensor<D>> 的 Module 实现在 param/tensor.rs：visit 调 visitor.visit_float(self)，map 调 mapper.map_float(self)。派生宏的递归到这里落地。

ParamId

ParamId（在 burn-std）只是一个 u64，ParamId::new() 用 IdGenerator::generate 取 8 个随机字节，不是递增计数器。Display 把它编码成 base32（BASE32_DNSSEC），所以打印出来像 9199rs1gnooc6，每次运行都不同。FromStr 还认旧格式的 UUID（哈希成 64 位），兼容老记录。

ParamId 是训练主线上认参数的唯一凭据：GradientsParams 按它存梯度，ModuleOptimizer 按它存 Adam 状态（见一个训练步第 4、5 步）。所以每个改参数值的操作都要保住 id：Param::map 先 consume 拆出 (id, 值, mapper)，算完新值再 from_mapped_value 用同一个 id 装回去。

懒初始化

LinearConfig::init 不分配权重。它调 Initializer::init_with，后者返回 Param::uninitialized：存下一个初始化闭包、目标设备、要不要梯度和形状（Uninitialized）。真正的值在 LazyInitState 的一个 SyncOnceCell 里，第一次 val() 经 Deref 进到 LazyInitState::val，由 get_or_init 跑闭包：

    fn val(&self) -> &T {
        self.value.get_or_init(|| {
            let mut init = self
                .initialization
                .as_ref()
                .expect("Should have an initialization when no state provided.")
                .write();
            let state = init.take().expect("Should exist when not initialized");
            state.initialize()
        })
    }

Uninitialized::initialize 传给闭包的「要不要梯度」是 self.is_require_grad && self.device.is_autodiff()：在普通设备上建出来的参数不要梯度，但「本来想要」这件事记在 is_active 里，留给 train()。

num_params、devices 只读 Uninitialized 里存的形状和设备（lazy_shape、lazy_device），不触发初始化。Param 的 Clone 复制的是 Arc，所以克隆和原件共享同一个懒状态。实验 lazy（debug，Flex）：

let device = Device::flex();
let lin = LinearConfig::new(3, 2).init(&device);
println!("after init:        weight initialized = {}", lin.weight.is_initialized());
println!("num_params = {:<5}  weight initialized = {}", lin.num_params(), lin.weight.is_initialized());
println!("devices = {:?}  weight initialized = {}", lin.devices(), lin.weight.is_initialized());

let copy = lin.clone();
let _w = lin.weight.val();
println!(
  "after lin.weight.val(): lin {} / copy {} / bias {}",
  lin.weight.is_initialized(),
  copy.weight.is_initialized(),
  lin.bias.as_ref().unwrap().is_initialized()
);
println!("ids: lin {} copy {}", lin.weight.id, copy.weight.id);
println!("Param Debug: {:?}", lin.weight);

after init:        weight initialized = false
num_params = 8      weight initialized = false
devices = [Device<Flex(Cpu)>]  weight initialized = false
after lin.weight.val(): lin true / copy true / bias false
ids: lin 9199rs1gnooc6 copy 9199rs1gnooc6
Param Debug: Param: 9199rs1gnooc6 - ParamMapper { load: false, save: false }

只取了 weight，bias 仍未初始化；克隆 copy 跟着原件一起变成已初始化，id 相同。

懒初始化有一个看得见的后果：随机数是在第一次取值时抽的，不是在 init 时。同一个种子，先建 a 再建 b，只是取值顺序反过来，两层的权重就对调了（同一个程序的后半段）：

for order in ["a then b", "b then a"] {
  device.seed(0);
  let a = LinearConfig::new(2, 2).with_bias(false).init(&device);
  let b = LinearConfig::new(2, 2).with_bias(false).init(&device);
  let (first, second) = if order == "a then b" { (&a, &b) } else { (&b, &a) };
  let _ = first.weight.val();
  let _ = second.weight.val();
  // print a and b
}

a then b: a = [0.42632765, 0.32684356, 0.07819015, 0.3867311]
a then b: b = [0.43747574, -0.670557, 0.24250603, 0.11901915]
b then a: a = [0.43747574, -0.670557, 0.24250603, 0.11901915]
b then a: b = [0.42632765, 0.32684356, 0.07819015, 0.3867311]

想让初始权重可复现，除了 seed，还要保证第一次前向（或第一次读权重）的顺序不变。

💡TIP
打个比方：定做家具

Config::init 像在家具店下单：店里只登记了尺寸、送货地址和做法（Uninitialized 的 shape、device、init 闭包），给你一张取货单（Param 和它的 ParamId）。复印取货单（clone）不会多做一件家具，几张单子对应同一个订单（共享的 Arc<LazyInitState>）。第一次有人拿单子来取货（val()），店里才开工（get_or_init），之后谁来取都是这一件。还没开工就搬家（to_device），只要改送货地址（下文的 map_to_device）；已经做好了再搬家，就得整件运过去。

比方到这里为止：真店里先下单的先开工；burn 里材料（随机数）是开工那一刻才从同一个料堆里取的，谁先来取货谁拿到前面的料，所以上面的 a、b 会对调。

Param 还带一个 ParamMapper，可以挂保存时和加载时的变换（save_mapper / load_mapper），以及 init_mapper：对懒参数，它把变换接在初始化闭包后面，不触发求值。Linear 的列优先布局用的就是这三个（见 nn 模块）。

Flag：不是张量的参数

Flag 是一个包着 bool 的类型，模块里以 Param<Flag> 出现，burn-nn 里用它的是 Dropout、Noise、RRelu、BatchNorm 的 training 字段，都是「训练时和推理时行为不同」的层。它有 ParamId，所以能被参数组（ParamGroup）按 id 或路径选中；它不在任何设备上，to_device / fork 原样返回；它不进记录。ModuleVisitor::visit_flag 和 ModuleMapper::map_flag 专门处理它，默认实现什么也不做。

遍历：ModuleVisitor 与 ModuleMapper

ModuleVisitor 和 ModuleMapper 的每个方法都有空的默认实现，用户只覆盖关心的那几个：

访问者（&Param）映射器（按值进、按值出）什么时候调
visit_float / visit_int / visit_boolmap_float / map_int / map_bool遇到对应种类的张量参数
visit_flagmap_flag遇到 Param<Flag>
enter_module / exit_module同名进出一个字段、Vec 元素、元组元素、枚举变体

container_type 参数是 "Struct:Linear"、"Enum:Block"、"Vec"、"Tuple"、"Array" 这样的字符串。Option 不调 enter_module，这样路径里是 bias 而不是 bias.Some。

实验 visit：一个带 Vec<Linear>、LayerNorm、Dropout、Option<Linear> 和一个 String 的模型，写一个打印路径的访问者：

#[derive(Module, Debug)]
struct Net {
  layers: Vec<Linear>,
  norm: LayerNorm,
  dropout: Dropout,
  head: Option<Linear>,
  name: String,
}

struct Printer {
  path: Vec<String>,
}

impl ModuleVisitor for Printer {
  fn enter_module(&mut self, name: &str, container_type: &str) {
    println!("{:indent$}enter {name} ({container_type})", "", indent = 2 * self.path.len());
    self.path.push(name.to_string());
  }

  fn exit_module(&mut self, _name: &str, _container_type: &str) {
    self.path.pop();
  }

  fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<D>>) {
    // print path, lazy_shape and id
  }

  fn visit_flag(&mut self, flag: &Param<Flag>) {
    // print path and value
  }
}

head 是 LinearConfig::new(8, 2).with_bias(false)，其余层带偏置：

enter layers (Struct:Net)
  enter 0 (Vec)
    enter weight (Struct:Linear)
      float layers.0.weight shape [4, 8] id ic13nbqi84dec
    enter bias (Struct:Linear)
      float layers.0.bias shape [8] id tumvmkapa4mm2
  enter 1 (Vec)
    enter weight (Struct:Linear)
      float layers.1.weight shape [8, 8] id 6ht90dskudvh6
    enter bias (Struct:Linear)
      float layers.1.bias shape [8] id r05ob8q2r0utu
enter norm (Struct:Net)
  enter gamma (Struct:LayerNorm)
    float norm.gamma shape [8] id 6rhfg7fb0pdkc
  enter beta (Struct:LayerNorm)
    float norm.beta shape [8] id 5rl41pdb2d544
enter dropout (Struct:Net)
  enter training (Struct:Dropout)
    flag  dropout.training = enabled
enter head (Struct:Net)
  enter weight (Struct:Linear)
    float head.weight shape [8, 2] id jsj0j7d1fqs6i
  enter bias (Struct:Linear)
num_params = 144
  • name: String 是常量，没有出现；Dropout 的 prob: f64 也没有。
  • head.bias 是 None：Linear 的派生代码照样进出 bias 这个字段，但 Option 里没东西，所以只有 enter 没有参数。
  • 路径用 . 连起来就是记录里的键名（layers.0.weight），存储和 PyTorch 导入时的键名重映射靠的就是这条路径（见存储与加载）。

映射器的例子：把每个浮点参数换成全零，map 之后 id 不变。

struct Zero;

impl ModuleMapper for Zero {
  fn map_float<const D: usize>(&mut self, param: Param<Tensor<D>>) -> Param<Tensor<D>> {
    param.map(|t| t.zeros_like())
  }
}

after map: id ic13nbqi84dec -> ic13nbqi84dec, sum = 0

ModuleOptimizer::step 就是一个这样的映射器，只是 map_float 里换的是 Adam 算出来的新张量（见优化器）。

train、valid、freeze：同一个类型的两种状态

valid 和 train 都在派生代码里逐字段递归，落到参数上：

  • Param<Tensor>::valid：self.base().without_autodiff().set_require_grad(false)，id 不变，is_active 原样保留。
  • Param<Tensor>::train：Tensor::from_inner(self.val()).set_require_grad(is_active)。from_inner 等于 autodiff()：打开 autodiff，用默认的检查点策略。
  • Param<Flag>::valid 把值设成 false、保留 is_active；train 把值恢复成 is_active。

is_active 就是「配置上想不想训练」，和张量当前是否要梯度分开存。这样 valid() 之后再 train() 能恢复原状，而 freeze()（把 is_active 也改掉）之后的 train() 不会把冻结解开。实验 valid（debug，Flex），Net 是一个 Linear 加一个 Dropout(0.5)，每行打印权重张量的 is_autodiff、is_require_grad、检查点策略和 dropout 标志：

flex.autodiff().gc() init    autodiff true  require_grad true  checkpointing Some(Balanced)       dropout enabled
  .valid()                   autodiff false require_grad false checkpointing None                 dropout disabled
  .valid().train()           autodiff true  require_grad true  checkpointing Some(Disabled)       dropout enabled
  .freeze()                  autodiff true  require_grad false checkpointing Some(Disabled)       dropout disabled
  .no_grad()                 autodiff true  require_grad false checkpointing Some(Disabled)       dropout enabled
flex init                    autodiff false require_grad false checkpointing None                 dropout enabled
  .train()                   autodiff true  require_grad true  checkpointing Some(Disabled)       dropout enabled
  .valid()                   autodiff false require_grad false checkpointing None                 dropout disabled

（.freeze() 和 .no_grad() 两行是对 .valid().train() 的结果调的。）

  • valid() 去掉 autodiff，train() 加回来，但加回来的是默认策略：Balanced 变成了 Disabled。Module::valid 的文档写明了这一点：检查点策略随 autodiff 一起去掉，train 不恢复，要接着训练就留着原来的模块。
  • freeze() 同时关掉梯度和 dropout；no_grad() 只关梯度，dropout 照常（部分微调时想要的就是这个）。
  • 在普通 Device::flex() 上建的模型，train() 之后就是 autodiff 张量、要梯度：Uninitialized 里记着 is_require_grad = true，is_active 是真，train() 把它兑现了。所以不必一开始就在 .autodiff() 设备上建模型。
  • 第一组里 dropout enabled 出现在普通 Flex 设备上，但这时 Dropout 并不会丢弃任何东西，原因见 nn 模块。
📝NOTE
冷知识：AutodiffModule 合进了 Module

0.21 里 valid() 不在 Module 上，而在单独的 AutodiffModule<B: AutodiffBackend> 上，返回的是另一个类型 Self::InnerModule（内层后端上的模块）。#4717 去掉后端参数后，AutodiffModule 变成了没有参数的 trait，valid 返回 Self，另有一个 from_inner。#5721「refactor(module)!: merge AutodiffModule into Module」（首次发布于 v0.22.0-pre.4）把它删掉，valid 和 train 成了 Module 的必需方法，trait 文档开头加了一句：「Training and validation use the same module type. This trait provides state transitions, not a guarantee that a value currently has autodiff enabled.」这就是上表里同一个 Net 类型能在两种状态之间来回切换的原因；代价是「这个模块能不能 backward」从编译期检查变成了运行期检查。

to_device 与 fork

两者对参数做的事（to_device、fork）：

    fn to_device(mut self, device: &Device) -> Self {
        let reparameterization = self.reparameterization.take();
        let base = self.map_to_device(device, |tensor| tensor.to_device(device));
        // ...
    }

    fn fork(mut self, device: &Device) -> Self {
        let reparameterization = self.reparameterization.take();
        let base = self.map_to_device(device, |tensor| {
            let is_require_grad = tensor.is_require_grad();
            let mut tensor = tensor.to_device(device).detach();

            if is_require_grad {
                tensor = tensor.require_grad();
            }

            tensor
        });
        // ...
    }

Tensor::to_device 在 autodiff 下是一个会被记进计算图的算子，结果是旧叶子的下游节点，不再是叶子；fork 多做一步 detach() 再 require_grad()，让新设备上的参数自己成为叶子。两者都先经过 map_to_device：参数还没初始化、又没有别的克隆共享它（Arc::get_mut 成功）时，只把 Uninitialized 里的目标设备改掉，闭包都不跑；否则才调传进来的搬运函数。改设备时用 with_autodiff_context_from 保留源设备的 autodiff 设置。

📝NOTE
冷知识：懒参数「改地址」是 0.22 加的

#5739「feat(core): a lazy parameter initializes on the device it moves to」（首次发布于 v0.22.0-pre.4）之前，to_device 一个没初始化的参数会先在原设备上初始化再拷过去。提交说明：「A model built lazily and then moved no longer initializes every parameter on its first device, and weights loaded after the move go straight to the new device.」同一个 PR 的后续提交修了一个由此引出的不一致：改地址时如果把目标设备整个存下来，搬到 autodiff 设备上的懒参数会变成可训练，而同样搬过去的已初始化参数不会，「reading a parameter before the move decided whether the move enabled gradients」。修法就是上面那句 with_autodiff_context_from。

派生出来的 to_device / fork 只是逐字段递归到参数这一层。另外两种参数更简单：Param<Flag> 两个方法都原样返回；Param<Tensor<D, Int>> 和 Param<Tensor<D, Bool>> 的 fork 直接调 to_device，源码注释是 // Don't support autodiff.。

❗IMPORTANT
动脑筋：to_device 和 fork 之后各有几个参数拿到梯度？

实验 fork（debug）：在 Device::flex().autodiff() 上建一个 Linear(4, 2)，先读一次 weight 和 bias 让它们初始化，然后分别用 to_device 和 fork 搬到 Device::metal(DeviceKind::DefaultDevice).autodiff()，在 Metal 上做一次前向、sum()、backward()，再 GradientsParams::from_grads(grads, &model).len()。第三组用一个没读过的新 Linear 直接 to_device。三组的 len() 各是几？

动脑筋答案

to_device  device Device<Autodiff { device: Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })), checkpointing: Disabled }>  require_grad false
to_device  GradientsParams::len() = 0
to_device: grad on the original flex weight = Some([4, 2])
fork       device Device<Autodiff { device: Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })), checkpointing: Disabled }>  require_grad true
fork       GradientsParams::len() = 2
lazy+to_device device Device<Autodiff { device: Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })), checkpointing: Disabled }>  require_grad true
lazy+to_device GradientsParams::len() = 2
  • to_device：0。参数已初始化，map_to_device 退回到 tensor.to_device(device)，Metal 上的权重是 Flex 叶子的下游节点，is_require_grad() 是假（只有叶子才是 Grad，见 Autodiff 装饰器）。GradientsParams 的访问者看到 !param.is_require_grad() 直接跳过。梯度去了哪里？程序里另做了一次前向和 backward，拿原来 Flex 模型的权重去取梯度，得到 Some([4, 2])：梯度沿着 to_device 这个算子流回了 Flex 上的旧叶子，而搬过去的模型里已经没有它们。用这样的模型训练，优化器什么也不更新，也不报错。
  • fork：2。detach() 再 require_grad()，Metal 上的权重和偏置是新叶子。
  • lazy+to_device：2。参数没初始化，只是改了地址，第一次前向时直接在 Metal 上建成叶子。同一行代码，结果取决于之前有没有读过参数，这就是 Module::to_device 文档里那句警告的来源：要在目标设备上训练，用 fork。

小结

  • Module 没有后端参数；八个必需方法（visit、map、collect_devices、to_device、fork、train、valid、materialize）由 #[derive(Module)] 逐字段递归生成，num_params、freeze 等都是在 visit / map 上现写的访问者或映射器。派生宏同时生成 Clone 和显示用的 trait。
  • 字段分类只看语法：认得出的基本类型是常量，#[module(skip)] 和常量同待遇，其余都当子模块；0.22 起自定义配置类型必须写 skip，0.22.1 的报错会提示这一点。裸 Tensor 字段不是参数。
  • Param = 随机 u64 的 ParamId + 共享的懒初始化状态 + 单独记着的 is_active。权重在第一次取值时才抽随机数；num_params、devices 不触发初始化。
  • ModuleVisitor / ModuleMapper 按 enter_module 拼出路径，map 保住 id；优化器和梯度整理都建立在这两个 trait 上。
  • valid / train 在同一个类型上切换 autodiff 和 Flag，检查点策略不会恢复；fork 让参数在新设备上成为叶子，to_device 对已初始化的参数不会。

💬 关联讨论

发起讨论
加载中...