Safetensors 模型加载速度测试指南

在Mac平台上，128 MiB的模型权重加载测试显示，Safetensors中位耗时1.378毫秒，普通torch.load约15.058毫秒，开启内存映射后降至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，避免切换到全局环境。

图示中操作入口为笔记本第二个代码单元，版本、系统及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水平。

结果与官方示例不符？可按照以下方法排查

1. Safetensors提升不明显：检查是否将全张量求和、模型构建、设备传输等操作计入计时，这些步骤会改变测试目标。
2. 第一次跑特别慢：导入库、动态链接、磁盘页缓存等开销可能落在第一轮，因此预热数值需单独记录，不可与正式7轮混合计算。
3. 每轮波动特别大：关闭同时运行的大型任务，增加测试轮数，并保存原始数据，而非仅保留最小值。
4. 想测冷启动速度：需单独设计冷缓存实验，并记录存储介质、文件系统及缓存清理方法。不同系统无法使用同一命令安全清理缓存。
5. 想测GPU场景：将设备传输与CUDA同步写入独立脚本，并记录GPU型号、驱动版本及CUDA版本；CPU测试结果不可直接套用于GPU。
6. 想测真实上线的速度：将模型实例化、权重绑定及首次推理纳入端到端计时，同时单独保留文件加载耗时。

完成版检查清单

1. 虚拟环境中能正常导入PyTorch、Safetensors和NumPy，版本和系统信息已记录。
2. 两种格式文件均使用同一组张量生成，张量数量、形状、类型、文件大小及SHA256已保存。
3. Safetensors文件的键名、元数据、形状及边界样本值可正常读取。
4. 三个加载入口均使用同一CPU、相同样本触碰方式及同一计时器。
5. 每种方法先预热一次，正式轮次轮换执行顺序，原始毫秒数未丢失。
6. 报告使用中位数解释结果，并明确区分普通PyTorch与开启mmap=True的情况。
7. 结论注明暖缓存、少量数据触碰的边界，未将本机测试倍数视为适用于所有设备的固定结论。
8. 五张Jupyter步骤截图均可正常打开，分别对应环境配置、文件生成、内容检查、原始计时及汇总结果步骤。

速度测试的核心价值在于明确测试边界，而非单纯追求高倍率。保留脚本、哈希、原始数据与环境版本，才能在设备或模型更换时准确追溯变化来源。