safetensors 与 mmap

主线位置：推理主线 safetensors 那条路上，VarBuilder 背后的数据来源。上一页讲到 from_mmaped_safetensors 的后端会 self.load(name, dev)；本页讲这个 load 从文件的哪几个字节读、复制到哪里、在 CPU 和 Metal 上各花多少时间和内存。代码在 candle-core/src/safetensors.rs，格式解析交给 safetensors crate（candle 0.11.0 依赖 0.8.0）。

sequenceDiagram
  participant U as VarBuilder
  participant M as MmapedSafetensors
  participant S as safetensors crate
  participant D as Device
  U->>M: from_mmaped_safetensors
  M->>M: File open 和 mmap
  M->>S: deserialize 只解析头部
  U->>M: load name dev
  M->>S: tensor name
  S-->>M: TensorView 借用 mmap
  M->>D: from_slice 复制一次
  D-->>U: Tensor
  U->>U: to_dtype 再比形状

本页的实验是 nn-exp 的 st（release），平台 M2 Max、macOS 27、rustc 1.99.0。RSS 用 ps -o rss= 读，Metal 显存用 MTLDevice 的 currentAllocatedSize（经 metal_device().current_allocated_size()）。

文件格式

safetensors 文件分三段：8 字节小端 u64 表示头部长度 N，接着 N 字节 UTF-8 的 JSON 头，最后是所有张量的原始字节首尾相接。直接读 SmolLM2 的文件（st header，自己读字节、用 serde_json 解析，不经过 candle）：

