Rust in Everything
基于 candle 0.11.0

safetensors

主线位置：推理主线第一步「读 config 和权重」，权重是 HuggingFace 目录里的 model.safetensors（或分片）。

版本：PyTorch 2.15.0 / transformers 5.22.0 / safetensors 0.8.0（CPU）、candle 0.11.0（CPU，accelerate；Metal）
flowchart TB
  fp["from_pretrained 目录"] --> auto["dtype auto 读 config"]
  auto --> zc["bf16 直接用文件页"]
  vbm["from_mmaped_safetensors"] --> ask["dtype 必须写明"]
  ask --> cp["每个张量复制一次"]
  cp --> conv{"和文件 dtype 同吗"}
  conv -->|同| keep["保留 BF16"]
  conv -->|不同| cast["再转一份"]
  idx["index.json"] -->|"自己读 weight_map"| vbm
  idx -->|"自动"| fp

实验代码：PyTorch 侧几个脚本（load_only.py、torch_mmap.py、make_shards.py、logits.py），candle 侧 crate weights-exp 的 bin load 和 f16conv（release）。权重是本地的 SmolLM2-135M-Instruct（model.safetensors，bf16，256.6 MiB）。平台 M2 Max、macOS 27、rustc 1.99.0。文件格式本身（8 字节长度、JSON 头、原始数据）见源码卷的 safetensors 与 mmap。

一个调用对三步

from transformers import AutoModelForCausalLM

model_dir = "SmolLM2-135M-Instruct"  # 本地的 SmolLM2-135M-Instruct 权重
model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.float32)

let model_dir = "SmolLM2-135M-Instruct";
let cfg: LlamaConfig = serde_json::from_slice(&std::fs::read(format!("{model_dir}/config.json"))?)?;
let cfg = cfg.into_config(false);
let vb = unsafe {
  VarBuilder::from_mmaped_safetensors(&[format!("{model_dir}/model.safetensors")], DType::F32, &dev)?
};
let model = Llama::load(vb, &cfg)?;

两边都读同一个 config.json 和同一个 model.safetensors，同一个 prompt 的最后一个位置的 logits 差 2.98e-5（最大 |ref| 20.6，容差 1e-4，见怎么读这一卷）。区别在 from_pretrained 替你做掉的那些事，candle 里每一件都要自己写：

from_pretrained(目录) 做的事candle 里
读 config.json，按 model_type 挑模型类serde_json 读成 LlamaConfig，模型类型自己选
在目录里找 model.safetensors、分片索引或 pytorch_model.bin把文件路径列表传给 from_mmaped_safetensors（.bin 走 from_pth，见 PyTorch 的 .bin / .pt）
dtype 默认 "auto"：用 config 里的 torch_dtypeDType 是必填参数
绑定 lm_head 和词嵌入Llama::load 看 cfg.tie_word_embeddings（见键名对不上）
缺的权重随机初始化，打印 LOAD REPORTcannot find tensor ... 或 panic
device_map、.to(device)&Device 在建 VarBuilder 时给定（见设备）

from_mmaped_safetensors 是 unsafe 的：它 mmap 文件，映射期间文件被别的进程改了，Rust 看到的「不可变」字节也会变。它只在建立时解析文件头，权重在模型构造函数按名字要的时候才读（VarBuilder）。

dtype：一个默认 bf16，一个没有默认

m = AutoModelForCausalLM.from_pretrained(model_dir)
print("default dtype:", m.dtype)
for d in ["float32", "bfloat16", "float16"]:
  m = AutoModelForCausalLM.from_pretrained(model_dir, dtype=getattr(torch, d))
  print(d, "->", m.dtype, m.model.layers[0].self_attn.q_proj.weight.dtype)

default dtype: torch.bfloat16
float32 -> torch.float32 torch.float32
bfloat16 -> torch.bfloat16 torch.bfloat16
float16 -> torch.float16 torch.float16

