Safetensors：新一代安全高效的模型权重存储格式详解

其中：

• dtype：指定张量的数据类型（如 F32, I32, BF16 等）。
• shape：张量的维度列表。
• data_offsets：一个包含两个整数的数组，分别表示该张量数据在二进制流中的起始位置和结束位置（相对于文件开头减去头部后的偏移，或者绝对偏移，具体取决于实现，通常指相对于文件开始位置的偏移量减去 header 长度）。这使得程序可以直接通过 mmap 定位到特定张量的内存地址。

2.2 解析文件头的代码示例

我们可以通过 Python 标准库来解析 Safetensors 文件的头部，无需依赖第三方库即可验证文件格式：

import struct
import json

def parse_safetensors_header(file_path):
    with open(file_path, 'rb') as f:
        # 读取前 8 字节获取 header 长度
        header_len_bytes = f.read(8)
        if len(header_len_bytes) < 8:
            raise ValueError("File too small")
        
        header_len = struct.unpack('<Q', header_len_bytes)[0]
        
        # 读取 header 内容
        header_json = f.read(header_len)
        header = json.loads(header_json.decode('utf-8'))
        
        return header

# 示例调用
header = parse_safetensors_header('model.safetensors')
print(f"Detected tensors: {len([k for k in header.keys() if not k.startswith('__')])} tensors")

三、不同模型权重格式对比分析

为了更直观地理解 Safetensors 的优势，我们从以下几个维度对比常见的模型存储格式：

3.1 Safetensors 与 ONNX 的区别

虽然两者都用于模型存储，但用途不同：

• Safetensors：专注于训练后权重的存储与快速加载。它是 Python pickle 的安全替代品，适用于同一框架内（如 PyTorch）的高效推理。
• ONNX：专注于跨框架模型交换。它将模型转换为中间表示形式，使得 PyTorch 训练的模型可以在 TensorFlow 或 C++ 环境中运行。

四、零拷贝技术深度解析

Safetensors 性能提升的核心在于**零拷贝（Zero-Copy）**技术。传统的文件读取流程通常涉及多次内存拷贝：磁盘 -> 内核缓冲区 -> 用户空间缓冲区 -> 应用程序变量。每次跨越用户空间和内核空间的边界都需要进行数据复制，消耗 CPU 周期和内存带宽。

4.1 传统读取流程

1. 系统调用 read()：内核将数据从磁盘复制到内核缓冲区。
2. 系统调用 memcpy()：内核将数据从内核缓冲区复制到用户空间缓冲区。
3. 应用程序处理数据。 此过程至少发生两次数据拷贝，并伴随上下文切换。

4.2 Safetensors 的零拷贝实现

Safetensors 利用操作系统的**内存映射（Memory Mapping, mmap）**功能。当打开 .safetensors 文件时，操作系统直接将文件内容映射到进程的虚拟地址空间。

• CPU 角色转变：CPU 不再负责搬运数据，而是直接通过指针访问映射后的内存区域。
• 效率提升：减少了不必要的内存分配和数据复制，显著降低了加载大型模型时的延迟和内存占用。
• 适用场景：特别适用于大模型推理场景，其中只需加载部分层或特定张量即可启动服务。

注意：零拷贝并非不进行拷贝。如果数据不在物理内存中，操作系统仍需将其从磁盘调入内存（Page Fault），但这部分工作由操作系统内核管理，对应用程序透明且高效。

五、Safetensors 实践指南

5.1 环境安装

pip install safetensors

5.2 保存与加载张量

保存张量

import torch
from safetensors.torch import save_file

tensors = {
    "embedding": torch.zeros((2, 2)),
    "attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")

加载张量

from safetensors import safe_open

tensors = {}
with safe_open("model.safetensors", framework="pt", device='cpu') as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)
print(tensors)

5.3 惰性加载（Lazy Loading）

对于超大模型，我们可能不需要一次性加载所有权重。Safetensors 支持切片读取：

from safetensors import safe_open

with safe_open("model.safetensors", framework="pt", device='cpu') as f:
    # 获取 embedding 层的切片
    tensor_slice = f.get_slice("embedding")
    vocab_size, hidden_dim = tensor_slice.get_shape()
    print(f"Shape: {vocab_size} x {hidden_dim}")
    
    # 仅读取部分数据，例如前 100 个词向量
    tensor = tensor_slice[:, :100]
    print(tensor.shape)

5.4 完整模型保存与加载

除了单个张量，Safetensors 也支持直接保存和加载 PyTorch 模型实例。

from torchvision.models import resnet18
from safetensors.torch import load_model, save_model
import torch

# 1. 创建预训练模型
model_pt = resnet18(pretrained=True)

# 2. 保存为 Safetensors 格式
save_model(model_pt, "resnet18.safetensors")

# 3. 加载模型（无需手动 load_state_dict）
model_st = resnet18(pretrained=False)
load_model(model_st, "resnet18.safetensors")

# 4. 验证一致性
img = torch.randn(2, 3, 224, 224)
model_pt.eval()
model_st.eval()

with torch.no_grad():
    result = torch.allclose(model_pt(img), model_st(img))
    print(f"Results match: {result}")

六、最佳实践与注意事项

6.1 设备管理

在使用 safe_open 时，务必指定 device 参数。如果不指定，张量将保留在 CPU 上，后续可能需要手动移动到 GPU，这会增加额外的拷贝开销。

# 推荐：直接在加载时指定设备
tensor = f.get_tensor("weight", device='cuda:0')

6.2 数据类型兼容性

确保源框架和目标框架支持相同的数据类型。虽然 Safetensors 支持多种 dtype，但在跨框架迁移时需注意精度损失（如 FP32 转 FP16）。

6.3 HuggingFace 集成

目前，Hugging Face 的 transformers 库已原生支持 Safetensors。在下载模型时，可以通过设置环境变量或配置优先使用 .safetensors 文件：

export HF_HUB_ENABLE_HF_TRANSFER=1
# 或使用 transformers 库自动选择
from transformers import AutoModel
model = AutoModel.from_pretrained("bert-base-uncased", trust_remote_code=True)

七、总结

Safetensors 作为一种新兴的模型权重存储格式，凭借其安全性、加载速度和内存效率，正在成为 AI 领域的事实标准之一。它有效解决了 pickle 带来的安全隐患，并通过零拷贝技术优化了大模型的推理性能。随着 HuggingFace 等主流平台的广泛采用，开发者应尽早掌握 Safetensors 的使用规范，以提升模型部署的可靠性与效率。

在实际项目中，建议将训练好的模型统一转换为 Safetensors 格式进行分发，特别是在涉及公共模型仓库的场景下，这不仅能保护用户免受潜在的攻击，还能显著提升推理服务的响应速度。