深入了解 大模型参数文件存储格式-Safetensors

什么是 Safetensors？

Safetensors 是一种专为存储和加载大型张量（tensors）而设计的文件格式。它由 Hugging Face 团队开发，旨在提供一种高效、安全且易于使用的方式来管理机器学习模型的权重和其他相关数据。Safetensors 文件通常以 .safetensors 为扩展名。

设计背景

在深度学习中，模型的权重和参数通常以二进制文件的形式存储，以便在推理和训练过程中快速加载。然而，传统的存储格式如 PyTorch 的 .pt 或 .pth 文件存在以下问题：

1. 安全性：这些文件通常使用序列化（serialization）技术，如 Python 的 Pickle 模块。这种方法允许在加载文件时执行任意代码，带来了潜在的安全风险，尤其是在处理来自不可信来源的文件时。
2. 效率：虽然这些格式适用于大多数应用场景，但在处理极大规模的模型文件时，加载速度和内存效率可能成为瓶颈。
3. 可移植性：不同框架和工具对模型文件格式的支持程度不一，导致跨平台和跨框架的模型共享变得复杂。

Safetensors 针对这些问题进行了优化，提供了一种更加安全和高效的存储解决方案。

Safetensors 的设计理念

Safetensors 的设计理念主要围绕以下几个方面：

1. 安全性：通过避免序列化任意代码，确保在加载文件时不会执行恶意代码。这使得 Safetensors 成为处理不可信文件的理想选择。
2. 高效性：优化文件的存储和加载机制，支持快速内存映射（memory mapping），提高大模型文件的加载速度。
3. 简洁性：采用简单明了的文件结构，便于理解和实现，同时确保跨平台和跨框架的兼容性。
4. 可扩展性：设计上支持未来的扩展和改进，适应不断发展的机器学习需求。

Safetensors 的文件结构

Safetensors 文件采用二进制格式，具有以下基本结构：

1. Header（头部）：包含文件的元数据，如格式版本、数据长度等，确保文件的完整性和正确解析。
2. Tensor Metadata（张量元数据）：描述每个张量的名称、数据类型、形状等信息，便于加载和使用。
3. Tensor Data（张量数据）：实际存储的张量数值，采用紧凑的二进制格式，支持快速加载和内存映射。

以下是一个简化的文件结构示意图：

+---------------------+
|      Header         |
+---------------------+
| Tensor Metadata 1   |
+---------------------+
| Tensor Data 1       |
+---------------------+
| Tensor Metadata 2   |
+---------