一文详解模型权重存储新格式 Safetensors

在日常AI模型训练过程中，需要好的模型权重通常需要以一种格式存储在磁盘中。比如：目前最流行的AI框架 PyTorch 使用 pickle 格式存储模型权重文件。但 PyTorch 文档中有一段话说明如下：使用 torch.load() 保存模型，除非 weights_only 参数设置为True （只加载张量、原始类型和字典），否则隐式使用 pickle 模块，这是不安全的。可以构造恶意的 pickle 数据，该数据将在 unpickling 期间执行任意代码。切勿在不安全模式下加载可能来自不受信任来源的数据或可能已被篡改的数据，仅加载您信任的数据。为了规避类似的问题，新的权重存储格式 Safetensors 应运而生。

Safetensors 简介

Safetensors 是一种用于安全地存储张量的新格式，非常简单，但速度仍然很快（零拷贝）。它是 pickle 格式的替代品，因为，pickle 格式不安全，可能包含可以执行的恶意代码。

Safetensors 内部格式

假设现在有一个名为 model.safetensors 的Safetensors文件，那么 model.safetensors 内部格式如下：

• 前面的 8 bytes是一个无符号的整数，表示 header 占的字节数。
• 中间的 N bytes是一个UTF-8编码JSON字符串，存储 header 的内容，里面为模型权重的元数据信息。
• 文件的剩余部分存储模型权重 tensor 的值。

以GPT2的Safetensors文件为例，其元数据信息可以通过如下代码获取：

import requests # pip install requests
import struct

def parse_single_file(url):
    # Fetch the first 8 bytes of the file
    headers = {'Range': 'bytes=0-7'}
    response = requests.get(url, headers=headers)
    # Interpret the bytes as a little-endian unsigned 64-bit integer
    length_of_header = struct.unpack('<Q', response.content)[0]
    # Fetch length_of_header bytes starting from the 9th byte
    headers = {'Range': f'bytes=8-{7 + length_of_header}'}
    response = requests.get(url, headers=headers)
    # Interpret the response as a JSON object
    header = response.json()
    return header

url = "https://huggingface.co/gpt2/resolve/main/model.safetensors"
header = parse_single_file(url)

print(header)
# {
#   "__metadata__": { "format": "pt" },
#   "h.10.ln_1.weight": {
#     "dtype": "F32",
#     "shape": [768],
#     "data_offsets": [223154176, 223157248]
#   },
#   ...
# }

不同模型权重格式大比拼

对于不同模型权重格式，从如下几个方面进行全面的对比：

• 安全性：是否可以使用随机下载的文件并期望不运行任意代码吗？
• 零拷贝：读取文件是否需要比原始文件更多的内存？
• 延迟加载：可以在不加载所有内容的情况下检查文件吗？并仅加载其中的一些张量而不扫描整个文件吗？
• 布局控制：延迟加载不一定足够，因为如果有关张量的信息分散在您的文件中，那么即使可以延迟访问该信息，您也可能必须访问大部分文件才能读取可用的张量（导致许多磁盘到 RAM 的拷贝）。因此，控制布局以快速访问单个张量非常重要。
• 无文件大小限制：文件大小有限制吗？
• 灵活性：是否可以以该格式保存自定义代码并能够在以后以零额外代码使用它吗？（意味着我们可以存储不仅仅是纯张量，还有代码）
• 支持Bfloat16：该格式是否支持原生 bfloat16（意味着不需要奇怪的解决方法）？因为 Bfloat16 数据格式在机器学习领域变得越来越重要。

下图展示了常见的存储格式的特性。

safetensors 和 ONNX 的不同

safetensors 和 ONNX 具有不同的用途。

safetensors 是一种简单、安全、快速的文件格式，用于存储和加载张量。它是 Python 的 pickle 实用程序的替代品，更加安全；而后者不安全，可能包含可以执行的恶意代码。

ONNX（开放神经网络交换）是一种用于表示深度学习模型的开放格式。它允许您用不同深度学习框架（例如：PyTorch、TensorFlow、Caffe2 等）加载模型以一种方式保存模型。这使得在不同框架之间共享模型变得更容易。

综上所述， safetensors 用于安全快速地存储和加载张量，而 ONNX 用于在不同深度学习框架之间共享模型。这同样适用于其他模型共享框架。

补充：零拷贝技术

零拷贝的主要任务就是避免CPU将数据从一块存储中拷贝到另一块存储，主要就是利用各种技术，避免让CPU做大量的数据拷贝任务，以此减少不必要的拷贝。或者借助其他的一些组件来完成简单的数据传输任务，让CPU解脱出来专注别的任务，使得系统资源的利用更加有效。

下面通过从磁盘读取数据并将其发送到套接字(socket)的示例来了解零拷贝（这种情况在大多数 Web 应用程序中经常发生）。为了完成此操作，内核会将数据读取到用户空间。

