对齐 logits

主线位置：推理主线的第三步「前向，logits 对上」。前两步（读 config 和权重、建模块）在这里落到 candle_transformers::models::llama 上，本页把它和 transformers 的 LlamaForCausalLM 一段段对，说明 SmolLM2 的 logits 为什么对得上、哪些地方对不上会静默出错。

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

flowchart TB
  cfg["config.json"] --> lc["LlamaConfig"]
  lc --> c["Config"]
  st["safetensors"] --> vb["VarBuilder"]
  c --> load["Llama load"]
  vb --> load
  load --> emb["wte 词嵌入"]
  emb --> blk["30 个 Block"]
  blk --> nrm["ln_f"]
  nrm --> last["取最后一个位置"]
  last --> head["lm_head"]
  head --> lg["logits 49152"]

实验代码：Python 侧 pt_logits.py、pt_pieces.py、pt_rope.py、pt_keep.py、cfg_probe.py；candle 侧 crate llm-exp 的 bin run、pieces、theta、flash。release 构建，平台 M2 Max、macOS 27、rustc 1.99.0。模型是本地的 SmolLM2-135M-Instruct 权重，两边都载成 f32；PyTorch 固定 torch.set_num_threads(8)。输入是 The capital of France is（5 个 token），对照程序本身见怎么读这一卷，这里不重贴。

结果先摆出来

transformers 的注意力有两种常用实现：eager（逐个算子写出来）和默认的 sdpa（调 torch.nn.functional.scaled_dot_product_attention）。四份输出两两比，最后一个位置的 49152 个 logits（最大 |ref| 20.6），容差 1e-4（依据见怎么读这一卷）：

transformers 换一种注意力实现，和自己差 3.34e-5；candle 和它差的也是这个量级。candle 离 eager 更近（2.98e-5 对 4.29e-5），因为 candle 的 llama 注意力就是「matmul、加 mask、softmax、matmul」四步，和 eager 的写法一样；sdpa 在 CPU 上走 PyTorch 的融合内核，求和顺序不同。下面逐段看两边为什么能对上。

config：candle 只读其中一部分

transformers 用 AutoConfig 把 config.json 读成 LlamaConfig，candle 用 serde 把同一个文件反序列化成自己的 LlamaConfig，再用 into_config 转成模型用的 Config：

from transformers import AutoConfig

cfg = AutoConfig.from_pretrained(model_dir)
print("rope_parameters", cfg.rope_parameters)
print("head_dim", getattr(cfg, "head_dim", None))
print("hidden_act", cfg.hidden_act, "tie", cfg.tie_word_embeddings)

let cfg: LlamaConfig = serde_json::from_slice(&std::fs::read(format!("{model_dir}/config.json"))?)?;
let cfg = cfg.into_config(false); // 参数是 use_flash_attn

rope_parameters {'rope_theta': 100000, 'rope_type': 'default'}
head_dim 64
hidden_act silu tie True

SmolLM2 的 config.json 顶层有 28 个键。candle 的结构体只有 13 个字段，serde 默认忽略多出来的键，不报错也不警告：

SmolLM2 的取值正好都落在 candle 写死的那一边，所以对得上。换一个变体，从源码看会是两种结果：head_dim 不等于 hidden_size / num_attention_heads 时，q_proj 的权重形状和 candle 期望的不同，加载时报错；带偏置的变体，投影用的是 linear_no_bias，文件里的偏置根本不读，加载成功、算出来的数不对。本页只在 SmolLM2 上跑过，这两种没有实测。

transformers 5 存出来的 config 没有 rope_theta

上表里最容易出事的是 rope_theta。SmolLM2 的 config.json 是 transformers 4.42.3 写的，rope_theta 在顶层。transformers 5.22.0 读进来以后把它挪进了 rope_parameters，再 save_pretrained 存回去，顶层就没有了：

cfg = AutoConfig.from_pretrained(model_dir)
cfg.save_pretrained("resaved")
saved = json.load(open("resaved/config.json"))
print({k: v for k, v in saved.items() if "rope" in k})

saved keys with rope: {'rope_interleaved': False, 'rope_parameters': {'rope_theta': 100000, 'rope_type': 'default'}}

candle 0.11.0 的 LlamaConfig 不认识 rope_parameters，rope_theta 走 default_rope 变成 10000。用这份重新存过的 config 跑同一个对照程序：

rope_theta=10000 dtype=F32
logits                   n=49152  max|diff|=6.10e-01
greedy                   n= 20  max|diff|=1.10e+04
worst=1.10e+04 tol=1e-04 MISMATCH

