Safetensors：新一代安全高效的模型权重存储格式详解
在日常 AI 模型训练与部署过程中，模型权重的存储格式至关重要。目前最流行的深度学习框架 PyTorch 默认使用 pickle 格式存储模型权重文件（.pth）。然而，PyTorch 官方文档明确指出：使用 torch.load() 保存模型时，除非将 weights_only 参数设置为 True，否则隐式使用 pickle 模块。pickle 是不安全的，可以构造恶意的 pickle 数据，在反序列化期间执行任意代码。
为了规避此类安全风险，Hugging Face 推出了新的权重存储格式 Safetensors。本文将深入解析 Safetensors 的内部机制、安全性优势、零拷贝技术原理以及在 PyTorch 中的实践应用。
一、Safetensors 简介
Safetensors 是一种用于安全地存储张量的新格式。它的设计目标是在保持简单性的同时实现高性能（支持零拷贝读取）。作为 pickle 格式的替代品，Safetensors 从根本上解决了反序列化漏洞问题，因为它不包含可执行的 Python 代码对象，仅存储纯数值数据。
1.1 核心特性
- 安全性：无法执行恶意代码，加载过程仅涉及内存映射和数据解析。
- 速度：支持零拷贝（Zero-Copy）读取，减少 CPU 和内存带宽消耗。
- 惰性加载：允许在不加载整个文件的情况下检查元数据或读取特定张量。
- 跨语言支持：基于 C++ 实现，提供 Python、Rust、C# 等多种语言的绑定。
二、Safetensors 内部文件格式
Safetensors 文件结构非常清晰，主要由两部分组成：头部信息（Header）和数据部分（Data）。假设有一个名为 model.safetensors 的文件，其内部结构如下：
- 文件头长度：前 8 个字节是一个无符号的 64 位整数（Little-endian），表示 JSON 头部信息的字节数。
- JSON 头部：紧接着的 N 个字节是 UTF-8 编码的 JSON 字符串，包含所有张量的元数据（如数据类型、形状、数据偏移量）以及可选的全局元数据（如
__metadata__）。 - 原始数据：文件的剩余部分直接存储模型权重 tensor 的二进制值，按照头部中定义的顺序排列。
2.1 头部元数据结构示例
以 GPT2 模型的 Safetensors 文件为例，其头部 JSON 内容通常包含以下字段：
{
"__metadata__": { "format": "pt" },
"h.10.ln_1.weight": {
"dtype": "F32",
"shape": [768],
"data_offsets": [223154176, 223157248]
}
}
其中：
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 的优势，我们从以下几个维度对比常见的模型存储格式：
| 特性 | Pickle (.pth) | Safetensors (.safetensors) | ONNX (.onnx) |
|---|---|---|---|
| 安全性 | 低（存在反序列化风险） | 高（纯数据，无代码执行） | 中（图定义，通常安全） |
| 零拷贝 | 否（需反序列化复制） | 是（内存映射） | 否 |
| 延迟加载 | 困难（需加载全部） | 支持（按 key 读取） | 支持（部分算子） |
| 布局控制 | 固定 | 灵活（可优化访问模式） | 固定（计算图） |
| 文件大小限制 | 受限于 Python 对象 | 无硬性限制 | 受限于 protobuf 大小 |
| 灵活性 | 高（可存自定义对象） | 中（仅张量 + 元数据） | 高（通用计算图） |
| Bfloat16 支持 | 需转换 | 原生支持 | 原生支持 |
3.1 Safetensors 与 ONNX 的区别
虽然两者都用于模型存储，但用途不同：
- Safetensors：专注于训练后权重的存储与快速加载。它是 Python
pickle的安全替代品，适用于同一框架内（如 PyTorch）的高效推理。 - ONNX：专注于跨框架模型交换。它将模型转换为中间表示形式，使得 PyTorch 训练的模型可以在 TensorFlow 或 C++ 环境中运行。
四、零拷贝技术深度解析
Safetensors 性能提升的核心在于**零拷贝（Zero-Copy）**技术。传统的文件读取流程通常涉及多次内存拷贝：磁盘 -> 内核缓冲区 -> 用户空间缓冲区 -> 应用程序变量。每次跨越用户空间和内核空间的边界都需要进行数据复制，消耗 CPU 周期和内存带宽。
4.1 传统读取流程
- 系统调用
read()：内核将数据从磁盘复制到内核缓冲区。 - 系统调用
memcpy()：内核将数据从内核缓冲区复制到用户空间缓冲区。 - 应用程序处理数据。 此过程至少发生两次数据拷贝，并伴随上下文切换。
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 格式进行分发，特别是在涉及公共模型仓库的场景下，这不仅能保护用户免受潜在的攻击，还能显著提升推理服务的响应速度。
