safetensors

Safetensors

这个仓库实现了一种新的简单格式，用于安全地存储张量（相对于pickle而言），同时仍然保持快速（零拷贝）。

安装

Pip

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

pip install safetensors

从源码安装

对于源码，你需要安装Rust

# 安装Rust
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
# 确保它是最新的并使用稳定版本
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的张量）。它们不在数据缓冲区中存储任何数据，但在头部中保留大小。从传统张量库（torch、tensorflow、numpy等）的角度来看，它们并不带来太多价值，但被接受，因为它们是有效的张量。
允许0阶张量（形状为[]的张量），它们仅仅是标量。
字节缓冲区需要完全索引，不能包含空洞。这防止了创建多格式文件。
字节序：小端序。
顺序：'C'或行主序。

又一种格式？

这个crate的主要理由是消除在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：目前似乎不推荐用于TF/Keras。除此之外，看起来实际上是一个很好的选择。存在一些经典的使用后释放问题：https://www.cvedetails.com/vulnerability-list/vendor_id-15991/product_id-35054/Hdfgroup-Hdf5.html。在安全性方面与pickle有很大不同。此外，代码量为21万行，而本库目前仅约400行。
SavedModel：特定于Tensorflow（包含TF图信息）。
MsgPack：没有布局控制来实现懒加载（在分布式设置中加载特定部分很重要）
Protobuf：有2GB的硬性文件大小限制
Cap'n'proto：不支持Float16 链接，因此需要使用字节缓冲区的手动包装器。布局控制似乎可行但不简单，因为缓冲区有限制链接。
Numpy (npz)：不支持bfloat16。容易受到zip炸弹攻击（DOS）。非零拷贝。
Arrow：不支持bfloat16。

注意事项

零拷贝：在机器学习中，没有真正的零拷贝格式，数据需要从硬盘传输到内存/GPU内存（这需要时间）。在CPU上，如果文件已经在缓存中，那么它可以真正实现零拷贝，而在GPU上没有这样的硬盘缓存，所以总是需要一次复制，但你可以避免在任何给定时间点在CPU上分配所有张量。SafeTensors对于头部不是零拷贝的。选择JSON是相当随意的，但由于反序列化所需时间远小于加载实际张量数据，而且可读性好，我选择了这种方式（而且空间远小于张量数据）。
字节序：小端序。这可以在以后修改，但目前感觉真的没有必要。
顺序：'C'或行主序。这似乎已经成为主流。如果需要，我们可以稍后添加该信息。
步幅：没有步幅，所有张量在序列化之前需要打包。我还没有看到在序列化格式中存储步幅张量有用的情况。

优势

由于我们可以发明一种新格式，我们可以提出额外的优势：

防止DOS攻击：我们可以精心设计格式，使得几乎不可能使用恶意文件对用户进行DOS攻击。目前，对头部大小有100MB的限制，以防止解析极大的JSON。此外，在读取文件时，保证文件中的地址不会以任何方式重叠，这意味着在加载文件时，内存中不应超过文件的大小。
更快的加载：PyTorch似乎是主要机器学习格式中加载最快的。然而，它在CPU上似乎还有一次额外的复制，我们可以通过使用torch.UntypedStorage.from_file来绕过这一点。目前，与pickle相比，这个库的CPU加载时间非常快。GPU加载时间与PyTorch等效或更快。使用torch进行内存映射在CPU上先加载，然后将所有张量移动到GPU，似乎也somehow更快（与torch pickle中的行为类似）。
懒加载：在分布式（多节点或多GPU）设置中，能够在各个模型上只加载部分张量是很好的。对于BLOOM，使用这种格式可以将在8个GPU上加载模型的时间从常规PyTorch权重的10分钟缩短到45秒。这真正加快了在模型上开发时的反馈循环。例如，当改变分布策略时（例如管道并行vs张量并行），你不必有单独的权重副本。

许可证：Apache-2.0