常用层

主线位置：推理主线的第二步「建模块」：transformers 模型里出现最多的那几种层，在 candle-nn 里各是什么、权重怎么摆、激活函数的名字怎么对。

版本：PyTorch 2.15.0 / transformers 5.22.0（CPU）、candle 0.11.0（CPU，accelerate；Metal）

flowchart TB
  sd["PyTorch state_dict"] --> st["safetensors"]
  st --> vb["VarBuilder"]
  vb --> lin["linear in out"]
  vb --> emb["embedding"]
  vb --> nrm["layer_norm rms_norm"]
  vb --> cnv["conv1d conv2d"]
  cfg["hidden_act 字符串"] -->|serde| act["Activation"]
  lin --> fwd["forward"]
  act --> fwd
  flag["train 参数"] --> drop["Dropout"]
  drop --> fwd

实验代码：PyTorch 侧 layers.py、acts.py、bert_ref.py、dropout.py；candle 侧 crate module-exp 的 bin layers、acts、bert_gelu、hidden_act、ce_dropout（candle 0.11.0，features metal、accelerate，release）。BERT 用本地的 bert-base-chinese 权重。平台 M2 Max、macOS 27、rustc 1.99.0。层内部怎么实现（Linear 的三维分支、融合的 RmsNorm 不记反向等），见 candle 卷的常用层，本页只看从 PyTorch 搬过来要对什么。

权重布局：和 PyTorch 一样

PyTorch 侧建六个层，LayerNorm / RMSNorm 的权重换成随机数（默认的 1 和 0 看不出读没读对），state_dict 加上前缀写成一个 safetensors 文件：

layers = {
  "lin": nn.Linear(4, 3),
  "emb": nn.Embedding(10, 4),
  "ln": nn.LayerNorm(4),
  "rms": nn.RMSNorm(4, eps=1e-5),
  "c1": nn.Conv1d(2, 3, kernel_size=3, padding=1),
  "c2": nn.Conv2d(2, 3, kernel_size=3, padding=1),
}

lin.weight [3, 4]
lin.bias   [3]
emb.weight [10, 4]
ln.weight [4]
ln.bias   [4]
rms.weight [4]
c1.weight [3, 2, 3]
c1.bias   [3]
c2.weight [3, 2, 3, 3]
c2.bias   [3]

candle 侧按 PyTorch 构造函数的参数顺序建同样的层，读同一个文件：

let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[path], DType::F32, &dev)? };
let lin = linear(4, 3, vb.pp("lin"))?;
let emb = embedding(10, 4, vb.pp("emb"))?;
let ln = layer_norm(4, 1e-5, vb.pp("ln"))?;
let rms = rms_norm(4, 1e-5, vb.pp("rms"))?;
let c1 = conv1d(2, 3, 3, Conv1dConfig { padding: 1, ..Default::default() }, vb.pp("c1"))?;
let c2 = conv2d(2, 3, 3, Conv2dConfig { padding: 1, ..Default::default() }, vb.pp("c2"))?;

weights: lin [3, 4] emb [10, 4] ln [4] rms [4] c1 [3, 2, 3] c2 [3, 2, 3, 3]
[Cpu] lin in [2, 5, 4] F32 -> [2, 5, 3] max_abs_diff=0e0
[Cpu] emb in [2, 3] I64 -> [2, 3, 4] max_abs_diff=0e0
[Cpu] ln  in [2, 5, 4] F32 -> [2, 5, 4] max_abs_diff=2.3841858e-7
[Cpu] rms in [2, 5, 4] F32 -> [2, 5, 4] max_abs_diff=1.1920929e-7
[Cpu] c1  in [2, 2, 7] F32 -> [2, 3, 7] max_abs_diff=2.3841858e-7
[Cpu] c2  in [2, 2, 5, 6] F32 -> [2, 3, 5, 6] max_abs_diff=2.3841858e-7

Metal 上六个差分别是 0、0、2.38e-7、2.38e-7、2.38e-7、1.19e-7。都是 f32 舍入级别的差，没有一个需要转置或改名：

Linear 的 [out, in] 是关键：candle 前向算 x @ w.t()，和 PyTorch 一样，所以 Hugging Face 的权重原样装进来。burn 的 Linear 权重是 [in, out]，从 PyTorch 过去要转置一次（定义模型）。

几处 PyTorch 有默认、candle 要写出来的地方：

