优化器与训练循环
主线位置：求导主线的最后一步「
GradStore→ 优化器」。前面两页讲了图怎么记、backward怎么算出GradStore；本页是candle-nn里的optim.rs和var_map.rs：优化器按Var从GradStore取梯度，用Var::set原地写回。最后用一个手写的训练循环把三段串起来。
sequenceDiagram participant T as 训练循环 participant M as Mlp participant G as GradStore participant O as AdamW participant V as VarMap 里的 Var T->>M: forward x M->>M: Linear Relu Linear 各记 Op T->>T: mse 再 backward T->>G: 按 TensorId 的梯度 T->>O: step grads O->>G: get var O->>V: set 新值 id 不变
本页的代码都在 candle-nn crate（candle-nn/src/{optim.rs,var_map.rs}，加起来约 320 行）。实验是 exp-candle/bp-ad 的 train，release，M2 Max、macOS 27、rustc 1.99.0。
VarMap：按名字存 Var
训练要从「一堆变量」开始。VarMap 是一张加了锁的名字到 Var 的表：
#[derive(Clone)]
pub struct VarMap {
data: Arc<Mutex<HashMap<String, Var>>>,
}
模型不直接碰它，而是经 VarBuilder::from_varmap 拿权重（VarBuilder 的其他后端和 pp 前缀见 VarBuilder）。candle_nn::linear(4, 16, vb.pp("l1")) 要一个名叫 l1.weight、形状 (16, 4) 的张量，最后落到 VarMap::get：
let shape = shape.into();
let mut tensor_data = self.data.lock().unwrap();
if let Some(tensor) = tensor_data.get(path) {
let tensor_shape = tensor.shape();
if &shape != tensor_shape {
candle::bail!("shape mismatch on {path}: {shape:?} <> {tensor_shape:?}")
}
return Ok(tensor.as_tensor().clone());
}
let var = init.var(shape, dtype, device)?;
let tensor = var.as_tensor().clone();
tensor_data.insert(path.to_string(), var);
Ok(tensor)
名字不存在就按初始化方式（linear 的权重是 Kaiming 正态、偏置是均匀分布）新建一个 Var 放进表里；返回给模型的是 var.as_tensor().clone()。这一行是整个训练链路能接上的关键：克隆的是 Arc<Tensor_>，模型里 Linear 持有的权重和 VarMap 里的 Var 是同一个 Tensor_，is_variable 为真，id 也相同。所以：
- 前向时权重是变量，每个算子都记
Op（见计算图怎么记下来）； backward后GradStore里以权重 id 为键的那一项，用Var去get就能取到；- 优化器
set写进的是同一块存储，模型下一步前向直接看到新值。
不需要遍历模型找参数，也不需要在两套 id 之间换算。VarMap 另外能 save / load safetensors 文件，用来存检查点。
Optimizer trait
Optimizer 只有四个必需方法：
pub trait Optimizer: Sized {
type Config: Sized;
fn new(vars: Vec<Var>, config: Self::Config) -> Result<Self>;
fn step(&mut self, grads: &candle::backprop::GradStore) -> Result<()>;
fn learning_rate(&self) -> f64;
fn set_learning_rate(&mut self, lr: f64);
fn empty(config: Self::Config) -> Result<Self> {
Self::new(vec![], config)
}
fn backward_step(&mut self, loss: &Tensor) -> Result<()> {
let grads = loss.backward()?;
self.step(&grads)
}
// ...
}
优化器在构造时拿到一组 Var（通常是 varmap.all_vars()），之后每次 step 收一张 GradStore。backward_step 把 backward 和 step 合成一步。trait 里没有学习率调度器、参数分组、梯度裁剪这些东西，要的话在训练循环里自己调 set_learning_rate、自己改梯度。
SGD
SGD 的文档注释写明「Contrary to the PyTorch implementation of SGD, this version does not support momentum」。构造时滤掉非浮点的变量，step 就是 θ ← θ − lr·g：
fn step(&mut self, grads: &candle::backprop::GradStore) -> Result<()> {
for var in self.vars.iter() {
if let Some(grad) = grads.get(var) {
var.set(&var.sub(&(grad * self.learning_rate)?)?)?;
}
}
Ok(())
}
grads.get(var) 经 Deref 把 &Var 当 &Tensor 用，按 id 查。var.sub(..) 算出一块新存储（var 是变量，这个结果会记一个 Op，但它马上被 set 复制进去后丢掉，图随之释放）。没有梯度的变量（这一步没用到）原样跳过。
AdamW
AdamW 给每个变量配两个同形状的 Var，存一阶矩和二阶矩：
#[derive(Debug)]
struct VarAdamW {
var: Var,
first_moment: Var,
second_moment: Var,
}
矩也用 Var，是为了能用 set 原地更新。step 的核心：
self.step_t += 1;
let lr = self.params.lr;
let lambda = self.params.weight_decay;
let lr_lambda = lr * lambda;
let beta1 = self.params.beta1;
let beta2 = self.params.beta2;
let scale_m = 1f64 / (1f64 - beta1.powi(self.step_t as i32));
let scale_v = 1f64 / (1f64 - beta2.powi(self.step_t as i32));
for var in self.vars.iter() {
let theta = &var.var;
let m = &var.first_moment;
let v = &var.second_moment;
if let Some(g) = grads.get(theta) {
// This involves locking 3 RWLocks per params, if the parameters are large this
// should not be an issue but this may be problematic with models with lots of
// small parameters.
let next_m = ((m.as_tensor() * beta1)? + (g * (1.0 - beta1))?)?;
let next_v = ((v.as_tensor() * beta2)? + (g.sqr()? * (1.0 - beta2))?)?;
let m_hat = (&next_m * scale_m)?;
let v_hat = (&next_v * scale_v)?;
let next_theta = (theta.as_tensor() * (1f64 - lr_lambda))?;
let adjusted_grad = (m_hat / (v_hat.sqrt()? + self.params.eps)?)?;
let next_theta = (next_theta - (adjusted_grad * lr)?)?;
m.set(&next_m)?;
v.set(&next_v)?;
theta.set(&next_theta)?;
}
}
- 权重衰减是解耦的（「W」）：先
θ·(1 − lr·λ)，再减去 Adam 的更新，不把λθ加进梯度。 - 偏差修正
scale_m、scale_v每步按step_t算一次，所有变量共用。 - 每个变量一串十来个张量算子，每个都是一次独立的内核调用（candle 不融合）。注释里担心的是另一笔开销：每个参数要锁三把
RwLock，参数很多、每个很小时会显眼。 ParamsAdamW的默认值是lr = 0.001、beta1 = 0.9、beta2 = 0.999、eps = 1e-8、weight_decay = 0.01。想要不带衰减的 Adam，把weight_decay设成 0。
一个手写的训练循环
两层 MLP（4 → 16 → 1，ReLU），合成回归数据 y = x·w + 0.5，64 个样本、4 个特征，和 burn 卷的一个训练步同一个设定。CPU 上 Device::set_seed 不可用（见下文），所以 x 用一个小的线性同余发生器生成，保证每个设备、每次运行都一样：
use candle_core::{DType, Device, Module, Tensor};
use candle_nn::{linear, loss, AdamW, Linear, Optimizer, ParamsAdamW, VarBuilder, VarMap, SGD};
struct Mlp {
l1: Linear,
l2: Linear,
}
impl Mlp {
fn new(vb: VarBuilder) -> candle_core::Result<Self> {
Ok(Self { l1: linear(4, 16, vb.pp("l1"))?, l2: linear(16, 1, vb.pp("l2"))? })
}
fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
self.l2.forward(&self.l1.forward(x)?.relu()?)
}
}
fn train(label: &str, dev: &Device, use_adamw: bool, steps: usize) -> anyhow::Result<()> {
if let Err(e) = dev.set_seed(42) {
println!("set_seed: {e}");
}
// ... x: (64, 4)，由线性同余发生器生成
let w = Tensor::new(&[[1f32], [-2.], [0.5], [3.]], dev)?;
let y = (x.matmul(&w)? + 0.5)?;
let varmap = VarMap::new();
let vb = VarBuilder::from_varmap(&varmap, DType::F32, dev);
let model = Mlp::new(vb)?;
// ... 打印 varmap 里的名字和标量总数
let vars = varmap.all_vars();
let mut adamw = AdamW::new(vars.clone(), ParamsAdamW { lr: 1e-2, weight_decay: 0.0, ..Default::default() })?;
let mut sgd = SGD::new(vars, 5e-2)?;
for step in 0..steps {
let pred = model.forward(&x)?;
let loss = loss::mse(&pred, &y)?;
let grads = loss.backward()?;
let n_grads = grads.get_ids().count();
if use_adamw {
adamw.step(&grads)?;
} else {
sgd.step(&grads)?;
}
let loss_value = loss.to_scalar::<f32>()?;
// ... 打印 step、loss、n_grads 和这一步的耗时
}
dev.synchronize()?;
// ...
Ok(())
}
main 依次跑 CPU + AdamW、CPU + SGD、Metal + AdamW，再在同一个进程里把 Metal + AdamW 跑一遍。模型没有 #[derive]，也没有参数列表：linear 在建层时把变量放进 VarMap，优化器从 all_vars() 拿到同样的四个 Var。candle_nn::loss::mse 是 (inp - target)?.sqr()?.mean_all()。
grads（GradStore 里的条目数）是几？模型有四个变量：两个权重、两个偏置。提示：上一页「GradStore」一节里，普通张量 k 的梯度去哪了？
输出
release，200 步：
set_seed: cannot seed the CPU rng with set_seed
cpu adamw: vars = ["l1.bias", "l1.weight", "l2.bias", "l2.weight"], params = 97
step 0 loss 5.301529 grads 6 2.056083ms
step 1 loss 5.108586 grads 6 54.292µs
step 2 loss 4.924667 grads 6 46.167µs
step 49 loss 0.165606 grads 6 36.167µs
step 99 loss 0.033818 grads 6 36.25µs
step 149 loss 0.015550 grads 6 33.833µs
step 199 loss 0.010279 grads 6 36.834µs
200 steps in 8.920625ms
set_seed: cannot seed the CPU rng with set_seed
cpu sgd : vars = ["l1.bias", "l1.weight", "l2.bias", "l2.weight"], params = 97
step 0 loss 3.743113 grads 6 25.75µs
step 1 loss 2.871971 grads 6 27.5µs
step 2 loss 2.318908 grads 6 28.125µs
step 49 loss 0.057856 grads 6 23.166µs
step 99 loss 0.021361 grads 6 23.25µs
step 149 loss 0.014078 grads 6 22.959µs
step 199 loss 0.011044 grads 6 22.708µs
200 steps in 4.687ms
metal adamw: vars = ["l1.bias", "l1.weight", "l2.bias", "l2.weight"], params = 97
step 0 loss 8.175325 grads 6 355.462709ms
step 1 loss 7.785774 grads 6 1.668584ms
step 2 loss 7.411994 grads 6 1.666541ms
step 49 loss 0.255462 grads 6 1.2515ms
step 99 loss 0.081451 grads 6 1.449167ms
step 149 loss 0.028029 grads 6 1.539542ms
step 199 loss 0.016447 grads 6 1.492667ms
200 steps in 673.329708ms
metal adamw (again): vars = ["l1.bias", "l1.weight", "l2.bias", "l2.weight"], params = 97
step 0 loss 8.175325 grads 6 1.539333ms
step 1 loss 7.785774 grads 6 1.554416ms
step 2 loss 7.411994 grads 6 1.613583ms
step 49 loss 0.255462 grads 6 1.627208ms
step 99 loss 0.081451 grads 6 1.583959ms
step 149 loss 0.028029 grads 6 1.573667ms
step 199 loss 0.016447 grads 6 1.555042ms
200 steps in 329.617583ms
- 三组都收敛：loss 从 3.7 到 8.2 之间降到 0.01 到 0.02。
params = 97：l1的权重 16×4 加偏置 16，l2的权重 1×16 加偏置 1。 - CPU 上不能设种子。
Device::Cpu的set_seed直接返回错误（CpuDevice::set_seed里是一句bail!("cannot seed the CPU rng with set_seed")），CPU 上的rand_uniform/rand_normal每次都用rand::rng()。所以 CPU 上的初始权重每次运行都不同：第 0 步的 loss 这次是 5.301529，再跑一次是 5.283688。Metal 上两次训练打印出来的每一步 loss 都相同，但这不是因为set_seed(42)起了作用：每个新的 Metal 设备都从同一个初始状态开始，而 2^32 以下的种子值实际上被忽略（set_seed(42)、set_seed(7)生成的随机数和不设一样），见取样里讲 Metalset_seed的一节。 - 每步的时间（含每步一次
to_scalar读回 loss）：CPU 上 AdamW 约 36 µs、SGD 约 23 µs，AdamW 每个变量要多做十来个小算子。Metal 上每步约 1.5 ms：这么小的模型，时间几乎全是逐个提交内核和每步一次等待 GPU 的开销。 - Metal 第一步慢。 第一次用到的每个内核都要先建计算管线（见 Metal 后端）。同一个进程里第二次训练，第 0 步就是 1.5 ms。不同进程之间第一步差别很大：上面这次 355 ms，之后三次分别是 20.8 ms、9.9 s、22.1 ms（原因本页没有追究）。
和 burn 的训练主线比
burn 的一个训练步要多两步换算：backward 得到按计算图节点 id 存的 Gradients，GradientsParams::from_grads 用派生出来的模块遍历把它换成按 ParamId 存的梯度；ModuleOptimizer::step 再 map 一遍模块，按 ParamId 取梯度和 Adam 状态，返回换上新参数的新模块（参数 id 不变）。candle 两步都省掉了：模型里的权重就是 Var 里的那个张量，张量 id 就是梯度的键，所以不需要第二套 id；优化器手里直接有 Vec<Var>，用 set 原地改，模型对象从头到尾是同一个，不被消耗也不被重建。代价是 candle 里没有 Module 派生那样的结构化遍历：哪些张量是参数，全凭它们是不是经 VarMap 建出来的 Var；优化器状态也只是跟着 Var 存在 Vec 里，没有按名字或 id 的状态表。两边在同一个 MLP 上的计时对照放在设计取舍。
动脑筋答案
是 6，不是 4。在训练循环第 0 步加一行，检查输入 x 和目标 y（都是普通张量）：
step 0: x.grad is_some=true y.grad is_some=true
step 0 loss 5.283688 grads 6 337.417µs
- 四个变量各一项。
x一项：第一层Linear::forward对二维输入是x.matmul(&self.weight.t()?)，Matmul的规则对两边都or_insert，x不在排序列表里，它的梯度一直留在表里。y一项：mse里的inp - target，Sub的规则同样给target一份梯度。- 中间结果（转置后的权重、广播后的偏置、ReLU 的输出……）的梯度在
backward里都被取走了。优化器只按自己的Var去查，多出来的两项不影响训练，只是每步多算了两个没人要的梯度（这里是一次 64×16 乘 16×4 的矩阵乘和一次逐元素取负）。
小结
VarMap是名字到Var的表；经VarBuilder::from_varmap建层时，模型拿到的权重和Var是同一个Tensor_，所以前向记图、GradStore按 id 查、set写回都对得上。Optimizertrait 只有new、step、读写学习率；SGD没有动量，θ ← θ − lr·g；AdamW每个变量配两个Var存矩，权重衰减是解耦的，默认weight_decay = 0.01。- 一步训练 = 前向（每个算子记
Op）→loss.backward()→optimizer.step(&grads)（按Var取梯度，set原地写回）。没有节点 id 到参数 id 的换算，模型不被重建。 - CPU 上
set_seed返回错误、初始化不可复现；Metal 每次从同一个状态开始所以可复现，但小于 2^32 的种子值不起作用。这个小 MLP 在 CPU 上每步 23 到 36 µs，Metal 上约 1.5 ms，Metal 的第一步要编译内核。