加载和前向都成功，logits 差 0.61，20 个贪心 token 前 18 个相同，最后两个变了（ E、space 变成 northern、 part）。这个 prompt 只有 24 个位置，RoPE 的低频分量还没转开；在一段 249 个 token 的英文上逐位置比 argmax，两种 rope_theta 有 40 个位置（16%）选了不同的 token。猜对下一个 token 的次数是 113 对 111，模型没有「坏掉」，只是安静地变成了另一个模型。

同样的事发生在 Llama 3 的 rope_scaling 上：transformers 5.22.0 把 rope_type: llama3 那一组参数也存进 rope_parameters，candle 读到的 rope_scaling 是 None。candle-transformers 0.11.0 里只有 gemma4 的 config 读 rope_parameters，origin/main 上的 llama.rs 和 0.11.0 相同。搬模型时先看一眼 candle 打印出来的 cfg.rope_theta，或者直接用模型仓库里原始的 config.json。

模块：名字一层层对上

transformers 的类和 candle 的结构体一一对应，VarBuilder::pp 的前缀就是 state_dict 的键名（机制见定义一个模块）：

每一层的残差写法相同。transformers：

residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states, _ = self.self_attn(hidden_states=hidden_states, ...)
hidden_states = residual + hidden_states
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states

candle 的 Block::forward：

        let residual = x;
        let x = self.rms_1.forward(x)?;
        let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?;
        let residual = &x;
        let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?;

MLP 是 SwiGLU，down_proj(silu(gate_proj(x)) * up_proj(x))，candle 的 Mlp::forward 也是这一行。两边哪一层开始有差、差多少，见逐层对数。

加载时有一处和 candle 的一般习惯不同：Llama::load 建 30 个 Block 用的是 .unwrap()，缺一个键不是返回 Err，而是 panic。把 config 的 num_hidden_layers 改成 31：

thread 'main' (9073933) panicked at .../candle-transformers-0.11.0/src/models/llama.rs:524:75:
called `Result::unwrap()` on an `Err` value: cannot find tensor model.layers.30.self_attn.q_proj.weight

退出码 101，调用方的 ? 接不住。词嵌入、model.norm、lm_head 缺了仍然是 Err。缺键、形状不对的其他报错见键名对不上。

RoPE：同一种约定，表只存一半

transformers 的 LlamaRotaryEmbedding 算出 [seq, head_dim] 的 cos / sin（前后两半相同），再用 rotate_half 把第 i 维和第 i + 32 维配成一对：