• eps 没有默认值。 nn.LayerNorm 默认 1e-5；nn.RMSNorm 默认 None，也就是按输入 dtype 取 torch.finfo(dtype).eps。candle 的 layer_norm / rms_norm 都要显式传。搬 transformers 模型时从 config.json 读（SmolLM2 是 rms_norm_eps: 1e-05，bert-base-chinese 是 layer_norm_eps: 1e-12）。eps 写错了会怎样，见逐层对数。

• 卷积的配置放进一个结构体。 padding、stride、dilation、groups 是 Conv1dConfig 的字段，padding 只有补零一种（没有 padding_mode），也没有 padding="same"。

• 没有不带批维的输入。 PyTorch 的 Conv1d 收 (C, L) 会当成一个样本，得 [3, 7]；candle 报错：

  Copy

  conv1d on [2, 7]: Err(unexpected rank, expected: 3, got: 2 ([2, 7]))
  

• Embedding 没有 padding_idx、max_norm。 前向就是 index_select，下标用 U32 / I64 都行（上面用的是 PyTorch 存下来的 I64）。

输入要连续

Linear 收一个切片视图时，PyTorch 照算，candle 可能报错。(1000, 1000) 的张量切出前 4 列交给 linear(4, 3)：

let x4 = x.narrow(1, 0, 4)?;
match x4.apply_t(&lin, true) { /* ... */ }
let y = x4.contiguous()?.apply_t(&lin, true)?;

narrowed [1000, 4] stride [1000, 1] apply_t(&linear, true): Err(MatMulUnexpectedStriding { lhs_l: Layout { shape: [1000, 4], stride: [1000, 1], start_offset: 0 }, rhs_l: Layout { shape: [4, 3], stride: [1, 4], start_offset: 0 }, bmnk: (1, 1000, 3, 4), msg: "non-contiguous lhs" })
after contiguous(): [1000, 3]

Metal 上是 Invalid matmul arguments [1000, 1] [1, 4] (1000, 3, 4)。二维输入走 Linear::forward 的最后一个分支，原样交给 matmul；开 accelerate 的 CPU 和 Metal 的 matmul 只收行步长等于列数（或转置）的布局。从融合的 QKV 输出里切出一段再进 Linear，就是这个形状。先 contiguous()。为什么、哪些步长不行，见形状与广播和 candle 卷的常见坑 Invalid matmul arguments。

激活函数：gelu 的两个版本

PyTorch 的 F.gelu(x) 默认是精确的 erf 版本，F.gelu(x, approximate="tanh") 是 tanh 近似。candle 的 Tensor 上也有两个方法，但名字反过来：x.gelu() 是 tanh 近似（Gelu：0.5 * v * (1.0 + tanh(SQRT_TWO_OVER_PI * v * (1.0 + 0.044715 * v * v)))），x.gelu_erf() 才是精确版（GeluErf）。开 accelerate 时 f32 的 gelu 走 vs_gelu，公式相同，tanh 交给 vForce。

在 [-8, 8] 上取 16001 个点，两边四种组合都算：

x = torch.linspace(-8, 8, 16001, dtype=torch.float32)
F.gelu(x)
F.gelu(x, approximate="tanh")

let x = refs.get("x").context("x")?.to_device(&dev)?;
for (label, y) in [("x.gelu()", x.gelu()?), ("x.gelu_erf()", x.gelu_erf()?)] {
  // max |y - F.gelu|, max |y - F.gelu(tanh)|, and against the f64 versions
}

f64 erf vs tanh: max 4.7324e-04 at x=-2.6990
f32 erf vs tanh: max 4.7350e-04 at x=2.6850
F.gelu vs f64: 4.5907997225924646e-07 F.gelu tanh vs f64: 4.252386354153259e-07

[Cpu] x.gelu()      vs F.gelu 4.735e-4 (x=2.685); vs F.gelu(tanh) 4.768e-7 (x=5.062); vs f64 erf 4.733e-4, f64 tanh 4.768e-7
[Cpu] x.gelu_erf()  vs F.gelu 4.768e-7 (x=2.045); vs F.gelu(tanh) 4.733e-4 (x=-2.694); vs f64 erf 4.768e-7, f64 tanh 4.733e-4
[Metal(MetalDevice(DeviceId(1)))] x.gelu()      vs F.gelu 4.735e-4 (x=2.685); vs F.gelu(tanh) 4.768e-7 (x=4.021); vs f64 erf 4.733e-4, f64 tanh 4.768e-7
[Metal(MetalDevice(DeviceId(1)))] x.gelu_erf()  vs F.gelu 2.384e-7 (x=2.168); vs F.gelu(tanh) 4.735e-4 (x=-2.694); vs f64 erf 4.768e-7, f64 tanh 4.735e-4