操作系统术语：

用户空间和内核空间是由操作系统分隔的两个虚拟内存区域，以提供内存保护和硬件保护，防止恶意或错误的软件行为。

用户空间是应用程序软件和一些驱动程序执行的内存区域。每个用户空间进程通常运行在自己的虚拟内存空间中，除非明确允许，否则不能访问其他进程的内存或内核空间。用户空间进程只能通过系统调用与内核交互，系统调用是内核向用户空间公开的一组函数。

内核空间是操作系统内核、内核扩展和大多数设备驱动程序运行的内存区域。内核空间程序运行在内核模式下，也称为管理模式，这是一种特权模式，允许访问所有CPU指令和硬件资源。

一旦数据加载到用户空间，它将再次执行内核调用，内核会将数据写入套接字(socket)。每次数据穿越user-kernel boundary时，都必须进行复制，这会消耗 CPU 周期和内存带宽。

而零拷贝请求即内核直接将数据从磁盘文件复制到套接字，而不经过应用程序。

下图为传统的拷贝操作流程。除了拷贝操作之外，这些系统调用会导致用户空间和内核空间之间发生大量上下文切换，使得这个过程变慢。

下图为零拷贝数据传输流程：

零拷贝的特点是 CPU 不全程负责内存中的数据写入其他组件，CPU 仅仅起到管理的作用。但注意，零拷贝不是不进行拷贝，而是 CPU 不再全程负责数据拷贝时的搬运工作。如果数据本身不在内存中，那么必须先通过某种方式拷贝到内存中（这个过程 CPU 可以不参与），因为数据只有在内存中，才能被转移，才能被 CPU 直接读取计算。

Safetensors 优势

上面对比了不同模型权重存储格式，从中可以发现 Safetensors 主要优势有：

• 安全：使用 torch.load() 加载模型权重可能会执行被插入的恶意代码（因为 pickle 模块是不安全的），不过可以设置weights_only=False 避免这个问题。而 Safetensors 天然就没有这个问题。

• 速度快：Safetensors 接口的速度比原生的 Pytorch 接口加载权重更快。

• CPU 上加载提速的原因：通过直接映射文件，避免了不必要的复制（零拷贝）
• GPU 上加载提速的原因：跳过不必要的CPU分配。

• 惰性加载：可以在不加载整个文件的情况下查看文件的信息或者只加载文件中的部分张量而不是所有张量。当我们有一个包含许多键和值对的大文件时，延迟加载非常重要。如果我们可以单独加载单个键的值，它将提高内存效率并且速度更快，否则我们将不得不将完整文件加载到内存中以检查任何键。

Safetensors 实践

安装

pip install safetensors

保存与加载张量

保存张量：

import torch
from safetensors.torch import save_file

tensors = {
    "embedding": torch.zeros((2, 2)),
    "attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")

加载张量：

from safetensors import safe_open

tensors = {}
# with safe_open("model.safetensors", framework="pt", device=0) as f:
with safe_open("model.safetensors", framework="pt", device='cpu') as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k) # loads the full tensor given a key
print(tensors)
# {'attention': tensor([[0., 0., 0.],
#          [0., 0., 0.]], device='cuda:0'),
#  'embedding': tensor([[0., 0.],
#          [0., 0.]], device='cuda:0')}

仅加载部分张量（延迟加载）：

tensors = {}
# with safe_open("model.safetensors", framework="pt", device=0) as f:
with safe_open("model.safetensors", framework="pt", device='cpu') as f:

    tensor_slice = f.get_slice("embedding")
    vocab_size, hidden_dim = tensor_slice.get_shape()
    print(vocab_size, hidden_dim)
    tensor = tensor_slice[:, :hidden_dim] # change the hidden_dim to load part of the tensor
    print(tensor)
# 2 2
# tensor([[0., 0.],
#         [0., 0.]])

将PyTorch模型保存为Safetensors格式

from torchvision.models import resnet18
from safetensors.torch import load_model, save_model
import torch

model_pt = resnet18(pretrained=True)

# 保存 state dict 为 safetensors格式
save_model(model_pt, "resnet18.safetensors")
# Instead of save_file(model.state_dict(), "model.safetensors")


# 加载没有权重的模型
model_st = resnet18(pretrained=False) 
load_model(model_st, "resnet18.safetensors")
# Instead of model.load_state_dict(load_file("model.safetensors"))


# 使用随机图像对初始模型和从safetensors中加载权重的新模型进行推理
img = torch.randn(2, 3, 224, 224)

model_pt.eval()
model_st.eval()


with torch.no_grad():
    print(torch.all(model_pt(img)==model_st(img))) # tensor(True)
    

结语

本文简要介绍了模型权重存储新格式 Safetensors，它具备安全、加载速度快等多个优点；并且可以在 HuggingFace 上面看到越来越多的模型使用Safetensors格式进行存储。