模型加载、Safetensors 与 Checkpoint Sharding

它位于磁盘格式与运行时权重之间

模型加载把 checkpoint 中的 bytes 转换成运行时参数 storage。它发生在请求 tokenization 和模型执行之前，决定启动峰值内存、加载时间、最终 dtype/layout 与每个 device/rank 的权重所有权；它不执行推理，也不等同 Backend Dispatch 或 KV Cache 分配。

一条典型路径是：

config + shard index + shard files
        ↓  检查文件、tensor names、shape/dtype、hash
CPU virtual address / OS page cache / tensor slice
        ↓  cast、dequantize/repack、transpose 或 TP slicing
pinned/pageable CPU staging buffer（可选）
        ↓  H2D copy / per-rank placement
GPU runtime parameter storage

最终设备上的权重可能与文件中的 bytes 不同：文件可以是量化 packed format，加载时却被展开；也可能在 CPU 先切成 TP shard、改变 dtype 或为 kernel 重排。因而“checkpoint 大小”“CPU RSS”“GPU allocated”和“最终 kernel 读取 bytes”必须分别测量。

Shard index 是名字到文件的路由表

当权重太大而拆成多个文件时，index 常包含：

{
  "metadata": {"total_size": 1152},
  "weight_map": {
    "embed.weight": "model-00001-of-00002.safetensors",
    "block.w1": "model-00001-of-00002.safetensors",
    "block.w2": "model-00002-of-00002.safetensors"
  }
}

加载器不应把 00001 理解成“第 1 层”；真正路由由 weight_map 定义。一个可验证的 manifest 至少要检查：

• 所有期望参数恰好出现一次，没有 missing/unexpected/duplicate names；
• 每个映射文件存在，header 可解析，tensor offset 不越界；
• shape、dtype 与模型配置一致，例如 block.w1:[H,I] 而不是转置后的 [I,H]；
• tied weights、aliases 或共享 storage 的约定被明确处理；
• shard 文件的 SHA-256 等 digest 与可信 manifest 一致。

Sharding 降低的是单文件大小和可选择读取范围，也便于按 rank/worker 并行加载。它不会自动减少全模型 payload；若每个 rank 仍读取并保留全部 shards，总 bytes 和峰值仍可能接近复制加载。

mmap、lazy loading 与 CPU staging 的边界

Safetensors safe_open 可把文件映射进进程地址空间，并按 tensor name 或 slice 访问所需区域。这里的“未立刻复制整文件”不等于“访问没有成本”：

• 首次触碰对应 pages 仍会触发 storage I/O 或 page fault；之后可能命中 OS page cache；
• 虚拟地址映射大小不等于 resident set size（RSS），RSS 也不等于 Python/PyTorch allocator 统计；
• get_tensor() 返回完整 tensor view/对象后，后续 cast、clone、pin memory 或 .to("cuda") 可能分配并复制完整 payload；
• 文件生命周期、映射对象和返回 tensor 的 ownership 必须由具体 library API 保证，不能假定任意自制 mmap view 都可在关闭文件后安全使用；
• 网络文件系统、容器 overlay、磁盘吞吐和并发 page faults 会改变启动延迟。

因此更准确的说法是：mmap 允许按需把文件区域纳入地址空间，get_slice 允许只物化某个 tensor slice；完整 lazy model loading 还要求构造模型时不先分配第二份随机权重、不立刻遍历/转换所有参数，并在放置到最终 device 后及时释放 staging storage。

CPU staging 常见两种方式：pageable CPU tensor 直接 .to(device)，或先放入 pinned memory 再用 non-blocking H2D copy 与其他工作重叠。Pinned memory 不是免费资源：它受主机内存与 OS 限制，过量会伤害系统；非阻塞 copy 也只有在 stream、buffer lifetime 和同步正确时才可能形成重叠。

TP 权重转换：尽量在传输前只取本 rank 所需 slice

以列并行线性层 $W:[K,N]=[8,12]$、TP world size $p=2$ 为例：

rank 0/1 分别拥有输出特征 [0:6] / [6:12]。若 checkpoint 保存全量 [8,12]，较低峰值的路径是用 tensor slice 在 CPU 端读取 [8,6]，再只传该 shard；先把完整 [8,12] 复制到每张 GPU 再切片，会造成临时复制容量与 PCIe/NVLink 流量。

但切分轴必须跟 张量并行 的数学契约一致：列并行通常切 output feature，行并行通常切 input/reduction dimension。量化 checkpoint 还要让 scale/zero-point group 与 shard 边界一致；若 group 跨 rank 边界，不能只切 payload 而忽略 metadata 或重新量化语义。

可运行的本地分片实验

实验问题：不用下载模型，能否创建两个 Safetensors shards 和 index，逐文件校验 SHA-256，按名字加载，直接读取 TP column slices，并证明 CPU→device staging 保持数值？

预期结果：所有原 tensor 与加载结果完全一致；两个 [8,6] slices 拼接回 [8,12] 且最大差为 0。临时 checkpoint 在脚本退出时删除，不修改项目数据。

import hashlib,json,platform
from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from safetensors import safe_open
from safetensors.torch import save_file

torch.manual_seed(163)

def sha256(path):
    digest=hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda:stream.read(1<<16),b""): digest.update(chunk)
    return digest.hexdigest()

