SafeTensors: 安全高效的张量存储与分发方案

SafeTensors简介

SafeTensors是由Hugging Face开发的一种新的张量存储格式,旨在为机器学习模型提供一种安全且高效的权重保存和加载方案。它的主要目标是取代PyTorch中默认使用的pickle格式,同时保持快速的加载速度和零拷贝的特性。

SafeTensors的核心优势包括:

1. 安全性: 与pickle格式不同,SafeTensors不会执行任意代码,从而避免了潜在的安全风险。
2. 高效性: 支持零拷贝和懒加载,可以显著提高大型模型的加载速度。
3. 跨平台兼容性: 同时支持Python和Rust实现,可以在不同的环境中使用。
4. 文件大小无限制: 不像某些格式存在文件大小限制,SafeTensors可以处理任意大小的模型。
5. 支持新兴数据类型: 原生支持bfloat16和fp8等新兴数据类型,无需额外的转换。

SafeTensors的技术细节

文件格式

SafeTensors文件由以下部分组成:

1. 头部大小: 8字节的无符号小端64位整数,表示头部的大小。
2. JSON头部: 包含张量的元数据信息,如数据类型、形状和数据偏移量等。
3. 数据缓冲区: 存储实际的张量数据。

这种设计允许快速访问文件中的特定张量,而无需加载整个文件,这对于分布式环境中的部分加载特别有用。

安全性考虑

SafeTensors格式的设计充分考虑了安全性:

• 不允许执行任意代码,避免了pickle等格式的安全隐患。
• 限制头部大小(最大100MB),防止解析过大的JSON导致的潜在DOS攻击。
• 保证文件中的地址不会重叠,避免在加载过程中超出文件大小的内存分配。

与其他格式的对比

SafeTensors并非是第一个尝试解决这个问题的格式。让我们来看看它与其他常见格式的对比:

从表中可以看出,SafeTensors在大多数关键特性上都表现出色,尤其是在安全性、效率和功能性的平衡上做得很好。

SafeTensors的使用

安装

SafeTensors可以通过pip轻松安装:

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 .

基本用法

以下是使用SafeTensors保存和加载张量的简单示例:

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")

# 加载张量
loaded_tensors = {}
with safe_open("model.safetensors", framework="pt", device="cpu") as f:
   for key in f.keys():
       loaded_tensors[key] = f.get_tensor(key)

这个例子展示了SafeTensors的简单性和直观性。它不仅易于使用,还能确保在加载过程中的安全性。

SafeTensors的性能优势

SafeTensors在性能方面表现出色,尤其是在处理大型模型时:

1. 快速加载: 与pickle相比,SafeTensors在CPU上的加载速度极快。在GPU上,加载速度与PyTorch相当或更快。

2. 内存效率: 通过使用torch.UntypedStorage.from_file,SafeTensors可以绕过CPU上的额外拷贝,进一步提高效率。

3. 懒加载支持: 在分布式环境中,SafeTensors允许只加载部分张量,这大大加快了模型初始化速度。例如,使用SafeTensors格式可以将BLOOM模型在8个GPU上的加载时间从10分钟缩短到45秒。

SafeTensors的未来发展

虽然SafeTensors已经在许多方面表现出色,但它仍在不断发展和改进。一些潜在的未来方向包括:

1. 进一步优化性能,特别是在大规模分布式环境中。
2. 增加对更多新兴数据类型的支持。
3. 改进与其他深度学习框架的集成。
4. 提供更多的工具和实用程序,以便更容易地在现有项目中采用SafeTensors。

结论

SafeTensors代表了机器学习模型权重存储和分发的重要进步。通过提供一个安全、高效且易于使用的格式,它解决了长期以来困扰研究人员和工程师的诸多问题。随着深度学习模型变得越来越大、越来越复杂,像SafeTensors这样的工具将在确保模型的安全分发和高效加载方面发挥关键作用。

对于那些关心模型安全性、加载性能和跨平台兼容性的开发者和研究人员来说,SafeTensors无疑是一个值得考虑的选择。随着它的不断发展和完善,我们可以期待看到它在机器学习社区中得到更广泛的采用。

要了解更多关于SafeTensors的信息,可以访问其GitHub仓库或查阅官方文档。无论您是在开发新的机器学习模型,还是在优化现有的部署流程,SafeTensors都可能为您提供一个有价值的工具,帮助您更安全、更高效地管理模型权重。