transformers 5.22.0 的 from_pretrained 文档写着 dtype 「defaults to "auto"」：先用 config.json 里的 dtype / torch_dtype，没有再看 checkpoint 里第一个浮点权重。SmolLM2 的 config 写的是 "torch_dtype": "bfloat16"，所以不写 dtype 就得到一个 bf16 模型。从老版本 transformers 过来的习惯是「不写就是 f32」，在 5.x 里不成立。

candle 的 from_mmaped_safetensors 必须给 DType，文件是什么都按你给的转：

$ load cpu bf16 st SmolLM2-135M-Instruct/model.safetensors
load 39.6 ms
Error: unsupported dtype BF16 for op matmul
$ load cpu f16 st SmolLM2-135M-Instruct/model.safetensors
load 115.8 ms
Error: the accelerate backend does not support f16 matmul

三种 dtype 都能加载；在 CPU 上，BF16 和 F16 第一次 matmul 就报错（F16 的报错来自 accelerate 后端），Metal 上三种都能前向。半精度怎么选、误差多大，见半精度。CPU 上推理用 DType::F32。

转成 F16：两边逐位相同

文件是 bf16，载成 F16 要换一种 16 位格式：BF16 是 8 位指数、7 位尾数，F16 是 5 位指数、10 位尾数，范围小、精度高。candle CPU 上 to_dtype 的这一分支先转 f32 再舍入成 f16：

            (Self::BF16(storage), DType::F16) => {
                let data = unary_map(storage, layout, |v| f16::from_f32(v.to_f32()));
                Ok(Self::F16(data))
            }

f16conv 按 VarBuilder 的路子把 272 个张量 load 再 to_dtype(F16)，写成 safetensors；Python 侧用 tensor.to(torch.float16) 转同样的 bf16 数据，逐位比：

f16_cpu.safetensors: 134515008 values, bitwise differs from torch .to(float16): 0
f16_metal.safetensors: 134515008 values, bitwise differs from torch .to(float16): 0
torch bf16->f16: inf 0, flushed to 0 23, subnormal 43183
bf16 |w| max 9.25, min nonzero 2.41562e-09
max |f16 - bf16| 2.98e-08

CPU 和 Metal 上 candle 的结果都和 torch 逐位相同。candle CPU 的做法是先精确转成 f32，再按就近舍入取 f16，torch 的结果和它一个值都不差。这份权重最大绝对值 9.25，远没到 F16 的上限 65504，没有溢出成 inf；但有 23 个极小的值（最小的是 2.4e-9）变成了 0，43183 个落进 F16 的次正规数，改变量最多 2.98e-8。所以 DType::F16 对这个模型几乎无损，换一个权重里有大数的模型就不一定了。

内存：bf16 时 transformers 不复制

/usr/bin/time -l 跑只加载、不前向的进程，每组 5 次取最小值。transformers 侧 torch.set_num_threads(8)，时间是进程内 from_pretrained 这一个调用；candle 侧是 from_mmaped_safetensors 加 Llama::load，Metal 上读时钟前 synchronize()。跑的时候 uptime 的负载在 6 到 8 之间（其他实验在并行跑）；各组 5 次之间的最大值见实验输出，下表只列最小值。

加载时间峰值 RSS
transformers，不写 dtype（bf16）615 ms336.0 MiB
transformers，torch.float32635 ms1090.6 MiB
transformers，torch.bfloat16647 ms336.5 MiB
transformers，torch.float16671 ms834.0 MiB
candle CPU，DType::F3267.2 ms841.1 MiB
candle CPU，DType::BF1637.6 ms522.1 MiB
candle CPU，DType::F16113.6 ms585.3 MiB
candle Metal，DType::F32 / BF16 / F1665.3 / 40.9 / 56.8 ms约 530 MiB

transformers 侧还要另算 import torch 和 transformers 的约 1.0 秒，加上解释器启动，进程的 real 在 2.0 到 2.2 秒，candle 进程的 real 就是加载时间。冷启动的完整对比在什么时候选 candle。

