Rust in Everything
基于 candle 0.11.0

from_pth 读到 0 个张量

主线位置：推理主线第一步「读 config 和权重」，权重是训练脚本用 torch.save 存的 .pt。

版本：PyTorch 2.15.0 / transformers 5.22.0（CPU）、candle 0.11.0（CPU，accelerate）
flowchart TB
  f[".pt 文件"] --> top{"data.pkl 顶层"}
  top -->|"state_dict"| ok["273 个张量"]
  top -->|"嵌套 dict"| zero["0 个张量"]
  top -->|"整个模型"| zero
  top -->|"TorchScript"| zero
  zero --> load["Llama load"]
  load --> err["cannot find tensor"]
  zero -->|"with_state"| ok

实验代码：PyTorch 侧 make_pt.py、fix_full.py，candle 侧 crate pitb-exp 的 bin p15_pth（release）。权重是本地的 SmolLM2-135M-Instruct。平台 M2 Max、macOS 27、rustc 1.99.0。四种文件各自读出什么、为什么，在 PyTorch 的 .bin / .pt 里逐个讲过，本页从「为什么没有报错」这一侧再走一遍。

PyTorch 的写法

训练脚本里存模型，常见的是这三种：

model = AutoModelForCausalLM.from_pretrained(model_dir)  # 默认 dtype="auto"，bf16
sd = model.state_dict()
torch.save(sd, f"{out}/sd.pt")
torch.save({"model": sd, "step": 100}, f"{out}/ckpt.pt")  # 带上训练步数的检查点
torch.save(model, f"{out}/full.pt")                       # 整个模型

另外有人用 torch.jit.save(torch.jit.script(m), "jit.pt") 导出 TorchScript（这里用一个 nn.Linear(3, 2) 试）。PyTorch 读回来：

obj = torch.load(f"{out}/{f}")
if f == "ckpt.pt":
  obj = obj["model"]

sd.pt    torch.load ok: OrderedDict, 273 keys
ckpt.pt  torch.load ok: OrderedDict, 273 keys
full.pt  torch.load UnpicklingError: Weights only load failed. This file can still be loaded, to 
full.pt  weights_only=False: LlamaForCausalLM
jit.pt   torch.jit.load ok: ['bias', 'weight']

检查点多一层 ["model"]，PyTorch 用户写这一下时不会多想。整个模型的文件在默认的 weights_only=True 下被拒（完整报错见 PyTorch 的 .bin / .pt），写上 weights_only=False 就能读。三种都读得出来，至少有一个明确的报错。

搬到 candle 的现象

把这几个文件交给 VarBuilder::from_pth，再 Llama::load。先用 pickle::read_pth_tensor_info 数一下读出了几个张量：

for f in ["sd.pt", "ckpt.pt", "full.pt", "full_sd.pt"] {
  let path = format!("{dir}/{f}");
  let n = pickle::read_pth_tensor_info(&path, false, None)?.len();
  let r = VarBuilder::from_pth(&path, DType::F32, &dev).and_then(|vb| Llama::load(vb, &cfg));
  match r {
    Ok(_) => println!("{f:8} tensors={n:3}  Llama::load ok"),
    Err(e) => println!("{f:8} tensors={n:3}  Llama::load Err: {e}"),
  }
}
let n = pickle::read_pth_tensor_info(format!("{dir}/jit.pt"), false, None)?.len();
println!("jit.pt   tensors={n:3}");

sd.pt    tensors=273  Llama::load ok
ckpt.pt  tensors=  0  Llama::load Err: cannot find tensor model.embed_tokens.weight
full.pt  tensors=  0  Llama::load Err: cannot find tensor model.embed_tokens.weight
full_sd.pt tensors=273  Llama::load ok
jit.pt   tensors=  0

（full_sd.pt 是下面「改法」里重新存的文件。）from_pth 对三种文件都返回 Ok。报错出现在 Llama::load 要第一个权重的时候，说的是「找不到 model.embed_tokens.weight」，读的人第一反应是键名前缀错了，去查 pp 和 rename_f，而真正的问题是文件里一个张量也没读出来。

原因

