Rust in Everything
基于 candle-core 0.11.0

cannot find tensor

主线位置：推理主线 safetensors 那条路的第一步，VarBuilder::from_mmaped_safetensors 之后模型的 load 按名字一个个取权重。pp 怎么拼前缀、get 怎么走到后端，见 VarBuilder；文件里的名字从哪来，见 safetensors 与 mmap。本页的例子在 candle-nn 和 candle-transformers 里。

flowchart TB
  load["Llama::load"] --> name["pp 拼出完整名字"]
  name --> look{"文件里有吗"}
  look -->|没有| e1["cannot find tensor 名字"]
  look -->|有| shape{"形状对吗"}
  shape -->|不对| e2["shape mismatch for 名字"]
  e2 --> blocks{"在 blocks 里"}
  blocks -->|是| boom["unwrap panic"]
  blocks -->|否| err["Err"]
  shape -->|对| ok["取到"]

实验 crate exp-candle/pb-pit，bin vbname，SmolLM2-135M 的 model.safetensors（BF16，按 F32 加载），CPU，debug，M2 Max、macOS 27、rustc 1.99.0。

现象

四种常见的写错法，每种单独一个进程：

// prefix: own layer code, forgot the "model." prefix
let attn = vb()?.pp("layers").pp(0).pp("self_attn");
let q = attn.get((576, 576), "q_proj.weight")?;

// double: Llama::load already adds "model.", the caller adds it again
let _ = Llama::load(vb()?.pp("model"), &cfg)?;

// tie: config written by hand, tie_word_embeddings left out
let mut c = config()?;
c.tie_word_embeddings = None;
let _ = Llama::load(vb()?, &c.into_config(false))?;

// kvheads: config written by hand, num_key_value_heads left out (GQA model)
let mut c = config()?;
c.num_key_value_heads = None;
let _ = Llama::load(vb()?, &c.into_config(false))?;

--- prefix
Error: cannot find tensor layers.0.self_attn.q_proj.weight
exit=1
--- double
Error: cannot find tensor model.model.embed_tokens.weight
exit=1
--- tie
Error: cannot find tensor lm_head.weight
exit=1
--- kvheads

thread 'main' (3936098) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-transformers-0.11.0/src/models/llama.rs:524:75:
called `Result::unwrap()` on an `Err` value: shape mismatch for model.layers.0.self_attn.k_proj.weight, expected: [576, 576], got: [192, 576]
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit=101
  • 前三种报的都是「找不到某个名字」。名字是拼好的完整路径，第二种多了一个 model.，一眼能看出；第三种的 lm_head.weight 在文件里本来就没有，问题不在名字而在配置。
  • 第四种是 panic：SmolLM2 用分组查询注意力（GQA），9 个查询头只有 3 个键值头，k_proj 是 [192, 576]；配置里漏了 num_key_value_heads，LlamaConfig::num_key_value_heads 回退成 num_attention_heads，按 [576, 576] 去要。错误本身是正常的 Err，被 unwrap 成了 panic。

出处

VarBuilder 的 path 把 pp 压进去的每一段用 . 连起来，再接上权重名；MmapedSafetensors::get 按这个完整名字查表，查不到就是 Error::CannotFindTensor { path }，消息只有这一个名字，不提示相近的名字，也不列出文件里有什么。名字找到之后，candle-nn 里 mmap 后端的 get 先把整个张量读出来、转好 dtype，再和要求的形状比，不对就是 shape mismatch for <名字>, expected: .., got: ..。