file size = 269060552 bytes
first 8 bytes = [40, 77, 00, 00, 00, 00, 00, 00] -> header length 30528
data starts at byte 30536 (30536 mod 8 = 0)
header head: {"__metadata__":{"format":"pt"},"model.embed_tokens.weight":{"dtype":"BF16","shape":[49152,576],"data_offsets":[0,56623104]},"model.layers.0.input_layernorm.wei
trailing bytes of header: "016]}}      "
__metadata__ = {"format":"pt"}
tensors = 272
dtypes = ["BF16"]
  model.embed_tokens.weight BF16 [49152, 576] [0, 56623104)
  model.layers.0.input_layernorm.weight BF16 [576] [56623104, 56624256)
  model.layers.0.mlp.down_proj.weight BF16 [576, 1536] [56624256, 58393728)
  ... last: model.norm.weight BF16 [576] [269028864, 269030016)
8 + header + last end = 269060552
tensors whose absolute start is not 4-byte aligned = 0

• 0x7740 就是 30528。头部是一个 JSON 对象：__metadata__ 放字符串键值（这里只有 "format":"pt"），其余每个键是张量名，值是 dtype、shape 和 data_offsets。
• data_offsets 是相对数据段起点的半开区间。按偏移排序后一个接一个，没有空隙；embed_tokens 是 49152 × 576 × 2 字节 = 56623104。最后一个张量的终点加上 8 + N 正好等于文件大小。
• 头部末尾有 6 个空格：写文件的一方把头部补齐到 8 的倍数，所以数据段从 8 字节对齐的位置开始。272 个张量（30 层 × 9 个 + embed_tokens + norm）全是 BF16，起点全都 4 字节对齐。

safetensors 0.8.0 的 SafeTensors::read_metadata 读头部时检查：N 不超过 100,000,000；JSON 能按它的 Metadata 结构反序列化（serde_json，见 serde 卷的 JSON 反序列化器）；每个张量的区间从上一个的终点开始，长度等于元素数乘 dtype 的位数除以 8；数据段的总长度和文件剩余部分一样长，多一个字节少一个字节都报错。这些检查都只看头部，不碰数据。

MmapedSafetensors：mmap 加 yoke

MmapedSafetensors：

pub struct MmapedSafetensors {
    safetensors: Vec<yoke::Yoke<SafeTensors_<'static>, memmap2::Mmap>>,
    routing: Option<HashMap<String, usize>>,
}

safetensors::SafeTensors<'data> 解析完头部后借用着整个文件的字节（TensorView 里的数据是 &'data [u8]），而文件字节属于 memmap2::Mmap。两者要放进同一个结构体，就成了自引用。yoke::Yoke<Y, C> 把「借用者」Y 和它借用的「车」C 绑在一起，对外是一个 'static 的值，取用时再通过 get() 把生命周期缩回来。new 打开文件、MmapOptions::new().map(&file)、在 try_attach_to_cart 的闭包里 SafeTensors::deserialize，就这么多。multi 对多个分片文件各做一遍，再建一张「张量名到文件序号」的 routing 表，同名张量以后面的文件为准。VarBuilder::from_mmaped_safetensors 走的是 multi，只有一个文件时也一样。

建立时只读了头部：实测 mmap + parse header 在 0.2 到 0.6 ms 之间，RSS 几乎不变（16 到 17 MiB）。函数是 unsafe 的，理由来自 memmap2：映射期间如果别的进程改了文件，Rust 看到的「不可变切片」就变了。

取一个张量：复制一次

load 先 get：按 routing 找到文件，向 SafeTensors 要一个 TensorView（dtype、形状、指向 mmap 的字节切片）。然后是 Load::load，即 convert：按文件里的 dtype 选元素类型。BF16、F16、F32 等直接对应；文件里的 U16 读成 U32（逐元素转换），I8 这类 candle 没有的 dtype 报 UnsupportedSafeTensorDtype。直接对应的走 convert_slice：

fn convert_slice<T: WithDType>(data: &[u8], shape: &[usize], device: &Device) -> Result<Tensor> {
    let size_in_bytes = T::DTYPE.size_in_bytes();
    let elem_count = data.len() / size_in_bytes;
    if (data.as_ptr() as usize).is_multiple_of(size_in_bytes) {
        // SAFETY This is safe because we just checked that this
        // was correctly aligned.
        let data: &[T] =
            unsafe { std::slice::from_raw_parts(data.as_ptr() as *const T, elem_count) };
        Tensor::from_slice(data, shape, device)
    } else {
        // ...
        let mut c: Vec<T> = Vec::with_capacity(elem_count);
        // ...
        unsafe {
            std::ptr::copy_nonoverlapping(data.as_ptr(), c.as_mut_ptr() as *mut u8, data.len());
            c.set_len(elem_count)
        }
        Tensor::from_slice(&c, shape, device)
    }
}

对齐时把 mmap 里的字节直接看成 &[T]，交给 Tensor::from_slice。复制发生在 from_slice 里面，Device::storage_from_slice 按设备分开：

• CPU：data.to_cpu_storage()，即 to_vec()，复制进一个新的 Vec<T>。CpuStorage 的每个变体都拥有自己的 Vec，没有能借用 mmap 的变体，所以 CPU 上不可能零拷贝。
• Metal：MetalStorage 的 storage_from_slice 用 with_data 建缓冲区，落到 new_buffer_with_data，即 Metal 的 newBufferWithBytes：复制进一个 StorageModeShared 的新缓冲区。Apple 芯片是统一内存，CPU 和 GPU 看的是同一块物理内存，但 candle 没有用不复制的接口，mmap 的页和 Metal 缓冲区是两份。中间没有经过 CPU 的 Vec，从 mmap 到 GPU 可见也只复制一次。

不对齐时多一次：先复制进 c，from_slice(&c, ..) 再复制一次（改用 Tensor::from_vec 可以省掉）。按上面的统计，SmolLM2 的文件走不到这个分支。npz 和 pth 后端是另一种情况：它们先读到 CPU 张量，再 to_device，到 Metal 是两次。

实测：时间、RSS、Metal 显存

st cpu / st metal：MmapedSafetensors::new 之后把 272 个张量全部 load 到设备上（保持 BF16），再丢掉 MmapedSafetensors，然后全部 to_dtype(F32)，再丢掉 BF16 的那一份。Metal 上读时钟前都先 synchronize()。以下是页缓存已热时的数字（第二次及以后的运行）：

• CPU 上 load 之后 RSS 涨了约 515 MiB，是两份 257 MiB：读过的 mmap 页（文件页，算在 RSS 里）和复制出来的 Vec。丢掉 MmapedSafetensors 解除映射，RSS 降回 275 MiB，张量还在。这直接说明张量和文件映射不共享内存。
• 本机第一次运行时 CPU 的 load 用了 315 ms，之后各次在 37 到 48 ms。差别应当是页缓存：第一次要从磁盘读文件（本页没有清空缓存去复现）。Metal 上第一次 to_dtype(F32) 用了 131 ms，包含 cast 内核的编译，表里是之后的数字。
• CPU 上丢掉 BF16 张量之后 RSS 没有降，内存没有还给系统；本页没有深究是哪一层分配器留住的。
• Metal 上 load 分配了 258 MiB，和数据量一致：new_buffer_with_data 按实际大小分配，不进复用池。
• Metal 上转出来的 F32 张量用的是 new_buffer 分配的 StorageModePrivate 缓冲区，不算在进程 RSS 里，所以 RSS 不变、Metal 分配涨了 638 MiB。F32 数据只有 513 MiB，多出来的是 buf_size 把每个缓冲区取整到 2 的幂：embed_tokens 的 F32 是 108 MiB，分到 128 MiB。
• 丢掉 BF16 张量时 Metal 分配没有马上降，synchronize() 之后才降：缓冲区池里引用计数为 1 的缓冲区在 wait_until_completed 里才被清掉。Metal 缓冲区池的细节见 Metal 后端。

BF16 转 F32

BF16 就是 F32 的高 16 位（符号、8 位指数、7 位尾数），转成 F32 只是在低位补 16 个 0，没有精度损失。实验里把 tensors() 返回的第一个张量转出来的 F32 再转回 BF16、再转 F32，和转回之前比，最大差是 0：BF16 到 F32 再到 BF16 是一个来回不丢信息的映射。代价全在内存和时间上：数据翻倍，CPU 上整个模型转一遍 41 ms，Metal 上 22 到 58 ms。

VarBuilder 的 dtype 和文件不同时，每次 get 都是「复制一次 BF16，再转一次 F32」，BF16 那份随即丢掉。SmolLM2 在 Metal 上用 BF16 是 258 MiB，用 F32 是 639 MiB（含 2 的幂取整）。CPU 上的 F32 算子更全、更快（例如 BF16 的 softmax 在 CPU 上会算错，见常用层），所以 CPU 上常用 F32，GPU 上常用 BF16 或 F16。

Err: unsupported safetensor dtype I8

写文件

反方向是 Tensor::save_safetensors 和 candle::safetensors::save，都交给 safetensors::serialize_to_file。candle 为 Tensor 实现了 safetensors::View，其中 data() 调 convert_back：先 flatten_all，再 to_vec1 读回 CPU，GPU 上的张量在这里同步并复制回主机；源码注释写着 CPU 上也多复制了一次，还有一个 unwrap 的 TODO（View::data 的签名不能返回错误）。

动脑筋答案

  let st = unsafe { MmapedSafetensors::new(ST)? };
  let names: Vec<String> = st.tensors().into_iter().map(|(n, _)| n).collect();
  let mut ts = Vec::new();
  for n in names.iter() {
    ts.push(st.load(n, dev)?.to_dtype(DType::F32)?);
  }
  dev.synchronize()?;
  // ... print rss
  drop(st);

[cpu] load+to_dtype(F32) per tensor (what VarBuilder F32 does): 67.866667ms, rss 16 -> 852 MiB
[cpu] mmap dropped: rss 595 MiB, metal allocated 0 MiB
[metal] load+to_dtype(F32) per tensor (what VarBuilder F32 does): 81.955667ms, rss 21 -> 286 MiB
[metal] mmap dropped: rss 29 MiB, metal allocated 639 MiB

• CPU：852 MiB 约等于 16 + 257（读过的 mmap 页）+ 513（F32 张量）+ 60 多（BF16 临时张量释放后没还给系统的部分）。丢掉 mmap 后是 595 MiB，减少的 257 MiB 正好是文件大小。三次运行分别是 849 / 848 / 852 和 592 / 592 / 595。
• Metal：F32 张量都在 Private 缓冲区里，不算 RSS；显存里是 639 MiB，三次都一样。RSS 里是读过的 mmap 页和还没被池清掉的 Shared 临时缓冲区，这部分每次不同（四次运行分别是 436、421、414、286 MiB，丢掉 mmap 后 89、64、60、29 MiB）。
• 时间：CPU 三次都在 67 到 68 ms；Metal 四次在 67 到 82 ms。都比「先全部读 BF16（约 40 ms）再全部转（20 到 60 ms）」两步分开做要快一点或持平：每个张量刚复制完就转，数据还在缓存里。

小结

• safetensors：8 字节小端长度、补齐到 8 字节的 JSON 头、首尾相接的原始数据；safetensors crate 只看头部就能验证整个文件的布局。
• MmapedSafetensors 用 yoke 把 mmap 和借用它的解析结果放进一个值，建立时只解析头部，零点几毫秒。
• 取张量时 CPU 复制进 Vec，Metal 复制进 Shared 缓冲区，各一次，都不是零拷贝；npz、pth 到 Metal 是两次。
• BF16 转 F32 无损，但内存翻倍；Metal 上的中间缓冲区取整到 2 的幂，SmolLM2 用 F32 占 639 MiB。
• 实测 SmolLM2（257 MiB）：热缓存下 CPU 和 Metal 读完都在 40 ms 左右，转 F32 再加 20 到 60 ms。