Python

Rust

safetensors

Safetensors

本仓库实现了一种新的简单格式，用于安全地存储张量（与 pickle 不同），同时保持高效（零拷贝）。

安装

Pip

您可以通过 pip 包管理器安装 safetensors：

pip install safetensors

从源码安装

要从源码编译，您需要 Rust 环境：

# 安装 Rust
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
# 确保 Rust 已更新并使用稳定版
rustup update
git clone https://github.com/huggingface/safetensors
cd safetensors/bindings/python
pip install setuptools_rust
pip install -e .

快速入门

import torch
from safetensors import safe_open
from safetensors.torch import save_file

tensors = {
   "weight1": torch.zeros((1024, 1024)),
   "weight2": torch.zeros((1024, 1024))
}
save_file(tensors, "model.safetensors")

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

Python 文档

格式说明

• 8 字节：N，一个无符号小端序 64 位整数，表示头部的大小。
• N 字节：一个 JSON UTF-8 字符串，代表头部数据。
  • 头部数据必须以 { 字符（0x7B）开头。
  • 头部数据可以末尾填充空白字符（0x20）。
  • 头部是一个字典，形如 {"TENSOR_NAME": {"dtype": "F16", "shape": [1, 16, 256], "data_offsets": [BEGIN, END]}, "NEXT_TENSOR_NAME": {...}, ...}，
    • data_offsets 指向张量数据相对于字节缓冲区起始位置的偏移量（不是文件中的绝对位置）， 其中 BEGIN 是起始偏移量，END 是结束偏移量的下一个位置（因此张量的总字节数为 END - BEGIN）。
  • 允许使用特殊键 __metadata__ 来存储自由格式的字符串到字符串映射。不允许使用任意 JSON 数据，所有值必须是字符串。
• 文件剩余部分：字节缓冲区。

注意事项：

• 不允许重复键。并非所有解析器都会遵守此规则。
• 通常情况下，JSON 的子集由该库的 serde_json 库隐式决定。任何不常见的用法，例如整数的特殊表示、换行符和 UTF-8 字符串中的转义序列，都可能在后续版本中被修改，以确保安全性。
• 张量值不会被检查，特别是文件中可能包含 NaN 和 +/-Inf。
• 允许空张量（其中一个维度为 0 的张量）。它们在数据缓冲区中不存储任何数据，但在头部保留其尺寸信息。虽然这些张量本身没有太多实际意义，但从传统张量库（PyTorch、TensorFlow、NumPy 等）的角度来看，它们仍然是合法的张量。
• 允许 0 阶张量（形状为 [] 的张量），它们实际上就是一个标量。
• 字节缓冲区必须完全索引，不能有空洞。这可以防止创建多态文件。
• 字节顺序：小端序。
• 存储顺序：C 语言顺序或行优先。
• 注意事项：出现了一些小于 1 字节的数据类型，这使得对齐变得复杂。对于这些数据类型，可能需要非传统的 API。

又一种格式？

本库的主要目的是消除在 PyTorch 中默认使用 pickle 的需求。目前已有许多其他机器学习格式以及更通用的文件格式。

让我们来看看一些替代方案，以及为什么这种格式被认为具有吸引力。以下是我个人的观点，可能带有一定的主观性：

格式	安全性	零拷贝	延迟加载	无文件大小限制	布局控制	灵活性	Bfloat16/Fp8
pickle (PyTorch)	✗	✗	✗	🗸	✗	🗸	🗸
H5 (Tensorflow)	🗸	✗	🗸	🗸	~	~	✗
SavedModel (Tensorflow)	🗸	✗	✗	🗸	🗸	✗	🗸
MsgPack (flax)	🗸	🗸	✗	🗸	✗	✗	🗸
Protobuf (ONNX)	🗸	✗	✗	✗	✗	✗	🗸
Cap'n'Proto	🗸	🗸	~	🗸	🗸	~	✗
Arrow	?	?	?	?	?	?	✗
Numpy (npy,npz)	🗸	?	?	✗	🗸	✗	✗
pdparams (Paddle)	✗	✗	✗	🗸	✗	🗸	🗸
SafeTensors	🗸	🗸	🗸	🗸	🗸	✗	🗸

• 安全性：我能否使用随机下载的文件，并确信不会执行任意代码？
• 零拷贝：读取文件是否需要比原始文件更多的内存？
• 延迟加载：我能否在不加载整个文件的情况下查看文件内容？并且只加载文件中的部分张量，而无需扫描整个文件（适用于分布式场景）？
• 布局控制：仅仅支持延迟加载并不足够，因为如果张量的信息分散在文件中，即使信息可以延迟访问，你也可能需要访问文件的大部分内容才能读取可用的张量（从而导致大量的磁盘到内存复制操作）。因此，控制布局以实现对单个张量的快速访问非常重要。
• 无文件大小限制：文件大小是否有上限？
• 灵活性：我能否将自定义代码保存到该格式中，并在以后无需额外代码即可使用它？（“~”表示可以存储除纯张量之外的内容，但不能存储自定义代码）
• Bfloat16/Fp8：该格式是否原生支持 Bfloat16/Fp8 数据类型（即不需要任何奇怪的变通方法）？这一点在机器学习领域越来越重要。

主要竞争对手

• Pickle：不安全，会执行任意代码
• H5：据称现在已不推荐用于 TensorFlow/Keras。不过从其他方面来看，它其实非常合适。存在一些经典的“使用后释放”问题：链接。在安全性上与 pickle 完全不在一个级别。此外，HDF5 有 21 万行代码，而该库目前只有约 400 行。
• SavedModel：TensorFlow 特定格式（包含 TF 图信息）。
• MsgPack：无法控制数据布局，因此无法实现延迟加载（这在分布式场景中加载特定部分时非常重要）。
• Protobuf：有严格的文件大小上限。
• Cap’n’proto：不支持 float16 数据类型 链接，因此需要手动封装字节缓冲区。虽然理论上可以控制数据布局，但并不容易，因为缓冲区本身存在限制 链接。
• Numpy (npz)：不支持 bfloat16 数据类型。易受 zip 恶意攻击（拒绝服务攻击）。非零拷贝。
• Arrow：不支持 bfloat16 数据类型。

备注

• 零拷贝：在机器学习领域，没有任何一种格式是真正意义上的零拷贝。数据总是需要从磁盘传输到内存或 GPU 内存中（这个过程需要时间）。在 CPU 上，如果文件已经缓存在高速缓存中，则可以实现真正的零拷贝；而在 GPU 上则没有类似的磁盘缓存机制，因此无论如何都需要进行一次拷贝，不过可以在任何时候避免将所有张量先分配到 CPU 上。SafeTensors 在处理头部时并非零拷贝。选择 JSON 格式其实比较随意，但由于反序列化所消耗的时间远远小于加载实际张量数据所需的时间，且 JSON 格式可读性强，所以选择了这种方式；另外，JSON 头部占用的空间也远小于张量数据本身。

• 字节序：小端序。未来可以修改，但目前看来并无必要。

• 存储顺序：C 语言风格的行优先顺序。这一选择似乎已经确定下来，如有需要，后续还可以添加相关信息。

• 步幅：不支持步幅操作，所有张量在序列化前都需要被完全打包。目前尚未遇到过在序列化格式中存储带步幅张量的实际需求。

• 小于 1 字节的数据类型：如今数据类型可以小于 1 字节，这使得对齐和寻址变得复杂。现阶段，当操作触发未对齐的读取时，库会直接报错。未来可能会为这些非标准操作设计更复杂的 API。

优势

由于我们可以自定义一种新的格式，因此可以提出以下额外优势：

• 防止 DOS 攻击：我们可以通过精心设计格式，几乎杜绝利用恶意文件发起 DOS 攻击的可能性。目前，为了防止解析超大 JSON 文件，头部大小被限制在 100MB 以内。此外，在读取文件时，保证文件中的地址不会有任何重叠，这意味着在加载文件时，内存占用永远不会超过文件本身的大小。

• 加载速度更快：在主流的机器学习文件格式中，PyTorch 的加载速度似乎最快。然而，PyTorch 在 CPU 上仍会进行一次额外的拷贝操作，而本库可以通过 torch.UntypedStorage.from_file 来避免这一拷贝。目前，使用本库在 CPU 上加载的速度比 pickle 快得多；而在 GPU 上，加载速度则与 PyTorch 相当甚至更快。先用 torch 的内存映射功能在 CPU 上加载数据，然后再将所有张量移动到 GPU 上，这种方法似乎也能进一步提升加载速度（与 PyTorch 的 pickle 格式表现类似）。

• 延迟加载：在分布式（多节点或多 GPU）环境中，能够只加载各个模型所需的部分张量是非常有用的。例如，对于 BLOOM 模型，使用这种格式后，原本需要用常规 PyTorch 权重在 8 张 GPU 上加载 10 分钟的流程，现在只需 45 秒即可完成。这大大加快了模型开发过程中的反馈循环。比如，在调整分布式策略（如流水线并行与张量并行之间切换）时，无需再为每种策略单独保存一份权重副本。

许可证：Apache-2.0

Safetensors 快速上手指南

Safetensors 是一种用于安全、快速存储张量（Tensor）的新格式，旨在替代存在安全风险的 Pickle 格式，同时支持零拷贝（zero-copy）和懒加载。

环境准备

• 系统要求：支持 Linux, macOS, Windows。
• 前置依赖：
  • Python 环境。
  • （仅源码安装需要）Rust 工具链。

安装步骤

推荐使用 pip 进行安装，简单且稳定。

方式一：通过 Pip 安装（推荐）

pip install safetensors

提示：国内用户若下载缓慢，可使用清华源或阿里源加速： pip install safetensors -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

如果你需要最新特性或参与开发，需先安装 Rust。

# 1. 安装 Rust
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
# 确保 Rust 更新至稳定版
rustup update

# 2. 克隆仓库并安装
git clone https://github.com/huggingface/safetensors
cd safetensors/bindings/python
pip install setuptools_rust
pip install -e .

基本使用

以下示例展示如何使用 PyTorch 保存和加载 .safetensors 文件。

import torch
from safetensors import safe_open
from safetensors.torch import save_file

# 1. 创建示例张量数据
tensors = {
   "weight1": torch.zeros((1024, 1024)),
   "weight2": torch.zeros((1024, 1024))
}

# 2. 保存为 safetensors 文件
save_file(tensors, "model.safetensors")

# 3. 加载文件
tensors = {}
# framework="pt" 表示加载为 PyTorch 张量, device="cpu" 指定设备
with safe_open("model.safetensors", framework="pt", device="cpu") as f:
   for key in f.keys():
       tensors[key] = f.get_tensor(key)

# 此时 tensors 字典中已包含加载后的张量
print(tensors["weight1"].shape)

核心优势说明：

• 安全性：不包含可执行代码，避免恶意 pickle 文件攻击。
• 高性能：支持内存映射（mmap），实现零拷贝加载，显著减少 CPU/GPU 内存占用并提升加载速度。
• 懒加载：可以只读取文件中特定的张量，无需加载整个模型，特别适合分布式训练场景。