read_pth_tensor_info 解析 data.pkl 之后，只在两种情况下往下走：顶层就是一个字典，或者顶层是 __torch__.Module 这个类建出来的对象。之后假定这个字典就是 state_dict，逐项尝试解析成张量：

        // If the object is a dict, then we can extract the tensor info from it.
        // NOTE: We are assuming that the `obj` is state_dict by this stage.
        if let Object::Dict(key_values) = obj {
            for (name, value) in key_values.into_iter() {
                match value.into_tensor_info(name, &dir_name) {
                    Ok(Some(tensor_info)) => tensor_infos.push(tensor_info),
                    Ok(None) => {}
                    Err(err) => eprintln!("skipping: {err:?}"),
                }
            }
        }
    }
    Ok(tensor_infos)
  • ckpt.pt：顶层是字典，但 "model" 的值是另一个字典、"step" 是整数，into_tensor_info 对它们都返回 Ok(None)，连 stderr 都不打。
  • full.pt：顶层是 transformers...LlamaForCausalLM 建出来的对象，类名不是 __torch__.Module，整个文件被 continue 跳过。
  • jit.pt：TorchScript 的类名是 __torch__.torch.nn.modules.linear Linear，同样不匹配（PyTorch 的 .bin / .pt 列出了这几个文件 data.pkl 的开头）。

结果都是一个空的 Vec，from_pth 拿它建出一个「什么都没有」的 VarBuilder，不检查是否为空。PyTorch 和 candle 的分歧在于：torch.load 返回的是对象本身，由调用方决定取哪一层；candle 要的是「名字 → 张量」的表，不是这种形状就当没有。origin/main 上这一段只改了 zip 文件名的读法，判断逻辑和 0.11.0 相同。

改法

嵌套一层的检查点，用 from_pth_with_state 给出键，对应 PyTorch 里的 ["model"]：

let vb = VarBuilder::from_pth_with_state(format!("{dir}/ckpt.pt"), DType::F32, "model", &dev)?;
let model = Llama::load(vb, &cfg)?;

ckpt.pt with state key "model": logits max|diff|=2.98e-5 (max|ref| 20.6), greedy same prefix 20/20
ckpt.pt with state key "state_dict": Some("key state_dict not found")

logits 和 transformers f32 差 2.98e-5（最大 |ref| 20.6，容差 1e-4），20 个贪心 token 全同，和读 safetensors 一样。键写错时这次有报错：key state_dict not found。只能取一层，{"model": {"backbone": sd}} 要在 Python 里先拆开。

整个模型和 TorchScript 没有 candle 侧的办法。在信得过来源的前提下，回到 Python 里只存 state_dict：

model = torch.load(f"{d}/full.pt", weights_only=False)
torch.save(model.state_dict(), f"{d}/full_sd.pt")

full_sd.pt keys: 273

上面 candle 那一段的 full_sd.pt tensors=273 Llama::load ok 就是这个文件。TorchScript 同理：torch.jit.load(f).state_dict() 再 torch.save。能改存的话，直接存成 safetensors 更好（safetensors）。

怎么检查

  • 加载 .pt 之前先数张量：pickle::read_pth_tensor_info(path, false, None)?.len() 是 0 就停下，不要等 cannot find tensor。把 verbose 换成 true 会打印整个 data.pkl 解析出的对象树，能看出顶层是 Dict、Build 还是别的。
  • 看到 cannot find tensor <第一个权重的名字>（对 llama 是 model.embed_tokens.weight）时，先怀疑「什么都没读到」，再怀疑键名。键名问题通常报的是某一个名字，空文件报的总是第一个。
  • 在 Python 一侧，type(torch.load(f, weights_only=False)) 是 dict / OrderedDict 且值都是张量，candle 才能直接读。

小结

  • torch.save({"model": sd, ...})、torch.save(model)、TorchScript 三种文件，from_pth 都返回 Ok，读出 0 个张量，到 Llama::load 才报 cannot find tensor model.embed_tokens.weight。
  • 原因是 read_pth_tensor_info 只认顶层字典（或 __torch__.Module），并假定它就是 state_dict；别的形状静默跳过。
  • 嵌套一层用 from_pth_with_state(path, dtype, "model", dev)，SmolLM2 的 logits 差 2.98e-5、贪心 20/20。整个模型和 TorchScript 要在 Python 里改存 state_dict 或 safetensors。
  • 加载前用 read_pth_tensor_info(..).len() 判空。

💬 关联讨论

发起讨论
加载中...