RmsNorm 之后梯度是 None

主线位置：求导主线的「每个结果张量记下 Op」一步。融合算子用 apply_op*_no_bwd 调用，结果张量不带 Op，backward 的拓扑排序走到这里就断了（记录规则见计算图怎么记下来，融合算子本身见常用层和自定义算子）。

flowchart TB
  w1["Var w1"] --> l1["Linear"]
  l1 --> rn{"RmsNorm 怎么调"}
  rn -->|forward 连续输入| cut["融合算子 不记 Op"]
  rn -->|forward_diff| keep["拼接 每步记 Op"]
  cut --> none["w1 和 g 无梯度"]
  keep --> some["全部有梯度"]

实验 crate exp-candle/pa-pit，bin p07_fused，debug；M2 Max、macOS 27、rustc 1.99.0。

现象

一个最小的训练：Linear(8 → 8) → RmsNorm → Linear(8 → 1)，三个参数 w1、g（RmsNorm 的权重）、w2 都是 Var，SGD 学习率 0.1，目标由一个结构相同、第一层不同的「老师」网络生成。参数用确定的值初始化，两种写法起点相同：

for step in 0..=1000 {
  let h = l1.forward(&x)?;
  let h = if diff { norm.forward_diff(&h)? } else { norm.forward(&h)? };
  let loss = (l2.forward(&h)? - &y)?.sqr()?.mean_all()?;
  let grads = loss.backward()?;
  // ... print grads.get(&w1).is_some() etc. at step 0, loss every 250 steps
  opt.step(&grads)?;
}

== Cpu
RmsNorm::forward
  grads present: w1 false g false w2 true
  step   0: loss 0.02811
  step 250: loss 0.00075
  step 500: loss 0.00075
  step 750: loss 0.00075
  step 1000: loss 0.00075
  w1 max |change| after 1000 steps: 0
RmsNorm::forward_diff
  grads present: w1 true g true w2 true
  step   0: loss 0.02811
  step 250: loss 0.00068
  step 500: loss 0.00062
  step 750: loss 0.00057
  step 1000: loss 0.00053
  w1 max |change| after 1000 steps: 0.042720333

Metal 上逐行相同（w1 的变化量是 0.04272057）。用 forward 时没有任何报错：backward 成功，优化器成功，loss 也在降，只是降到 0.00075 就停了，因为 RmsNorm 之前的 w1 和 RmsNorm 自己的 g 一次都没动过。

出处

RmsNorm 的 Module::forward：

impl Module for RmsNorm {
    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        if xs.is_contiguous() {
            crate::ops::rms_norm(xs, &self.0.weight, self.0.eps as f32)
        } else {
            self.0.forward(xs)
        }
    }
}

ops::rms_norm 最后一行是 xs.apply_op2_no_bwd(alpha, &RmsNorm { eps })。apply_op2_no_bwd 建的结果张量用 BackpropOp::none()：既不记 xs，也不记 alpha。Linear 的输出总是连续的，所以训练循环里永远走融合路径。

backward 从 loss 往回找，经过 l2 到达 w2，再往前遇到一个没有 Op 的张量，停下。GradStore 里只有 w2。SGD::step 对每个变量 if let Some(grad) = grads.get(var)，拿不到就跳过，不报错。

同样不记反向的融合算子（candle-nn 0.11.0）：ops::softmax_last_dim、ops::rms_norm、ops::layer_norm（LayerNorm::forward 在输入连续、减均值、带偏置时走它）、rotary_emb::rope / rope_i / rope_thd、ops::sdpa。candle-transformers 的模型在推理路径上大量用它们（quantized_llama 用了 rms_norm、rope_i、softmax_last_dim，见量化 llama 的一次前向），拿这些模型代码直接做微调会遇到同样的问题。

修法

训练时用有反向的版本：

上面的输出里，forward_diff 那一组三个参数都有梯度，loss 持续下降。forward_diff 直接调 LayerNorm::forward 的拼接路径，每一步都是普通张量运算，都记 Op。它的文档注释写的是「Faster variant of the forward kernel, this can only be used on contiguous tensors though」，和它的行为正好相反，别被误导（常用层一页有分析）。

带偏置的 LayerNorm 没有 forward_diff。要训练它，就自己写拼接版本，或者调 ops::layer_norm_slow。

同一份模型代码既要推理又要训练时，给层加一个开关，训练时走 _slow / forward_diff。拼接版本在 Metal 上慢 9 到 13 倍（实测见常用层），推理不要用它。

怎么确认没中招

• 第一步之后数一遍：vars.iter().filter(|v| grads.get(v).is_none()).count() 应该是 0。这是最直接的检查，不依赖 loss 曲线。
• loss 下降后停在一个平台、某些参数和初始值完全一样：先怀疑梯度断了，而不是学习率。
• 在模型代码里搜 softmax_last_dim、rms_norm、layer_norm、rope、sdpa，以及 RmsNorm / LayerNorm 的 forward 调用。
• burn 里梯度为 None 的九种情形见 burn 卷的这一页；candle 多了「融合算子」这一种。

小结

• candle-nn 的融合算子（softmax_last_dim、rms_norm、layer_norm、rope、sdpa）都不记反向，梯度在它们那里断掉，没有报错。
• RmsNorm::forward 在连续输入上走融合算子：它前面的参数和它自己的权重都拿不到梯度，SGD 静默跳过。
• 训练用 forward_diff 或 _slow 版本；第一步后检查每个变量都有梯度。