配对对了，两边差 2.4e-7 到 4.8e-7，是 f32 的舍入；配错了，差 4.73e-4，就是两个公式在 x ≈ ±2.7 处本身的差（f64 里算也是 4.7324e-4）。所以 PyTorch 里的 F.gelu(x) 搬过来要写 x.gelu_erf()，F.gelu(x, approximate="tanh") 才写 x.gelu()。

单个元素差 4.7e-4 看着不大，进了模型会放大。bert-base-chinese 的 config.json 写的是 "hidden_act": "gelu"（erf）。两句中文（补齐到 18 个 token）跑 transformers 的 BertModel，再在 candle 里分别用 HiddenAct::Gelu（gelu_erf）和 HiddenAct::GeluApproximate（gelu），比 last_hidden_state（只比非填充位置）：

[Cpu] candle Gelu vs transformers gelu: max_abs_diff=1.7166138e-5 (max |ref|=12.881842)
[Cpu] candle GeluApproximate vs transformers gelu: max_abs_diff=1.9316673e-2 (max |ref|=12.881842)
[Cpu] candle GeluApproximate vs transformers gelu_pytorch_tanh: max_abs_diff=1.6212463e-5 (max |ref|=12.880053)

Metal 上三行是 1.44e-5、1.93e-2、1.81e-5。激活对上时差 1.7e-5；用错一个激活，12 层之后差 1.93e-2，是容差 1e-4 的近两百倍，和 transformers 那边把 hidden_act 改成 gelu_pytorch_tanh 之后的差（1.93e-2）一样大。第三行说明 candle 的 tanh 版和 transformers 的 tanh 版是对得上的，差别全在公式。tanh 版在 CPU 上快一倍，candle-examples 的 bert 有 --approximate-gelu 开关，速度和句向量的影响见 candle 卷的 BERT。

hidden_act 字符串对 Activation

transformers 按 config.json 的 hidden_act 字符串查 ACT2FN 表。candle 的 Activation 是一个 serde::Deserialize 枚举（rename_all = "lowercase"，另有几个 alias），同一个字符串能直接反序列化成它。把 transformers 5.22.0 ACT2CLS 里的 24 个名字逐个交给 serde_json，能解析的再和 ACT2FN[name] 的输出比：

match serde_json::from_str::<Activation>(&format!("\"{name}\"")) {
  Err(e) => println!("{name:18} -> Err({e})"),
  Ok(act) => { /* act.forward(&x) vs ACT2FN[name](x) */ }
}

gelu               -> Gelu: max diff vs transformers 4.768e-7 (at x=2.045)
gelu_10            -> Err(unknown variant `gelu_10`, expected one of `gelu`, `gelu_new`, `newgelu`, `relu`, `relu2`, `relu6`, `silu`, `sigmoid`, `hardsigmoid`, `swiglu`, `swish`, `mish`, `hardswish`, `elu`, `leakyrelu`, `gelu_pytorch_tanh`, `gelupytorchtanh` at line 1 column 9)
gelu_new           -> NewGelu: max diff vs transformers 4.768e-7 (at x=5.062)
gelu_pytorch_tanh  -> GeluPytorchTanh: max diff vs transformers 4.768e-7 (at x=5.062)
...
silu               -> Silu: max diff vs transformers 4.768e-7 (at x=2.836)
...
swish              -> Swish: max diff vs transformers 4.768e-7 (at x=2.385)
tanh               -> Err(unknown variant `tanh`, expected one of ...)

（CPU；省略的行见下表，Err 行里 expected one of 之后的列表都一样。）

• 能解析的 11 个语义都对：gelu → gelu_erf，gelu_new 和 gelu_pytorch_tanh → gelu（tanh），swish 在 transformers 里就是 nn.SiLU，candle 的 Swish 是 x * sigmoid(x)，同一个函数。
• 解析不了的会让整个 config 反序列化失败，而不是静默换一个函数。其中 quick_gelu 是 transformers 里 CLIPTextConfig 和 CLIPVisionConfig 的默认值。
• leaky_relu 和 elu 在 candle 里带参数：只认 {"leakyrelu": 0.01}、{"elu": 1.0} 这样的写法，字符串 "leakyrelu" 报 invalid type: unit variant, expected newtype variant。transformers 的 leaky_relu 用的是 nn.LeakyReLU 的默认斜率 0.01。
• 这张表只对 candle_nn::Activation 有效。candle-transformers 的模型多半有自己的枚举：bert 的 HiddenAct 只有 gelu、geluapproximate、relu 三个，一个写着 "gelu_new" 的 BERT 系 config 报 unknown variant `gelu_new`, expected one of `gelu`, `geluapproximate`, `relu` 。