Llama::load 自己会加 model. 前缀，所以调用方传进去的应该是没有前缀的根 VarBuilder：

    pub fn load(vb: VarBuilder, cfg: &Config) -> Result<Self> {
        let wte = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?;
        let lm_head = if cfg.tie_word_embeddings {
            Linear::from_weights(wte.embeddings().clone(), None)
        } else {
            linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?
        };
        let ln_f = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?;
        let blocks: Vec<_> = (0..cfg.num_hidden_layers)
            .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), cfg).unwrap())
            .collect();
  • tie_word_embeddings 为假时去取 lm_head.weight；SmolLM2 的输出层和词嵌入共享权重，文件里没有这个张量。LlamaConfig 里它是 Option<bool>，into_config 用 unwrap_or(false) 补默认值，所以配置里漏写等于「不共享」。
  • 30 层的 Block::load 在 map 里 unwrap()，任何一层的名字或形状不对都是 panic，前面三样（词嵌入、lm_head、最后的 RmsNorm）才是 Err。这一行从 2023 年 #794 把模型搬进 candle-transformers 时就是这样。返回 Result 的函数里藏一个 unwrap，我认为这是上游的小 bug：调用方的 ? 接不住，服务进程会因为一个坏的模型文件整个退出。

修法

先看文件里到底有什么名字、什么形状，再对代码：

let st = unsafe { MmapedSafetensors::new(format!("{DIR}/model.safetensors"))? };
let mut names: Vec<(String, Vec<usize>)> =
  st.tensors().into_iter().map(|(n, v)| (n, v.shape().to_vec())).collect();
names.sort();
for (n, s) in names.iter().filter(|(n, _)| n.starts_with("model.layers.0.self_attn")) {
  println!("  {n} {s:?}");
}
println!("lm_head.weight present: {}", names.iter().any(|(n, _)| n == "lm_head.weight"));
let vb = vb()?;
println!("vb.contains_tensor(\"model.embed_tokens.weight\") = {}", vb.contains_tensor("model.embed_tokens.weight"));
println!("vb.pp(\"model\").contains_tensor(\"model.embed_tokens.weight\") = {}", vb.pp("model").contains_tensor("model.embed_tokens.weight"));
let _ = Llama::load(vb, &config()?.into_config(false))?;

272 tensors; layer 0 attention:
  model.layers.0.self_attn.k_proj.weight [192, 576]
  model.layers.0.self_attn.o_proj.weight [576, 576]
  model.layers.0.self_attn.q_proj.weight [576, 576]
  model.layers.0.self_attn.v_proj.weight [192, 576]
lm_head.weight present: false
vb.contains_tensor("model.embed_tokens.weight") = true
vb.pp("model").contains_tensor("model.embed_tokens.weight") = false
Llama::load with config.json as is: ok
  • 前缀：文件里的名字从 model. 开始，Llama::load 也从 model. 开始拼，所以传根 vb。自己写层时，vb.contains_tensor(..) 在当前前缀下查，结果为假就说明前缀拼错了，比等 load 报错快。
  • 配置：不要手写 LlamaConfig，直接 serde_json::from_str 读模型自带的 config.json，tie_word_embeddings、num_key_value_heads 这类字段就不会漏。上面最后一行用原样的 config.json 加载成功。
  • 名字真的不一样（例如老模型的 gamma / beta）：用 rename_f 改名，见 cannot find tensor embeddings.word_embeddings.weight。
  • panic：在服务进程里加载不可信的模型文件时，Llama::load 外面套一层 std::panic::catch_unwind，把 panic 变回错误。

怎么确认没中招

  • 加载之前打印一次 MmapedSafetensors::tensors() 的前几个名字和你代码里第一个要的名字，前缀是否一致一目了然。
  • 报 cannot find tensor X 时，在名字列表里搜 X 的最后两段（如 q_proj.weight），看文件里它前面是什么。
  • shape mismatch for ... 的 expected 来自配置，got 来自文件：两者差一个整数倍（这里 576 对 192 是 3 倍）多半是头数的配置错了。

小结

  • VarBuilder 按 pp 拼出的完整名字取权重，找不到报 cannot find tensor <名字>，不给候选；形状不对报 shape mismatch for <名字>。
  • Llama::load 自己加 model. 前缀；配置漏了 tie_word_embeddings 会去找不存在的 lm_head.weight，漏了 num_key_value_heads 会按错误的形状要 k_proj。
  • Llama::load 里每一层的加载是 unwrap()，层内的错误是 panic 不是 Err。先列出文件里的名字和形状，配置从模型自带的 config.json 读。

💬 关联讨论

发起讨论
加载中...