速度比較

Safetensors 非常快。讓我們透過載入 gpt2 權重來將它與 PyTorch 進行比較。若要執行 GPU 基準測試，請確保您的機器擁有 GPU，或者如果您正在使用 Google Colab，請確保已選擇 GPU 執行階段。

開始之前，請確保您已安裝所有必要的程式庫。

pip install safetensors huggingface_hub torch

讓我們從匯入所有將要使用的套件開始。

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

下載 gpt2 的 safetensors 與 torch 權重

>>> sf_filename = hf_hub_download("gpt2", filename="model.safetensors")
>>> pt_filename = hf_hub_download("gpt2", filename="pytorch_model.bin")

CPU 基準測試

>>> start_st = datetime.datetime.now()
>>> weights = load_file(sf_filename, device="cpu")
>>> load_time_st = datetime.datetime.now() - start_st
>>> print(f"Loaded safetensors {load_time_st}")

>>> start_pt = datetime.datetime.now()
>>> weights = torch.load(pt_filename, map_location="cpu")
>>> load_time_pt = datetime.datetime.now() - start_pt
>>> print(f"Loaded pytorch {load_time_pt}")

>>> print(f"on CPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X")
Loaded safetensors 0:00:00.004015
Loaded pytorch 0:00:00.307460
on CPU, safetensors is faster than pytorch by: 76.6 X

此加速是因為該程式庫透過直接映射檔案來避免不必要的複製。實際上，使用純 pytorch 也是可以做到的。目前顯示的加速數據取得環境為：

作業系統：Ubuntu 18.04.6 LTS
CPU：Intel(R) Xeon(R) CPU @ 2.00GHz

GPU 基準測試

>>> # This is required because this feature hasn't been fully verified yet, but 
>>> # it's been tested on many different environments
>>> os.environ["SAFETENSORS_FAST_GPU"] = "1"

>>> # CUDA startup out of the measurement
>>> torch.zeros((2, 2)).cuda()

>>> start_st = datetime.datetime.now()
>>> weights = load_file(sf_filename, device="cuda:0")
>>> load_time_st = datetime.datetime.now() - start_st
>>> print(f"Loaded safetensors {load_time_st}")

>>> start_pt = datetime.datetime.now()
>>> weights = torch.load(pt_filename, map_location="cuda:0")
>>> load_time_pt = datetime.datetime.now() - start_pt
>>> print(f"Loaded pytorch {load_time_pt}")

>>> print(f"on GPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X")
Loaded safetensors 0:00:00.165206
Loaded pytorch 0:00:00.353889
on GPU, safetensors is faster than pytorch by: 2.1 X

這種加速之所以有效，是因為該程式庫能夠跳過不必要的 CPU 配置。遺憾的是，據我們所知，這在純 pytorch 中是無法複製的。該程式庫的工作原理是將檔案進行記憶體映射，使用 pytorch 建立空的張量，並直接呼叫 cudaMemcpy 將張量移動到 GPU 上。目前顯示的加速數據取得環境為：

作業系統：Ubuntu 18.04.6 LTS。
GPU：Tesla T4
驅動程式版本：460.32.03
CUDA 版本：11.2

在 GitHub 上更新