随着huggingface transformers的流行，越来越多的模型采用了.safetensors的文件存储格式。它有什么特点？与传统的torch.save保存的文件有何不同？本文将详细梳理相关细节。
.safetensors这是huggingface设计的一种新格式，大致就是以更加紧凑、跨框架的方式存储Dict[str, Tensor]，主要存储的内容为tensor的名字（字符串）及内容（权重）。
其官网对文件格式的详细内容进行了解释，本质上就是一个JSON文件加上若干binary形式的buffer。对于tensor而言，它只存储了数据类型、形状、对应的数据区域起点和终点。因此它只支持存储dense and contiguous tensor，也就是稠密张量（没有stride，没有view等信息）。本质上它就像存储了一大块连续的内存，甚至可以直接把文件映射到内存里（使用Python的mmap模块）。
这种方式其实有一个问题，就是没有存储数据的大小端信息。目前大部分机器都是小端（little-endian）的，如果你在用大端机器（big-endian），建议就不要跑LLM了…… 很多代码都会容易出问题的。
理论上这种数据格式是可以支持多个tensor共享存储的，只需要把多个tensor的offsets设置成一样就行了。但是huggingface为了保持对多个框架的兼容性（同时支持pytorch、TensorFlow、jax），没有实现这一特点。如果真的有两个tensor共享存储，它只存第一个tensor。
使用一份简单的Python代码就可以将.safetensors文件转成内存中的Dict[str, Tensor]类型的对象：
from safetensors import safe_open
tensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:
for k in f.keys():
tensors[k] = f.get_tensor(k)这样得到的结果基本上就可以当pytorch的state_dict用了。
.safetensors支持五种框架，包括pytorch、TensorFlow、flax（jax）、paddle（paddlepaddle）、numpy。对每个框架都提供了save/save_file/load/load_file这四个函数：
load( data: bytes, device = 'cpu' ) → Dict[str, Tensor]
load_file( filename: Union[str, os.PathLike], device = 'cpu' ) → Dict[str, Tensor]
save( tensors: Dict[str, Tensor], metadata: Optional = None ) → bytes
save_file( tensors: Dict, filename: Union[str, os.PathLike], metadata: Optional = None ) → None
对于numpy，函数没有device参数；Tensor类型根据具体的框架替换成框架的具体类，例如torch.Tensor或者np.array。
为了支持pytorch里的weight sharing（例如常见的word embedding与lm_head共享权重），safetensors专门为pytorch设计了两个函数：
load_model(model: Module, filename: Union[str, os.PathLike], strict = True ) → (missing, unexpected)
save_model( model: Module, filename: str, metadata: Optional = None, force_contiguous: bool = True )
以一个简单的共享权重的模型为例：
from torch import nn
class Model(nn.Module):
def __init__(self):
super().__init__()
self.a = nn.Linear(100, 100)
self.b = self.a
def forward(self, x):
return self.b(self.a(x))
model = Model()
from safetensors.torch import save_model, safe_open, load_model
save_model(model, "model.safetensors")
tensors = {}
with safe_open("model.safetensors", framework="pt") as f:
for k in f.keys():
tensors[k] = f.get_tensor(k)
print(list(tensors.keys())) # ['a.bias', 'a.weight']
model_restored = Model()
load_model(model_restored, "model.safetensors")其实model.safetensors文件里完全丢失了b的信息，只保留了a的数据。这就要求我们加载模型的时候，由model_restored来提供b与a共享数据的信息。
我们大致也能猜到load_model的实现方式：
def load_model(model, filename):
data = load_file(filename)
model_data = model.state_dict()
for k, v in data.items():
model_data[k].copy_(v).safetensors支持只加载元数据然后再加载tensor的一部分，例如：
from safetensors import safe_open
tensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:
tensor_slice = f.get_slice("embedding")
vocab_size, hidden_dim = tensor_slice.get_shape()
tensor = tensor_slice[:, :hidden_dim]执行f.get_slice之后，tensor_slice并没有立马读取全部的数据，只有在被索引的时候才读出必要的那部分数据。这一特性在分布式训练中很有用，一张显卡上面可能只加载权重的一部分。
由于.safetensors的元数据与实际数据的存储是分开的，因此有可能不读取全部数据就能够看到权重的元信息。例如，当我们访问 https://huggingface.co/openai-community/gpt2 时，点击权重旁边的展开按钮就可以看到所有权重的名字、数据类型以及形状。
这一功能还是挺实用的，不用下载模型就可以看权重的一些维度参数。
torch.save的存储格式torch.save是我们日常使用的函数。然而，恐怕很少有人知道它的原理。
实际上，torch.save保存的文件就是一个zip文件。是的，没错，就是个压缩包。
import torch
x = torch.tensor([0, 1, 2, 3, 4])
torch.save(x, 'tensor.pt')我们以上述代码为例，对存储的文件直接unzip tensor.pt，就可以看到以下内容：
tensor
├── byteorder
├── data
│ └── 0
├── data.pkl
└── versiontorch.save可以看做通用的Python object存储方式，实际上它不仅可以存储tensor，还可以存储函数对象，甚至可以存储torchscript编译后的内容……而且在加载时它能正确保留内存共享、内存切片等信息。
torch.save的功能太复杂，这里无法详细列出。大致来说，它是把每一个Python object使用pickle进行保存，然后打包成一个zip压缩文件，同时对pytorch内部的tensor存储、内存共享等细节进行了部分优化。
根据以上内容，我们可以看出，torch.save保存的内容比.safetensors多很多。所以，从torch.save转到.safetensors是比较简单的，官方也给出了转换脚本。
总的来说，.safetensors是针对huggingface transformers类模型的权重存储需求而设计的一种简单的跨框架文件格式，抛弃了torch.save支持的非常多大模型用不到的特性，所以才能够变得简单、高效。
