Safetensors 模型加载速度测试指南

同样一份128 MiB的模型权重，在这台Mac上，Safetensors返回张量字典的中位耗时是1.378毫秒，普通的torch.load要15.058毫秒；要是给PyTorch也开内存映射，中位耗时就降到1.724毫秒。做这类测试记录的时候，必须同时写清楚四个信息：文件格式、加载参数、操作系统缓存情况、实际触碰了多少数据。光抄个“快了多少倍”的数字，根本没什么参考价值。

先搞清楚测试的终点和边界

咱们这次测试的对比对象，是三个CPU端的加载入口：safetensors.torch.load_file、普通torch.load，以及开了mmap=True的torch.load。三种方法读取的都是同一组张量存出来的文件，每种方法先预热一次，再测7轮；每轮轮换方法的测试顺序，减少固定先后次序带来的偏差。

计时是从调用加载函数开始，到读完首尾两个样本值截止。它测的是“暖缓存状态下，加载API返回结果+触碰少量数据一共要多久”，既不是冷盘读取的时间，也不是模型跑完第一次推理的时间。不管是Safetensors还是PyTorch的内存映射模式，都有可能延迟读取张量页面，所以千万别把这几毫秒的数字，当成整个模型已经完完整整进物理内存了。

步骤一：搭个独立的Python环境

操作入口是空目录里的终端。macOS或者Linux系统跑下面的命令就行；要是用Windows PowerShell，把激活命令换成.venvScriptsActivate.ps1：

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip torch safetensors numpy

装完之后，跑个版本检查：

python -c "import torch, safetensors; print(torch.__version__); print(safetensors.__version__)"

成功的标志是能正常输出两行版本号，而且不会报ModuleNotFoundError。要是安装的时候找不到兼容的安装包，先确认下Python版本和系统架构，再去PyTorch安装页选对应平台的命令；要是导入PyTorch的时候提示缺NumPy，就在当前这个虚拟环境里补装numpy，别切到全局的Python去装。

图里的操作入口是笔记本的第二个代码单元。版本、系统、CPU测试状态都已经输出了，末尾的SUCCESS说明依赖都能正常导入；就算MPS显示可用，也不影响咱们这轮的测试标准，因为所有加载操作都明确指定用CPU。

步骤二：生成两份内容一模一样的权重文件

操作入口是虚拟环境里的Python脚本。咱们先生成8个2048 × 2048的float32张量，总原始数据量是128 MiB，再分别存成Safetensors和PyTorch格式的文件：

from pathlib import Path
import torch
from safetensors.torch import save_file

data_dir = Path("benchmark_data")
data_dir.mkdir(exist_ok=True)

tensor_count = 8
elements_per_tensor = 2048 * 2048
tensors = {}

for index in range(tensor_count):
    tensor = torch.arange(elements_per_tensor, dtype=torch.float32)
    tensors[f"layer_{index:02d}"] = tensor.reshape(2048, 2048).add_(index)

save_file(
    tensors,
    data_dir / "model.safetensors",
    metadata={"tensor_count": "8", "total_mib": "128"},
)
torch.save(tensors, data_dir / "model.pt")

成功的标志是目录里多出两个大概128 MiB的文件，而且能分别算出它们的SHA256。要是save_file报错说张量不连续，先给对应的张量调用contiguous()就行；要是磁盘空间不够，可以减少张量数量或者缩小边长，但两种格式必须用同一组张量，不能各自随机生成。

文件大小差不多，只能说明存的内容规模相当，证明不了内容是可读的。图里同时记录了张量数量、形状、数据类型和文件哈希，后面要重复测试的话，就能核对输入是不是一致了。

步骤三：计时前先检查键名和样本值对不对

操作入口是刚生成的model.safetensors文件。咱们用safe_open读取文件头、键名还有首尾的样本值：

from safetensors import safe_open

with safe_open(
    "benchmark_data/model.safetensors",
    framework="pt",
    device="cpu",
) as handle:
    keys = list(handle.keys())
    metadata = handle.metadata()
    first = handle.get_tensor(keys[0])
    last = handle.get_tensor(keys[-1])
    edge_sample = float(first[0, 0]) + float(last[-1, -1])

print(keys)
print(first.shape, first.dtype)
print(metadata)
print(edge_sample)

成功的标志是键名从layer_00一直连续到layer_07，第一个张量的形状是2048 × 2048、类型是float32，样本值算出来是4194310.0。要是报SafetensorError，先把不完整的文件删掉重新生成；要是键名或者样本值对不上，就别往下做计时了——不然两种加载方法比的根本不是同一个输入，测了也白测。

这一步不算在最终的计时里。它把“文件存在”这个最低标准，提升到了“键名、形状、类型、边界值都能正常读取”，免得你在损坏的文件或者错误的路径上，测出个看起来正常其实没用的耗时。

步骤四：预热后轮换顺序测7轮

操作入口是刚才通过检查的两份权重文件。三个加载函数都统一用CPU、只加载权重、用同样的方式读取样本：

import gc
import statistics
import time
import torch
from safetensors.torch import load_file