def rotate_half(x):
  x1 = x[..., : x.shape[-1] // 2]
  x2 = x[..., x.shape[-1] // 2 :]
  return torch.cat((-x2, x1), dim=-1)

q_embed = (q * cos) + (rotate_half(q) * sin)

candle 在 Cache::new 里一次算好 max_position_embeddings 行、head_dim / 2 列的表，apply_rotary_emb 取出 index_pos 开始的几行，交给融合算子 rotary_emb::rope：

        let cos = cache.cos.narrow(0, index_pos, seq_len)?;
        let sin = cache.sin.narrow(0, index_pos, seq_len)?;
        candle_nn::rotary_emb::rope(x, &cos, &sin)

rope 和 rotate_half 是同一种配对（transformers 里叫非交错，config.json 的 rope_interleaved: false）；另一个函数 rope_i 把第 2i 和 2i + 1 维配成一对。用 transformers 的 apply_rotary_pos_emb 对一个随机的 q（[1, 9, 5, 64]）算出参照，candle 用同样的 rope_theta 建表，三个函数各算一遍：

let (cos, sin) = cos_sin(100000.0, 64, 5, &dev)?; // 和 llama.rs 同样的算法
println!("rope      max|diff| {:.3e}", maxdiff(&rope(q, &cos, &sin)?)?);
println!("rope_slow max|diff| {:.3e}", maxdiff(&rope_slow(q, &cos, &sin)?)?);
println!("rope_i    max|diff| {:.3e}", maxdiff(&rope_i(q, &cos, &sin)?)?);
let (c10k, s10k) = cos_sin(10000.0, 64, 5, &dev)?;
println!("rope theta=10000 max|diff| {:.3e}", maxdiff(&rope(q, &c10k, &s10k)?)?);

cos [5, 32]  vs transformers cos[:, :32] max|diff| 1.192e-7
rope      max|diff| 2.384e-7
rope_slow max|diff| 2.384e-7
rope_i    max|diff| 6.356e0
rope theta=10000 max|diff| 7.436e-1

• rope 和 rope_slow（用 narrow、cat 拼出来的参照实现）都和 transformers 差 2.4e-7，只是 f32 舍入；cos 表本身差 1.2e-7，来自 powf 和 PyTorch 的 base ** (arange / dim) 舍入不同。
• rope_i 差 6.36，和 q 本身一个量级：配对方式错了，位置信息完全乱掉。HF 格式的 Llama 权重配 rope。GGUF 文件里的 Llama 配 rope_i，原因和 candle 怎么按架构名选约定见量化 llama 的一次前向。
• rope_theta 错成 10000，这一层就差 0.74，就是上一节那份 config 的效果。

rope 是融合算子，要求输入连续、不记反向（没有 no_grad），所以 llama.rs 在 transpose(1, 2) 之后先调 .contiguous() 再进 apply_rotary_emb。

GQA：9 个查询头共用 3 个 KV 头

SmolLM2 的 num_attention_heads 是 9，num_key_value_heads 是 3，k_proj / v_proj 只输出 3 × 64 = 192 维。算注意力前要把 KV 头复制成 9 个，关键是复制的顺序：查询头 0、1、2 读 KV 头 0，还是 0、3、6 读 KV 头 0。两边各把 [0, 1, 2] 三个头复制 3 倍：

from transformers.models.llama.modeling_llama import repeat_kv

kv = torch.arange(3.0).reshape(1, 3, 1, 1)
print("repeat_kv heads:", repeat_kv(kv, 3).flatten().tolist())

use candle_transformers::utils::repeat_kv;

let kv = Tensor::arange(0f32, 3., &dev)?.reshape((1, 3, 1, 1))?;
println!("repeat_kv heads: {:?}", repeat_kv(kv, 3)?.flatten_all()?.to_vec1::<f32>()?);

repeat_kv heads: [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0]
repeat_kv heads: [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0]

顺序相同，都等于 torch.repeat_interleave(x, 3, dim=1)。实现不同：transformers 是 expand 再 reshape，candle 的 repeat_kv 是在序列维上 cat 3 份再 reshape，源码注释说这样避开了一次带步长的复制。

因果 mask：加法 mask 对 where_cond

transformers 的 mask 由 create_causal_mask 统一生成，形式跟着注意力实现走：

m = create_causal_mask(config=model.config, inputs_embeds=emb, attention_mask=None,
                       past_key_values=None, position_ids=pos)
print(attn, "mask:", m if m is None else (tuple(m.shape), m.dtype))
if m is not None:
  print(m[0, 0])

eager mask: ((1, 1, 5, 5), torch.float32)
tensor([[ 0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00, -3.4028e+38],
        [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00]])
sdpa mask: None

• eager 拿到一个 [1, 1, 5, 5] 的浮点 mask，能看的位置是 0，不能看的是 f32 的最小值，在 matmul 之后加到分数上。
• sdpa 拿到 None：没有 padding 时 transformers 不建 mask，调 scaled_dot_product_attention 时传 is_causal=True，由 PyTorch 的内核自己处理。

candle 的 mask 是 u8，1 表示不能看，由 build_causal_mask 生成：

build_causal_mask(5, 0):
[[0, 1, 1, 1, 1],
 [0, 0, 1, 1, 1],
 [0, 0, 0, 1, 1],
 [0, 0, 0, 0, 1],
 [0, 0, 0, 0, 0]]
Tensor[[5, 5], u8]

CausalSelfAttention::forward 用它把分数换成 -inf，不是加：

            let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
            let att = if seq_len == 1 {
                att
            } else {
                let mask = cache.mask(seq_len, index_pos)?.broadcast_as(att.shape())?;
                masked_fill(&att, &mask, f32::NEG_INFINITY)?
            };
            let att = candle_nn::ops::softmax_last_dim(&att)?;

masked_fill 是这个文件自己定义的，内部是 mask.where_cond(&on_true, on_false)；candle-core 没有 masked_fill（什么时候选 candle）。加一个很大的负数和换成 -inf，softmax 之后这些位置都是 0，结果相同。seq_len == 1（decode 一步）时 candle 不建 mask：只有一个查询，它本来就能看所有已缓存的键。

注意力实现：candle 只有一种

transformers 用 attn_implementation 在加载时选注意力实现；candle 的 llama.rs 只有上面那种「两次 matmul 加 softmax」，不调 candle_nn::ops::sdpa，Metal 上也一样（sdpa 只有 Metal 实现，见什么时候选 candle）。into_config(true) 打开的 use_flash_attn 走的是 CUDA 的 candle-flash-attn，没开 flash-attn feature 时，加载成功，第一次前向 panic：

loaded

thread 'main' (9075567) panicked at .../candle-transformers-0.11.0/src/models/llama.rs:258:5:
not implemented: compile with '--features flash-attn'

不开 feature 时的 flash_attn 函数体就是一个 unimplemented!。transformers 那边写 attn_implementation="flash_attention_2" 而环境里没有这个包，from_pretrained 当场抛 ImportError：「FlashAttention2 has been toggled on, but it cannot be used due to the following error: the package for FlashAttention2 doesn't seem to be installed.」

只算最后一个位置

LlamaForCausalLM.forward 默认对每个位置都算 logits，logits_to_keep 可以只留最后几个；candle 的 Llama::forward 没有选项，固定只取最后一个位置：

with torch.inference_mode():
  full = model(ids).logits
  last = model(ids, logits_to_keep=1).logits
print("logits shape default", tuple(full.shape), "logits_to_keep=1", tuple(last.shape))

logits shape default (1, 5, 49152) logits_to_keep=1 (1, 1, 49152)

        let x = self.ln_f.forward(&x)?;
        let x = x.i((.., seq_len - 1, ..))?.contiguous()?;
        let logits = self.lm_head.forward(&x)?;
        logits.to_dtype(DType::F32)

candle 返回 [1, 49152]，squeeze(0) 之后和 Python 的 logits[0, -1] 对应。生成时这样省掉了 lm_head 对其余位置的计算（49152 × 576 的矩阵乘每个位置一次）。代价是拿不到每个位置的 logits：算一段文本的困惑度、或者像逐层对数那样看中间结果，要自己写前向（从头写 Llama）。

另外 forward 最后总是 to_dtype(DType::F32)：模型用 BF16 / F16 加载时，返回的 logits 也是 F32（半精度）。

动脑筋答案

with torch.inference_mode():
  full = model(ids).logits[0, -1]
  last = model(ids, logits_to_keep=1).logits[0, -1]
  h = model.model(ids).last_hidden_state[0]  # [5, 576]，已过最后的 norm
  w = model.lm_head.weight  # [49152, 576]，和 embed_tokens 共用
  five = (h @ w.T)[-1]
  one = (h[-1:] @ w.T)[0]
print(f"logits[0, -1] vs logits_to_keep=1: max|diff|={(full - last).abs().max().item():.3e}")
print(f"(h @ W.T)[-1] vs h[-1:] @ W.T:     max|diff|={(five - one).abs().max().item():.3e}")
print(f"same storage as embed_tokens: {w.data_ptr() == model.model.embed_tokens.weight.data_ptr()}")

logits[0, -1] vs logits_to_keep=1: max|diff|=2.480e-05
(h @ W.T)[-1] vs h[-1:] @ W.T:     max|diff|=2.480e-05
same storage as embed_tokens: True

• 差 2.48e-5，和 candle 对 transformers 的 2.98e-5 是同一个量级。
• 第二行说明差别全在 lm_head 那一次矩阵乘：同一个隐藏状态，作为 5 行矩阵的最后一行去乘，和单独作为 1 行去乘，BLAS 分块和求和顺序不同，结果就差这么多。
• 所以 candle「只算最后一个位置」和 transformers「算全部再取最后一个」之间，有一部分差别来自矩阵形状本身，不来自哪一边的实现。

小结

• transformers eager 对 sdpa 差 3.34e-5；candle CPU 对 eager 差 2.98e-5、对 sdpa 差 4.29e-5，Metal 3.53e-5 / 3.72e-5，容差 1e-4，20 个贪心 token 全同。
• candle 的 LlamaConfig 只读 13 个键，head_dim、hidden_act、偏置、rope_interleaved 都写死；rope_theta 缺省 10000。transformers 5.22.0 存出来的 config 把它放进 rope_parameters，candle 静默用 10000，logits 差 0.61。
• RoPE 两边都是非交错配对：candle 的 rope 和 transformers 差 2.4e-7，换成 rope_i 差 6.36。GQA 的头复制顺序相同。
• mask：transformers eager 加一个最小值 mask，sdpa 不建 mask 用 is_causal；candle 用 u8 mask 和 where_cond 换成 -inf，单 token 时不建。
• Llama::forward 只返回最后一个位置、总是 F32；Llama::load 缺层 panic；use_flash_attn 没开 feature 时第一次前向 panic。