def write_checkpoint(root,state):
    shards={
        "model-00001-of-00002.safetensors":{
            "embed.weight":state["embed.weight"],"block.w1":state["block.w1"]},
        "model-00002-of-00002.safetensors":{"block.w2":state["block.w2"]},
    }
    weight_map={}; checksums={}
    for filename,tensors in shards.items():
        path=root/filename; save_file(tensors,path,metadata={"format":"pt"})
        checksums[filename]=sha256(path)
        weight_map.update({name:filename for name in tensors})
    index={"metadata":{"total_size":sum(t.numel()*t.element_size() for t in state.values())},
           "weight_map":weight_map,"sha256":checksums}
    (root/"model.safetensors.index.json").write_text(json.dumps(index,indent=2))
    return index

def load_tensor(root,index,name):
    filename=index["weight_map"][name]; path=root/filename
    assert sha256(path)==index["sha256"][filename]
    with safe_open(path,framework="pt",device="cpu") as handle:
        return handle.get_tensor(name)

def load_column_parallel_slice(root,index,name,rank,world_size):
    filename=index["weight_map"][name]; path=root/filename
    assert sha256(path)==index["sha256"][filename]
    with safe_open(path,framework="pt",device="cpu") as handle:
        tensor_slice=handle.get_slice(name); shape=tensor_slice.get_shape()
        assert shape[1]%world_size==0; width=shape[1]//world_size
        return tensor_slice[:,rank*width:(rank+1)*width]

def main():
    state={"embed.weight":torch.randn(12,8),
           "block.w1":torch.randn(8,12),"block.w2":torch.randn(12,8)}
    with TemporaryDirectory(prefix="glossary-checkpoint-") as directory:
        root=Path(directory); index=write_checkpoint(root,state)
        assert set(index["weight_map"])==set(state)
        loaded={name:load_tensor(root,index,name) for name in state}
        for name in state: torch.testing.assert_close(loaded[name],state[name])

        tp_shards=[load_column_parallel_slice(root,index,"block.w1",rank,2)
                   for rank in range(2)]
        reconstructed=torch.cat(tp_shards,dim=1)
        torch.testing.assert_close(reconstructed,state["block.w1"])
        device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
        staged=loaded["block.w2"].to(device)
        torch.testing.assert_close(staged.cpu(),state["block.w2"])

        names=sorted(index["sha256"])
        prefixes={name:index["sha256"][name][:12] for name in names}
        print(f"python={platform.python_version()} torch={torch.__version__}")
        print(f"safetensors_shards={names} total_payload_bytes={index['metadata']['total_size']}")
        print(f"weight_map={index['weight_map']}")
        print(f"CPU_loaded_block_w1={tuple(loaded['block.w1'].shape)} dtype={loaded['block.w1'].dtype}")
        print(f"TP_column_shards={[tuple(t.shape) for t in tp_shards]} "
              f"reconstructed_max_diff={(reconstructed-state['block.w1']).abs().max():.3e}")
        print(f"staging cpu->{device} tensor={tuple(staged.shape)} "
              f"max_diff={(staged.cpu()-state['block.w2']).abs().max():.3e}")
        print(f"sha256_prefixes={prefixes}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/model_loading_reference.py，需要当前环境已有的 safetensors 包。2026-08-13 在 Python 3.11.12、Safetensors 0.7.0、PyTorch 2.13.0+cu130 上，两个 TP slices 均为 [8,6]，拼接与 CPU→CUDA staging 最大差都为 0，checkpoint tensor payload 合计 1152 bytes。

代码有意没有把性能数字附会到真实模型：临时文件很小，OS page cache 几乎主导，无法代表多 GiB checkpoint、网络存储或多进程加载。它证明的是 routing、slice、ownership、hash 与数值不变量。

峰值容量与失败场景

若 checkpoint payload 为 $M_{file}$，最终每 rank 参数为 $M_{rank}$，最坏的朴素路径可能同时存在：模型初始化参数、完整 CPU checkpoint、cast 后 CPU copy、完整 GPU copy 与最终 shard。更有用的峰值核对式是：

这些项生命周期可能重叠或复用，RSS、mapped virtual memory、PyTorch allocated/reserved 口径也不同，不能机械相加成一个跨平台精确值。生产优化目标是缩短重叠生命周期，并避免“完整复制后再切分”。

常见失败包括：

• index 缺少参数或指向错误 shard，但 loader 只验证文件能打开；
• 文件 tensor 是 [N,K]，模型期望 [K,N]，shape 恰好方阵时错误被掩盖；
• dtype cast 或量化展开使 GPU 容量远大于 checkpoint 文件；
• 每个 TP rank 都从共享存储读取全模型，引发 I/O 放大和 page-cache 竞争；
• 多进程同时 pin 大量 CPU buffers，导致主机内存压力；
• .to(cuda) 后仍保留完整 CPU copies，启动成功但可服务 batch 容量下降；
• 只验证 SHA-256 一致，却没有验证 digest 是否来自可信 manifest；
• 参数加载正确，但 tokenizer、RoPE、Transformer block/vocabulary、adapter 或 量化配置不匹配，模型语义仍错误。

实际放置完成后，还需用代表性前向验证 dtype/layout 能被预期 backend 消费；否则加载器可能成功，运行时却 fallback 或首次请求才触发昂贵 repack。模型加载与 backend dispatch 因此相邻，但一个负责 storage/ownership，另一个负责算子实现选择。

MoE/Expert Parallel还要求 checkpoint manifest 能把 expert id 映射到 owning rank；只按层号均分文件可能让每个 rank 读取不需要的 experts，或让 expert/TP shard 轴错位。

Multi-LoRA Serving在共享 base 之外动态加载小 adapter shards；需要同样的安全格式、hash、shape/target module 校验与 TP placement，且 adapter lifecycle 比 base checkpoint 更频繁。

参考资料

• Safetensors format
• Safetensors — Torch API
• Safetensors — Tensor slicing
• Hugging Face Transformers — Big Model Inference
• PyTorch — CUDA memory management