safe_path = "benchmark_data/model.safetensors"
torch_path = "benchmark_data/model.pt"

def load_safetensors():
    return load_file(safe_path, device="cpu")

def load_pytorch():
    return torch.load(
        torch_path,
        map_location="cpu",
        weights_only=True,
    )

def load_pytorch_mmap():
    return torch.load(
        torch_path,
        map_location="cpu",
        weights_only=True,
        mmap=True,
    )

loaders = {
    "safetensors": load_safetensors,
    "pytorch": load_pytorch,
    "pytorch_mmap": load_pytorch_mmap,
}

def timed(name):
    gc.collect()
    started = time.perf_counter_ns()
    data = loaders[name]()
    sample = float(data["layer_00"][0, 0]) + float(data["layer_07"][-1, -1])
    elapsed_ms = (time.perf_counter_ns() - started) / 1_000_000
    if sample != 4194310.0:
        raise RuntimeError("sample mismatch")
    return elapsed_ms

for name in loaders:
    timed(name)

names = list(loaders)
measurements = {name: [] for name in names}

for round_index in range(7):
    offset = round_index % len(names)
    order = names[offset:] + names[:offset]
    for name in order:
        measurements[name].append(timed(name))

for name, values in measurements.items():
    print(name, statistics.median(values), min(values))

成功的标志是7轮测试下来，三种方法都能得到正的耗时数，而且样本校验没报错。要是普通torch.load提示安全问题，先确认代码里是不是留了weights_only=True，而且测试文件是你自己生成的；要是旧版PyTorch不支持mmap参数，要么先升级PyTorch，要么就把这组对照删掉，并且在报告里写清楚版本限制，别悄咪咪把它当成普通加载来算。

从原始数据行能看到，Safetensors每轮大概1.34至1.42毫秒，普通PyTorch大概14.14至15.90毫秒，PyTorch开内存映射的话大概1.69至1.76毫秒。轮换顺序没法消除所有系统噪声，但三种方法在每一轮里的位置都不一样，能减少固定顺序带来的偏差。

步骤五：看结果要盯中位数，别只看最快的那轮

操作入口是7轮测试的原始数据。把每种方法的中位数、最小值、高位耗时都写到JSON里，同时算一下PyTorch耗时除以Safetensors耗时的比值。成功的标志是报告能从保存好的JSON里重新读出来，而不是靠屏幕上临时打印的那一行数字。

在这台机器上，普通PyTorch的中位耗时是Safetensors的10.92 倍，开了PyTorch内存映射之后，这个比值变成1.25 倍。前一个差距主要来自普通反序列化和内存拷贝的成本，后一个对比就更接近两种内存映射路径的API返回开销了。只有7轮测试的话，高位耗时适合用来发现异常抖动，还代表不了稳定的生产环境P95水平。

结果和官方示例不一样？这么查

• Safetensors没快多少：先看看你是不是把全张量求和、模型构建、设备传输这些操作也算进计时里了——这些动作会把测试的问题都给改了。
• 第一次跑特别慢：导入库、动态链接、磁盘页缓存这些开销都可能落在第一轮，所以预热的数值要单独记，别和正式的7轮混在一起算。
• 每轮波动特别大：把同时在跑的大型任务都关了，多测几轮，同时把原始数据都存下来，别只留个最小值。
• 想测冷启动速度：得单独设计冷缓存实验，还要记录清楚存储介质、文件系统、缓存清理方法。不同的系统没法用同一条命令安全清理缓存。
• 想测GPU场景：把设备传输和CUDA同步都写到单独的脚本里，还要记录GPU型号、驱动版本、CUDA版本；CPU的测试结果不能直接套到GPU上。
• 想测真实上线的速度：把模型实例化、权重绑定、第一次推理都算进端到端计时里，同时单独保留文件加载这一项的耗时。

完成版检查清单

• 虚拟环境里能正常导入PyTorch、Safetensors和NumPy，版本和系统信息都已经记录好了。
• 两种格式的文件都是用同一组张量生成的，张量的数量、形状、类型、文件大小还有SHA256都已经存好了。
• Safetensors文件的键名、元数据、形状、边界样本值都能正常读取。
• 三个加载入口都用同一个CPU、同一种样本触碰方式、同一个计时器。
• 每种方法都先预热一次，正式轮次轮换执行顺序，原始的毫秒数都没丢。
• 报告用中位数来解释结果，而且明确区分了普通PyTorch和开mmap=True的情况。
• 结论里注明了暖缓存、少量数据触碰的边界，没把本机测出来的倍数说成所有设备都适用的固定结论。
• 五张Jupyter步骤截图都能正常打开，分别对应环境配置、文件生成、内容检查、原始计时、汇总结果这几个步骤。

速度测试真正有价值的地方，不是测出个好看的倍数，而是让别人知道这个倍数是在什么样的加载路径、缓存条件、数据触碰范围下得出来的。把脚本、输入哈希、原始轮次数据、环境版本都保留好，以后换电脑、换模型、升级库的时候，才能知道变化到底来自哪里。