Python

Rust

safetensors

Safetensors

此仓库实现了一种新的简单格式，用于安全地存储张量（与 pickle 相反）且仍然快速（零拷贝）。

安装

Pip

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

pip install safetensors

从源码

对于源代码，您需要 Rust

# Install Rust
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
# Make sure it's up to date and using stable channel
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 可能存在于文件中
• 允许空张量（维度为 1 的张量大小为 0）。它们在数据缓冲区中不存储任何数据，但在头部中保留大小。它们并不真正带来很多值，但由于从传统的张量库（torch、tensorflow、numpy 等）的角度来看是有效的张量，因此被接受。
• 允许 0 级张量（形状为 [] 的张量），它们只是一个标量。
• 字节缓冲区需要完全索引，不能包含空白。这可以防止创建多语言文件。
• 字节序：小端。
• 顺序：'C' 或行主序。

还有另一种格式？

这个库的主要理由是去除在 PyTorch 中使用 pickle 的需要，这是默认使用的。还有其他由机器学习和更通用的格式使用的格式。

让我们看看替代方案以及为什么这个格式被认为是有趣的。这是我非常个人且可能存在偏见的看法。

• Safe: 我能否使用随机下载的文件，并期望不会运行任意代码？
• Zero-copy: 读取文件是否需要比原始文件更多的内存？
• Lazy loading: 我能否在不加载所有内容的情况下检查文件？以及在分布式设置下只加载文件中的某些张量而无需扫描整个文件（分布式设置）？
• Layout control: Lazy loading可能不足以满足需求，因为如果张量的信息散布在您的文件中，那么即使信息是懒加载的，您可能仍然需要访问大部分文件来读取可用的张量（导致许多DISK -> RAM复制）。控制布局以保持对单个张量的快速访问非常重要。
• No file size limit: 文件大小有限制吗？
• Flexibility: 我能否以格式保存自定义代码，并在以后无需额外代码的情况下使用它？（~表示我们可以存储比纯张量更多的内容，但不能存储自定义代码）
• Bfloat16/Fp8: 格式支持本地的bfloat16/fp8（这意味着不需要奇怪的解决方案）？这在机器学习领域变得越来越重要。

主要反对意见

• Pickle: 不安全，会运行任意代码
• H5: 看起来现在不鼓励用于TF/Keras。实际上似乎很合适。一些经典的使用后免费问题：[链接](https://www.cvedetails.com/vulnerability-list/vendor_id-15991/product_id-35054/Hdfgroup-Hdf5.html)。在安全性方面与pickle相比处于非常不同的水平。此外，210k行代码与当前库的约400行代码相比。
• SavedModel: Tensorflow特定（它包含TF图信息）。
• MsgPack: 没有布局控制以启用懒加载（对于分布式设置中加载特定部分非常重要）
• Protobuf: 有硬限制的最大文件大小
• Cap'n'proto: 不支持Float16 链接，因此需要使用字节缓冲区的手动包装器。布局控制似乎可能，但并非易事，因为缓冲区有局限性 链接。
• Numpy (npz): 不支持 bfloat16。容易受到zip炸弹（DOS）的影响。不是zero-copy。
• Arrow: 不支持 bfloat16。

注意

• Zero-copy: 在机器学习中没有真正的zero-copy格式，它需要从磁盘到RAM/GPU RAM（这需要时间）。在CPU上，如果文件已经在缓存中，则可以真正实现zero-copy，而在GPU上没有这样的磁盘缓存，因此总是需要复制，但可以绕过在任何给定点上分配所有张量在CPU上。SafeTensors对于标题不是zero-copy。选择JSON相当随意，但鉴于反序列化所需的时间远远小于加载实际张量数据所需的时间，并且是可读的，所以我选择了这种方式（此外，空间远远小于张量数据）。

• Endianness: Little-endian。这可以在以后修改，但现在感觉非常不必要。

• Order: 'C'或行主序。这似乎已经获胜。如果需要，我们可以稍后添加这些信息。

• Stride: 无步进，所有张量在序列化之前都需要打包。我还没有看到任何情况表明在序列化格式中存储具有步进的张量是有用的。

优点

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

• 防止DOS攻击：我们可以设计格式，使其几乎不可能使用恶意文件对用户进行DOS攻击。目前，为了防止解析非常大的JSON，对标题大小有限制为100MB。此外，在读取文件时，保证文件中的地址不会以任何方式重叠，这意味着在加载文件时，您不应该超过文件在内存中的大小

• 更快的加载：PyTorch似乎在主要的机器学习格式中是最快的文件加载方式。然而，它似乎在CPU上有一个额外的副本，我们可以通过使用torch.UntypedStorage.from_file来绕过这个库。目前，与pickle相比，使用这个库的CPU加载时间非常快。GPU加载时间与PyTorch相当或更快。首先在CPU上使用torch进行内存映射加载，然后将所有张量移动到GPU上似乎也更快（类似于torch pickle的行为）

• 懒加载：在分布式（多节点或多GPU）环境中，能够只加载各种模型的部分张量是非常好的。对于BLOOM，使用此格式可以从10mn将模型加载到8个GPU上，使用常规PyTorch权重只需要45秒。这确实加快了开发模型时的反馈循环。例如，在更改分布策略（例如管道并行与张量并行）时，不需要有单独的权重副本。

许可证：Apache-2.0