F32 时两边的峰值差不多是同一个构成：映射进来、被读过的 bf16 文件页（257 MiB）加 f32 权重（513 MiB），transformers 再多 Python 运行时的 200 多 MiB。bf16 时反过来了：transformers 只有 336 MiB，candle 是 522 MiB。原因是 transformers 在 dtype 和文件一致时不复制：

print(f"after import: rss {rss():.0f} MiB")
m = AutoModelForCausalLM.from_pretrained(model_dir, dtype=dtype)
print(f"after from_pretrained({sys.argv[2]}): rss {rss():.0f} MiB, peak {peak():.0f} MiB")
e = m.model.embed_tokens.weight
q = m.model.layers[0].input_layernorm.weight
d = m.model.layers[0].mlp.down_proj.weight
print("embed storage bytes", e.untyped_storage().nbytes(), "ptr gap embed->layers.0.input_layernorm",
      q.untyped_storage().data_ptr() - e.untyped_storage().data_ptr(),
      "-> down_proj", d.untyped_storage().data_ptr() - q.untyped_storage().data_ptr())
with torch.no_grad():
  s = sum(float(p.sum()) for p in m.parameters())
print(f"after touching all params: rss {rss():.0f} MiB")

after import: rss 222 MiB
after from_pretrained(bfloat16): rss 327 MiB, peak 327 MiB
embed storage bytes 56623104 ptr gap embed->layers.0.input_layernorm 56623104 -> down_proj 1152
after touching all params: rss 581 MiB
after import: rss 221 MiB
after from_pretrained(float32): rss 834 MiB, peak 1090 MiB
embed storage bytes 113246208 ptr gap embed->layers.0.input_layernorm 22805915648 -> down_proj -22692636672
after touching all params: rss 834 MiB
  • bf16 时，相邻三个参数的存储地址之差是 56623104 和 1152，正好是文件头里 embed_tokens（49152 × 576 × 2 字节）和 input_layernorm（576 × 2 字节）的长度：参数直接指向映射进来的文件，一个接一个排着。from_pretrained 返回时 RSS 只涨了 105 MiB，因为文件页还没被读过；把每个参数求一次和，页被换进来，RSS 才到 581 MiB。
  • f32 时地址毫无规律，是新分配的内存，加载完 RSS 就是 834 MiB，求和也不再涨。
  • candle 两种情况都复制：CPU 上每个张量 to_vec() 进自己的 Vec，CpuStorage 没有能借用 mmap 的变体，所以 BF16 也是「文件页 257 MiB + 副本 257 MiB」。机制和逐步的 RSS 在源码卷的 safetensors 与 mmap。好处是加载完就全在内存里，第一次前向不会再缺页；代价是 bf16 多一份。
  • Metal 一行的 RSS 不能和 CPU 比：F32 和 F16 的权重在 Private 缓冲区里，不算进 RSS（源码卷的同一页有 Metal 显存的数字）。

分片：自己读 index.json

大模型的权重分成几个文件，旁边一个 model.safetensors.index.json 写着每个张量在哪个文件里。SmolLM2 只有一个文件，用 safetensors 的 Python 库按文件顺序切成 3 片，做一份分片的副本：

with safe_open(f"{src}/model.safetensors", framework="pt") as f:
  meta = f.metadata()
  tensors = {k: f.get_tensor(k) for k in f.keys()}
# ... 按累计字节数分成 n 组 shards
for i, sh in enumerate(shards):
  name = f"model-{i + 1:05d}-of-{len(shards):05d}.safetensors"
  save_file(sh, f"{out}/{name}", metadata=meta)
  for k in sh:
    weight_map[k] = name
json.dump({"metadata": {"total_size": total}, "weight_map": weight_map},
          open(f"{out}/model.safetensors.index.json", "w"), indent=2)

model-00001-of-00003.safetensors 40 tensors 88483968 bytes
model-00002-of-00003.safetensors 115 tensors 88503552 bytes
model-00003-of-00003.safetensors 117 tensors 92042496 bytes
total_size 269030016

