Safetensors项目速度优势深度解析：与PyTorch的性能对比

Safetensors项目速度优势深度解析：与PyTorch的性能对比

safetensors Simple, safe way to store and distribute tensors 项目地址: https://gitcode.com/gh_mirrors/sa/safetensors

前言

在深度学习模型部署和推理过程中，模型权重的加载速度直接影响着整体效率。本文将深入分析Safetensors这一高效张量存储格式在加载速度方面的显著优势，并通过详实的基准测试数据展示其与PyTorch原生格式的性能差异。

环境准备

在开始性能对比前，我们需要准备测试环境：

pip install safetensors huggingface_hub torch

测试方法

我们选择GPT-2模型权重作为测试对象，分别对比Safetensors格式和PyTorch原生格式在CPU和GPU上的加载速度。

数据准备

import os
import datetime
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import torch

# 下载两种格式的模型权重
sf_filename = hf_hub_download("gpt2", filename="model.safetensors")
pt_filename = hf_hub_download("gpt2", filename="pytorch_model.bin")

CPU性能测试

在CPU环境下，我们观察到Safetensors展现出惊人的加载速度优势：

# Safetensors加载测试
start_st = datetime.datetime.now()
weights = load_file(sf_filename, device="cpu")
load_time_st = datetime.datetime.now() - start_st

# PyTorch加载测试
start_pt = datetime.datetime.now()
weights = torch.load(pt_filename, map_location="cpu")
load_time_pt = datetime.datetime.now() - start_pt

# 结果对比
print(f"CPU加载速度对比：Safetensors比PyTorch快 {load_time_pt/load_time_st:.1f} 倍")

典型测试结果：

• Safetensors加载时间：0.004秒
• PyTorch加载时间：0.307秒
• 速度提升：76.6倍

技术原理

Safetensors之所以能在CPU上实现如此显著的性能提升，主要归功于以下技术特点：

1. 内存映射技术：直接映射文件到内存，避免了不必要的内存拷贝
2. 零拷贝优化：在加载过程中最小化数据移动操作
3. 精简格式设计：专为高效加载而优化的二进制格式

GPU性能测试

在GPU环境下，Safetensors同样展现出明显优势：

# 启用GPU加速优化
os.environ["SAFETENSORS_FAST_GPU"] = "1"

# 预热GPU
torch.zeros((2, 2)).cuda()

# 执行测试（代码与CPU测试类似，设备改为"cuda:0"）

典型测试结果：

• Safetensors加载时间：0.165秒
• PyTorch加载时间：0.354秒
• 速度提升：2.1倍

GPU优化原理

Safetensors在GPU上的性能优势来自以下创新设计：

1. 直接GPU拷贝：通过cudaMemcpy直接将张量数据复制到GPU
2. CPU分配跳过：避免了PyTorch默认流程中的中间CPU分配步骤
3. 异步传输优化：充分利用GPU的并行处理能力

性能影响因素

实际测试中，性能提升幅度会受到以下因素影响：

1. 硬件配置：

   • CPU型号和架构
   • GPU型号和CUDA版本
   • 存储设备速度（SSD/NVMe等）

2. 软件环境：

   • 操作系统版本
   • 驱动程序版本
   • CUDA/cuDNN版本

3. 模型特性：

   • 权重文件大小
   • 张量数量和维度

最佳实践建议

基于测试结果，我们推荐：

1. 生产环境优先使用Safetensors：特别是对延迟敏感的应用场景
2. 大型模型部署：模型越大，使用Safetensors带来的收益越明显
3. 频繁加载场景：如模型服务需要频繁重新加载权重的情况

结论

Safetensors通过创新的文件格式设计和高效的加载实现，在模型权重加载速度上显著优于PyTorch原生格式。无论是在CPU还是GPU环境下，都能带来可观的性能提升，这对于提高深度学习工作流的整体效率具有重要意义。

对于追求极致性能的开发者来说，将模型权重转换为Safetensors格式是一个简单而有效的优化手段。

safetensors Simple, safe way to store and distribute tensors 项目地址: https://gitcode.com/gh_mirrors/sa/safetensors