Dropout：train 是参数，不是模式

PyTorch 的 nn.Dropout 看模块上的 training 标志，由 model.train() / model.eval() 整棵树一起切换。新建的模块默认是训练模式，from_pretrained 出来的是推理模式：

fresh nn.Dropout training: True
Dropout(0.5) train(): zero fraction 0.4998, max 2.0
Dropout(0.5) eval(): zero fraction 0.0000, max 1.0
Dropout(1.0) train: 0.0
BertModel.from_pretrained training: False dropout p 0.1

candle 没有这个标志。Dropout::forward 每次调用收一个 train: bool：

let drop = Dropout::new(0.5);
for train in [true, false] {
  let y = drop.forward(&x, train)?;
  // fraction of zeros, max
}

[Cpu] Dropout(0.5) train=true: zero fraction 0.5003, max 2
[Cpu] Dropout(0.5) train=false: zero fraction 0.0000, max 1
apply_t(&drop, false) == x: true
Dropout(1.0) train=true: Err(dropout probability has to be in [0, 1), got 1)

Metal 上 train=true 的零占比是 0.4996。行为和 PyTorch 一样（以 1/(1-p) 放大留下的元素），区别在接口：

• Dropout 实现的是 ModuleT（forward_t(&self, xs, train)），不是 Module。照着别的层写 x.apply(&drop) 编译不过：

  Copy

  error[E0277]: the trait bound `Dropout: Module` is not satisfied
      --> src/bin/ce_dropout.rs:7:19
       |
     7 |   let y = x.apply(&drop)?;
       |             ----- ^^^^^ the trait `for<'a> Fn(&'a candle_core::Tensor)` is not implemented for `Dropout`
       |             |
       |             required by a bound introduced by this call
  ...
  note: required by a bound in `candle_core::Tensor::apply`
  

  要写 x.apply_t(&drop, train)。反过来，任何 Module 都自动是一个忽略 train 的 ModuleT（impl<M: Module> ModuleT for M），所以一个要区分训练和推理的模块，forward 写成 forward_t(&self, x, train)，把 train 一路传给子层。candle-nn 里实现 ModuleT 的只有 Dropout、BatchNorm 和 FuncT。

• p 的范围是 [0, 1)。PyTorch 允许 p=1.0（输出全 0），candle 返回 Err。

• 推理时传 false 就是原样返回（xs.clone()，只多一个引用）。candle-transformers 的 llama 里没有 Dropout；bert 有一个自己的 Dropout，forward 里只有 // TODO 和 Ok(x.clone())，拿它微调就等于没有 dropout。

和 burn 对照一句：burn 也没有 eval() 标志，dropout 看输入张量在不在 autodiff 设备上（训练和推理模式）；candle 是调用方每次明说。

小结

• Linear、Embedding、LayerNorm、RmsNorm、Conv1d/2d 的构造参数顺序、权重键名和形状都和 PyTorch 相同，Linear 是 [out, in]，读 state_dict 不用转置（burn 要转）；六个层前向和 PyTorch 差 ≤ 2.4e-7。
• eps 要显式传，卷积没有 padding_mode 和不带批维的输入，Embedding 没有 padding_idx；切片视图进 Linear 在 accelerate / Metal 上报 matmul 步长错误，先 contiguous()。
• x.gelu() 是 tanh 近似，F.gelu 默认的 erf 版本是 x.gelu_erf()。配错时单点差 4.73e-4，bert-base-chinese 的隐藏状态差 1.93e-2；配对时差 1.7e-5。
• config.json 的 hidden_act 能直接反序列化成 candle_nn::Activation：24 个 transformers 名字里 11 个能解析、语义都对，13 个（quick_gelu、gelu_fast、tanh、leaky_relu…）报 unknown variant；模型自己的枚举（bert 的 HiddenAct）认的更少。
• Dropout 是 ModuleT，每次调用传 train；x.apply(&drop) 是 E0277，p=1.0 返回 Err；没有 model.eval()，train 要一路传下去。