transformers 不用改代码，from_pretrained(分片目录) 自己找到索引；logits 和单文件逐位相同（compare.py 容差 0，差 0.00e+00），加载 626 ms、峰值 1090.5 MiB，和单文件一样。

candle 的 from_mmaped_safetensors 本来就接一组路径（MmapedSafetensors::multi 为每个文件建一个映射，再建一张「名字到文件」的表），但不认索引文件，要自己从 weight_map 里取出文件名：

fn shard_files(index: &str) -> Result<Vec<std::path::PathBuf>> {
  let dir = std::path::Path::new(index).parent().context("index dir")?;
  let json: serde_json::Value = serde_json::from_slice(&std::fs::read(index)?)?;
  let map = json["weight_map"].as_object().context("weight_map")?;
  let mut files: Vec<String> = map.values().filter_map(|v| v.as_str().map(String::from)).collect();
  files.sort();
  files.dedup();
  Ok(files.iter().map(|f| dir.join(f)).collect())
}

let files = shard_files(&format!("{model_dir}/model.safetensors.index.json"))?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&files, DType::F32, &dev)? };

shards: [Some("model-00001-of-00003.safetensors"), Some("model-00002-of-00003.safetensors"), Some("model-00003-of-00003.safetensors")]
load 510.9 ms
logits                   n=49152  max|diff|=2.98e-05
worst=2.98e-05 tol=1e-04 OK

和 transformers 比差 2.98e-5，和 candle 读单文件的 logits 比差 0。510.9 ms 是三个新文件第一次读（页缓存是冷的）；之后 5 次的最小值是 67.6 ms、峰值 841.1 MiB，和单文件相同。分片只改变文件的切法，不改变数据。

candle 仓库的 candle-examples 里有一个做同样事情的 hub_load_local_safetensors，它属于示例 crate，不在 candle-core / candle-nn 的 API 里，自己写这十行就够。少给一个分片时：

$ load cpu f32 st model-00001-of-00003.safetensors noforward
Error: cannot find tensor model.norm.weight

Llama::load 先取词嵌入和最后的 norm，再逐层取，所以缺第三片时报的是 model.norm.weight。

❗IMPORTANT
动脑筋：把 model.safetensors.index.json 本身当成 safetensors 文件传给 from_mmaped_safetensors（像给 from_pretrained 传一个路径那样），会得到什么？

动脑筋答案

$ load cpu f32 st model.safetensors.index.json noforward
Error: path: "model.safetensors.index.json" header too large
  • safetensors 文件的前 8 个字节是小端 u64 的头部长度。索引是 JSON，前 8 个字节是 {\n "met，读成 u64 是 8387229874382703227。
  • safetensors 0.8.0 规定头部不超过 100,000,000 字节（MAX_HEADER_SIZE），超过就是 HeaderTooLarge，显示为 header too large；candle 在前面加上了路径。
  • 所以看到 header too large，先检查传进去的是不是 safetensors 文件。传一个目录则是 path: "sharded" Invalid argument (os error 22)：from_mmaped_safetensors 只认文件。

小结

  • from_pretrained(目录) 读 config、挑文件、定 dtype、绑定权重、补缺；candle 里 config 用 serde_json 读，文件列表、DType、设备都要写明。
  • transformers 5.22.0 的 dtype 默认 "auto"，SmolLM2 不写就是 bf16；candle 没有默认。CPU 上 BF16 / F16 能加载、不能 matmul。
  • 文件是 bf16、载成 bf16 时，transformers 的参数直接指向映射的文件页（加载后 RSS 只涨 105 MiB），candle 复制一份（522 MiB 对 336 MiB）；F32 时两边都要复制和转换，峰值 841 MiB 对 1091 MiB。
  • BF16 转 F16，candle（CPU、Metal）和 torch 1.3 亿个值逐位相同；这份权重有 23 个值变成 0。
  • 分片：transformers 自己读索引；candle 从 weight_map 取文件名传给 from_mmaped_safetensors，logits 和单文件逐位相同。

💬 关联讨论

发起讨论
加载中...