veTurboIO

火山引擎研发的一款用于高性能读写 PyTorch 模型文件的 Python 库。该库实现了主要基于 safetensors 文件格式，实现高效的存储与读取张量数据。

安装

可以直接通过以下方式进行安装：

pip install veturboio -f https://veturbo-cn-beijing.tos-cn-beijing.volces.com/veturboio/index.html

Tips: 该指令会优先下载与当前 Python、PyTorch 版本匹配的 whl 文件，如果没有找到匹配的 whl 文件，会自动下载源码进行编译安装。当使用源码安装时，可增加 --no-build-isolation 来使用当前的运行环境进行编译并安装（否则会尝试创建虚拟环境）。

如果已经安装失败，可以尝试通过下载源码进行安装：

cd veturboio
python setup.py get_libcfs
python setup.py install

快速开始

import torch
import veturboio

tensors = {
   "weight1": torch.zeros((1024, 1024)),
   "weight2": torch.zeros((1024, 1024))
}

veturboio.save_file(tensors, "model.safetensors")

reloaded_tensor = veturboio.load("model.safetensors", map_location="cpu")

# check if the tensors are the same
for k, v in tensors.items():
    assert torch.allclose(v, reloaded_tensor[k])

使用锁页内存加速连续加载数据到GPU

import torch
import veturboio

tensors1 = {
   "weight1": torch.zeros((1024, 1024)),
   "weight2": torch.zeros((1024, 1024))
}

veturboio.save_file(tensors1, "model1.safetensors")

tensors2 = {
   "weight1": torch.zeros((1024, 1024)),
   "weight2": torch.zeros((1024, 1024))
}

veturboio.save_file(tensors2, "model2.safetensors")

helper = veturboio.init_io_helper()
reloaded_tensor1 = veturboio.load("model1.safetensors", map_location="cuda:0", use_pinmem=True, helper=helper)
# the map_location may be different
reloaded_tensor2 = veturboio.load("model2.safetensors", map_location="cuda:0", use_pinmem=True, helper=helper) 

# check if the tensors are the same
for k, v in tensors1.items():
    assert torch.allclose(v.cuda(), reloaded_tensor1[k])
for k, v in tensors2.items():
    assert torch.allclose(v.cuda(), reloaded_tensor2[k])

使用SFCS读写模型并启用加解密

该库可以读写 SFCS 上的模型文件，在此情况下可以启用加解密能力。读写 SFCS 上的模型文件和加解密所需的敏感信息，有两种获取方式：(1) 火山引擎可信服务的 unix domain socket，(2) 环境变量。在没有挂载可信服务 uds 的情况下，可以使用下面的环境变量：

环境变量名	含义
SFCS_ACCESS_KEY	SFCS文件系统的AK
SFCS_SECRET_KEY	SFCS文件系统的SK
SFCS_NAMENODE_ENDPOINT_ADDRESS	SFCS文件系统name节点地址
VETUROIO_KEY	加解密的128位数据密钥的base64编码
VETUROIO_IV	加解密的128位初始向量的base64编码

挂载可信服务 uds 或者配置好环境变量后，可以参考下面代码在读写 SFCS 上的模型文件时启用加解密：

import torch
import veturboio

tensors = {
   "weight1": torch.zeros((1024, 1024)),
   "weight2": torch.zeros((1024, 1024))
}

# use cpu to encrypt
veturboio.save_file(tensors, "sfcs://model.safetensors", use_cipher=True)

# use cpu to decrypt if map_location is cpu
reloaded_tensor1 = veturboio.load("sfcs://model.safetensors", map_location="cpu", use_cipher=True)

# use gpu to decrypt if map_location is cuda
reloaded_tensor2 = veturboio.load("sfcs://model.safetensors", map_location="cuda:0", use_cipher=True)

# check if the tensors are the same
for k, v in tensors.items():
    assert torch.allclose(v, reloaded_tensor1[k])
for k, v in tensors.items():
    assert torch.allclose(v, reloaded_tensor2[k])

转换现有的 PyTorch 文件

python -m veturboio.convert -i model.pt -o model.safetensors

性能测试

直接运行

bash bench/io_bench.sh

可以得到如下结果

fs_name    tensor_size     veturboio load_time(s)             torch load_time(s)            
shm        1073741824      0.08                               0.63                              
shm        2147483648      0.19                               1.26                              
shm        4294967296      0.36                               2.32

也可以进一步根据以下命令的参数说明调整使用参数

python bench/io_bench.py -h

特性

多线程高性能读取文件；
zero-copy 读取，不额外花费内存；
支持直接加载到 CUDA；
bfloat16 数值类型支持；
支持固定 pin-memory 用于让 GPU 快速反复读取大文件；
兼容 PyTorch 标准格式（无性能提升）；
兼容 safetensors 格式；
特殊加密格式存储；

收益

标准的 PyTorch 模型文件会经过 zip 与 pickle 两次操作，这两个操作极大的抑制了读取的速度，同时 unpickle 也会带来潜在的不安全性。我们使用一种自定义的模型格式来存储 tensor 数据，希望可以改善 PyTorch 标准格式所存在的这些问题。目前已经实现的优点有：

多线程读取：当前文件对象主要的存放点为云端存储，单一进程无法达到云存储的带宽上限，必须使用多线程读取才能达到最大的读取速度。PyTorch 标准格式的读取速度受限于 pickle 解析速度，远无法达到云存储的速度上限；
云端适配：基于火山引擎的云端存储（vePFS、SFCS）特性，最大化的利用了云端存储的带宽；
安全性：不再使用 pickle 对象，避免了 pickle 的安全性问题；

更新记录

前往 CHANGELOG 了解更多。

numb3r3/veTurboIO