优化器
主线位置：训练主线的第四、五步。
loss.backward()（求导主线）交回按节点 id 存的Gradients，GradientsParams::from_grads把它换成按ParamId存，ModuleOptimizer::step遍历模块、给每个参数换上新值。一个训练步已经按顺序走过一遍；本页讲这两步内部的结构：为什么分成两层、类型参数去了哪、Adam 的状态放在哪，以及Learner里的同一步和手写的是不是同一回事。
flowchart TB
gp["GradientsParams 按 ParamId"] --> step["ModuleOptimizer step"]
step --> mapper["module.map 遍历"]
mapper --> mf["map_float 每个 Param"]
mf --> has{"有梯度吗"}
has -->|没有| keep["原样返回"]
has -->|有| st["按 ParamId 取状态"]
st --> clip["梯度裁剪"]
clip --> lr["按参数取学习率"]
lr --> dyn["step_dyn 按 rank 分派"]
dyn --> adam["Adam step"]
adam --> back["状态存回 表"]
back --> restore["恢复 autodiff 标记"]
restore --> newp["同一 ParamId 的新 Param"]
burn-optim 里有 SGD、Adam、AdamW、Adagrad、Adafactor、Adan、LAMB、L-BFGS、Lion、Muon、RMSprop（crate 文档的列表），本页以 Adam 为例。实验 crate 是 train-learner（从模板 base 复制，Flex 设备，debug 构建），平台 M2 Max、macOS 27、rustc 1.99.0。
两层：Optimizer 和 ModuleOptimizer
写一个优化器只需要实现逐张量的 Optimizer：
pub trait Optimizer: Send + Sync + Clone + 'static {
/// The state of the optimizer for a single parameter of rank `D`.
///
/// It implements [`RecordState`] (which itself requires `Send + Sync + 'static`) so it can be
/// decomposed into named tensors and scalars for the burnpack format.
type State<const D: usize>: Clone + RecordState;
// ...
fn step<const D: usize>(
&self,
lr: LearningRate,
tensor: Tensor<D>,
grad: Tensor<D>,
state: Option<Self::State<D>>,
) -> (Tensor<D>, Option<Self::State<D>>);
// ...
fn to_device<const D: usize>(state: Self::State<D>, device: &Device) -> Self::State<D>;
}
step 只看一个张量：拿参数、梯度和上一步的状态（第一步是 None），交回新参数和新状态。状态由调用方保管，所以实现者不用管存盘；trait 上的文档注释列出了实现者不用操心的事：缺梯度的参数、记录的读写、模块的结构、张量是否被追踪。这些都归 ModuleOptimizer：
#[derive(Clone)]
pub struct ModuleOptimizer {
optimizers: Vec<OptimizerGroup>,
param_context: HashMap<ParamId, OptimizationContext>,
}
optimizers 是若干参数组，每组一个优化器和一个可选的梯度裁剪；param_context 是每个参数的状态，键是 ParamId。AdamConfig::init() 先 build() 出 Adam，再用 From<O: Optimizer> 包成只有一组（ParamGroup::all()）的 ModuleOptimizer；配置里有梯度裁剪时只有 init 会装上，build 出来的裸 Adam 不带。要让一部分参数用别的优化器，用 with_group 追加一组：一个参数匹配多组时最后加的那组生效，第一组必须匹配全部参数作为兜底，追加时匹配到的参数的已有状态会被清掉。
0.21 的对应物是 SimpleOptimizer<B> 加一个对优化器、模块和后端都泛型的 OptimizerAdaptor<O, M, B>（v0.21.0 的 optim/simple/）。0.22 里张量没有后端参数了（#4717），这一层也跟着改成不泛型：Learner 的字段直接是 optim: ModuleOptimizer，一个优化器对象可以给任何模块用。
类型去了哪：DynOptimizer 和 DynState
ModuleOptimizer 不泛型，但 Adam::State<D> 带着一个常量参数 D，不同秩的参数状态类型不同，没法放进同一个 HashMap。burn 的做法是把状态和秩一起擦掉：DynState 是 Arc<dyn Any + Send + Sync> 加一个 rank: usize；对象安全的 DynOptimizer 收发的是非泛型的 BridgeTensor（见 BridgeTensor 与 bridge）和 DynState。每个 O: Optimizer 都自动实现它，在里面把运行期的秩变回编译期常量：
fn step_dyn(
&self,
rank: usize,
lr: LearningRate,
tensor: BridgeTensor,
grad: BridgeTensor,
state: Option<DynState>,
) -> (BridgeTensor, Option<DynState>) {
dispatch_rank!(rank, D => {
let (tensor, state) = self.step(
lr,
Tensor::<D>::from_bridge(tensor),
Tensor::<D>::from_bridge(grad),
state.map(|state| state.downcast::<O::State<D>>()),
);
(tensor.into_bridge(), state.map(|state| DynState::create(state, D)))
})
}
dispatch_rank! 是一个 match，0 到 8 各一个分支，每个分支里 const D: usize = n; 再展开同一段代码；其他值 panic!("Unsupported tensor rank for optimizer state: {other}")。所以优化器只支持秩不超过 8 的参数。DynState::downcast 用 Arc::try_unwrap：只有一个持有者时直接移出状态，不克隆张量句柄。
实测一个 Param<Tensor<9>> 的模块（train-learner 的 rank9，debug）：GradientsParams 能收到梯度，step 时 panic：
grads: 1
thread 'main' (3303952) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-optim-0.22.1/src/optim/module/base.rs:173:9:
Unsupported tensor rank for optimizer state: 9
step：一次 map 走完整个模块
ModuleOptimizer::step 收一个学习率（impl Into<ModuleLearningRate>，传 f64 就是所有参数同一个值）、模块和 GradientsParams，调 module.map(&mut ModuleOptimizerMapper::new(..))。GradientsParams 是一个 TensorContainer<ParamId>，怎么从 Gradients 换过来见一个训练步；多设备训练时换成 MultiGradientsParams，它从第 id % 设备数 个设备开始轮着找第一个有这个梯度的设备，把其他设备上的同一个梯度搬过去相加，所以同一个参数整个训练期间都在同一个设备上累加。
map 每遇到一个浮点参数就调一次 map_float，进出子模块时 enter_module / exit_module 维护一个路径栈（l1.weight 这样的路径，参数组按它匹配）：
fn map_float<const D: usize>(&mut self, param: Param<Tensor<D>>) -> Param<Tensor<D>> {
let (id, tensor, mapper) = param.consume();
let grad = self.grads.remove(id);
let tensor = if let Some((grad, device)) = grad {
let context = ParameterContext::capture(&tensor);
let entry = self.states.remove_entry(&id);
// ...
let path = self.path.join(".");
let (optim, grad_clipping, existing_dyn_state) = match entry.map(|(_, s)| s) {
Some(OptimizationContext {
optim,
grad_clipping,
state,
..
}) => (optim, grad_clipping, Some(state)),
None => {
let (optim, grad_clipping) =
self.optimizer_from_param(id, Some(path.as_str())).clone();
(optim, grad_clipping, None)
}
};
// ...
let lr = self.lr_module.lr_from_param(id, Some(path.as_str()));
let (tensor, state) = optim.step_dyn(
D,
lr,
tensor.inner().into_bridge(),
clipped_grad.into_bridge(),
existing_dyn_state.map(|s| optim.to_device_dyn(s, &device)),
);
if let Some(state) = state {
self.states.insert(
key.unwrap_or(id),
OptimizationContext {
optim,
path: Some(path),
state,
grad_clipping,
},
);
}
context.restore(Tensor::from_bridge(tensor), id)
} else {
tensor
};
Param::from_mapped_value(id, tensor, mapper)
}
按顺序：
grads.remove(id)取不到梯度的参数原样返回，也不建状态。冻结的参数、这次前向没用到的参数都走这条路。- 先
remove_entry再insert，状态在表里只有一份，不用克隆。第一次见到的参数按路径找参数组，之后用记在OptimizationContext里的那个优化器，不再每步匹配（结构体上的注释：「to avoid matching every step」）。 - 参数被搬到梯度所在的设备，状态用
to_device_dyn也搬过去。 - 更新用的是
tensor.inner()：参数先脱下 autodiff 这一层，在内层后端上算tensor - delta，这一步不进计算图。 ParameterContext在更新前记下require_grad、是否分布式、梯度检查点策略三样，更新后用Tensor::from_inner包回 autodiff 张量、再把三样加回去。Param::from_mapped_value(id, ..)：ParamId不变，下一步才能按同一个 id 找到梯度和状态。
ParameterContext 来自 PR #5618（「fix(optim): keep a parameter's checkpointing strategy across an update」，首次进入 v0.22.0-pre.4，正式版 0.22.0）。之前 map_float 用 Tensor::from_inner 把更新后的参数包回去，这会给它默认的检查点策略。提交说明写道：在 balanced 检查点下训练的模型，第一次优化器 step 之后所有被更新的参数都换成了默认策略，而输入和冻结的参数还是 balanced，下一次前向拒绝合并两种策略，「Training with gradient checkpointing could not get past its first step」。L-BFGS 写回参数时有同样的问题。第一版修复加了一个公开的 from_inner_with_strategy；评审后改成 crate 内部的 ParameterContext，一次带上三样状态，L-BFGS 也改用它，不再对每个写回的参数强加 require_grad，于是冻结的参数经过一步仍然是冻结的。测试 step_keeps_the_gradient_checkpointing_strategy 在 balanced 设备上连走两步。检查点策略本身见梯度检查点。
Adam 的状态存在哪
AdamConfig 的默认值：beta_1 = 0.9、beta_2 = 0.999、epsilon = 1e-5、amsgrad = false，权重衰减和梯度裁剪默认没有。每个参数的状态是 AdamState<D>，里面是 AdaptiveMomentumState<D>：步数 time、一阶矩 moment_1、二阶矩 moment_2，开了 AMSGrad 再加一个 max_moment_2。两个矩和参数同形状，所以 Adam 的状态占两倍参数的内存。transform 更新完矩之后这样算步长：
let time = state.time as i32;
let bias_correction2_sqrt = (1.0 - self.beta_2.powi(time)).sqrt();
let combined_factor = bias_correction2_sqrt / (1.0 - self.beta_1.powi(time));
// ...
let grad = state.moment_1.clone().mul_scalar(combined_factor).div(
v_to_use
.clone()
.sqrt()
.add_scalar(self.epsilon * bias_correction2_sqrt),
);
这是 m̂ / (√v̂ + ε) 上下同乘 √(1 − β₂ᵗ) 的写法：两个偏差修正因子合成一个标量 combined_factor，张量上只剩一次乘、一次开方、一次加、一次除；ε 跟着乘了同一个因子，结果和教科书公式一致。Adam::step 再乘学习率、从参数里减掉。
这些状态在 ModuleOptimizer 的 param_context 里，不在模块里。要存下来得单独走优化器自己的记录：to_record 用 #[derive(RecordState)] 生成的 state_flatten 把每个状态拆成带名字的张量和标量，名字以 ParamId 的十进制值开头，再额外记一个 __rank 标量（load_record 靠它知道该用哪个 D 重建状态）和参数路径（靠它把参数分回正确的组）。save 把这些写成一个 burnpack 文件。手写循环跑完 21 步后存下来，用 burn_pack::Reader 读回，l1.weight 的条目是：
hand-optim.bpk: 8 tensors, 8 scalars
tensor 18402649998147242804.momentum.moment_1
tensor 18402649998147242804.momentum.moment_2
scalar 18402649998147242804.__rank = UInt(2)
scalar 18402649998147242804.momentum.time = UInt(21)
metadata {"18402649998147242804": "l1.weight", "2021826497617271949": "l2.bias", "4872954564708272749": "l1.bias", "708128169117982014": "l2.weight"}
4 个参数张量各两个矩，共 8 个张量；各一个 __rank 和一个 time，共 8 个标量；路径表放在文件的元数据里。键名里的字段路径（momentum.moment_1）就是 AdamState → AdaptiveMomentumState 的嵌套字段名。
问： model.save_file(..) 存下的模型，读回来接着训练，Adam 是不是从头开始？
答： 只读模型的话是。模型文件里只有参数值和每个参数的 ParamId（见存储与加载），Adam 的矩在 ModuleOptimizer 里，要另存 optim.save(..)、读回时 AdamConfig::new().init().load(..)。两边能对上，靠的是加载模型时恢复记录里的 ParamId，而不是用新建模型时随机生成的 id：优化器状态的键就是这个 id。Learner 的 checkpoint 正是这样分三个文件存模型、优化器和学习率调度器，见 Learner 与数据。
学习率调度
LrScheduler 只有 step(&mut self) -> LearningRate 和存取记录两件事；LearningRate 是 f64 的别名，f64 自己实现了这个 trait，返回常数。现成的有常数、阶梯、指数、线性、余弦退火、Noam，以及把几个按顺序或按乘积组合的两种。ModuleLrScheduler 给参数组各配一个调度器，每次 step 交出一个 ModuleLearningRate，map_float 里按 ParamId 和路径查出这个参数的学习率（同样是最后匹配的组生效）。Learner 每次优化器 step 前先调一次调度器的 step；开了梯度累积时，只在真正更新的那一步调。
实验：手写循环和 Learner
train-learner 把同一个两层 MLP（4 → 16 → 1，97 个参数）在同一份合成数据上训练两次：一次手写循环，一次交给 Learner。数据是 200 个样本，y = x · [1, −2, 0.5, 3] + 0.5，x 由一个固定种子的线性同余发生器生成，DataLoaderBuilder 每批 32 个、不打乱，训练 3 个 epoch，每个 epoch 7 批（最后一批 8 个），共 21 步。两次都从同一个 model0.clone() 出发（参数值和 ParamId 都相同），都用 AdamConfig::new().init()，学习率 1e-2。
fn hand_loop(model0: &Mlp, device: &Device) -> (Vec<f32>, Mlp) {
let loader = builder()
.set_device(device.clone())
.build(InMemDataset::new(samples(200, 1)));
let mut model = model0.clone();
let mut optim = AdamConfig::new().init();
let mut losses = Vec::new();
for _epoch in 1..=EPOCHS {
for batch in loader.iter() {
let out = model.forward_regression(batch.unwrap());
losses.push(out.loss.clone().into_scalar());
let grads = GradientsParams::from_grads(out.loss.backward(), &model);
model = optim.step(LR, model, grads);
}
}
optim.save("artifacts/hand-optim.bpk").unwrap();
(losses, model)
}
impl TrainStep for Mlp {
type Input = Batch;
type Output = RegressionOutput;
fn step(&self, batch: Batch) -> TrainOutput<RegressionOutput> {
let out = self.forward_regression(batch);
LEARNER_LOSSES.lock().unwrap().push(out.loss.clone().into_scalar());
TrainOutput::new(self, out.loss.backward(), out)
}
}
Learner 那边每一步的 loss 由 TrainStep::step 推进一个全局 Mutex<Vec<f32>>；SupervisedTraining 的其余配置和输出在 Learner 与数据。
Learner 里的一步拆开看：train_step 调用户的 TrainStep::step，TrainOutput::new 里就是 GradientsParams::from_grads(grads, module)；然后 lr_step 从调度器取学习率（这里是 f64 常数），optimizer_step 调 TrainStep::optimize，它的默认实现是 optim.step(lr_module, self, grads)。和手写循环是同一串调用。
f32::to_bits）相同？两个 DataLoaderBuilder 都加上 .shuffle(7) 再跑呢？Flex、debug、单线程数据加载。先想好再看本页末尾的答案。动脑筋答案
不打乱时（train-learner 不带参数，debug）：
model0: 97 params, l1.weight id 6j5mg0dcb1hvu
hand loop: 21 steps in 103.698ms
...
learner: done in 109.276541ms
step hand learner equal
0 4.47128010 4.47128010 true
1 3.73394918 3.73394918 true
2 6.43713617 6.43713617 true
6 2.71958590 2.71958590 true
13 2.40191126 2.40191126 true
20 1.97708082 1.97708082 true
all 21 losses bit-identical: true
final l1.weight bit-identical: true
trained model id kept: true
21 步全部逐位相同，训练完的 l1.weight 也逐位相同，ParamId 没变。同样的初值、同样顺序的批、同一串算子在同一个后端上，结果是确定的；Learner 只是把手写循环里那几行搬进了 TrainStep 和 Learner 的方法，中间多出来的验证 epoch 用的是 model.valid() 的副本，不碰参数。
都加 .shuffle(7)（train-learner 7）：
step hand learner equal
0 4.40936899 2.40452671 false
1 3.28642321 6.34260845 false
2 4.17087030 4.18700695 false
6 5.02957439 3.63703752 false
13 3.13688731 1.26208472 false
20 3.37988091 1.49062657 false
all 21 losses bit-identical: false
从第 0 步就不同。差别不在优化器，在数据顺序：Learner 的单设备策略开始训练前对训练集的加载器调一次 to_device，而 BatchDataLoader::to_device 用 rng.fork() 派生出一个新的随机数发生器，打乱顺序就和原加载器不一样了。让手写循环也先 loader.to_device(device)（环境变量 HAND_TO_DEVICE=1），两边又逐位相同：
### args: 7, HAND_TO_DEVICE=1
step hand learner equal
0 2.40452671 2.40452671 true
1 6.34260845 6.34260845 true
2 4.18700695 4.18700695 true
6 3.63703752 3.63703752 true
13 1.26208472 1.26208472 true
20 1.49062657 1.49062657 true
all 21 losses bit-identical: true
final l1.weight bit-identical: true
加载器为什么这样派生随机数、多线程加载时批的顺序又怎样，见 Learner 与数据。
小结
- 优化器分两层：
Optimizer只写一个张量的一步，状态由调用方传入传出；ModuleOptimizer负责遍历模块、缺梯度的参数、参数组、梯度裁剪、学习率和状态表。 ModuleOptimizer不泛型：状态擦成DynState（Arc<dyn Any>加秩），DynOptimizer用dispatch_rank!把秩变回常量参数，只支持秩 0 到 8，秩 9 的参数在step时 panic。map_float按ParamId取梯度和状态，在内层后端上更新、不进计算图，ParameterContext把require_grad、分布式标记和检查点策略带回去，ParamId不变。- Adam 的两个矩和步数存在
ModuleOptimizer的表里，键是ParamId；要续训得另存优化器记录，模型加载时恢复ParamId才能对上。 Learner的一步就是TrainStep::step→from_grads→ 调度器step→ModuleOptimizer::step；不打乱数据时和手写循环逐位相同，.shuffle时差在to_device派生的随机数。