Comparaison des vitesses

Safetensors est vraiment rapide. Comparons-le à PyTorch en chargeant les poids de gpt2. Pour exécuter le benchmark GPU, assurez-vous que votre machine dispose d’un GPU ou que vous avez sélectionné GPU runtime si vous utilisez Google Colab.

Avant de commencer, assurez-vous d'avoir installé toutes les bibliothèques nécessaires :

pip install safetensors huggingface_hub torch

Commençons par importer tous les paquets qui seront utilisés :

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

Télécharger les poids safetensors & torch pour gpt2 :

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

Benchmark 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

Cette accélération est due au fait que cette bibliothèque évite les copies inutiles en mappant directement le fichier. Il est en réalité possible de le faire en pytorch pur. L'accélération actuellement affichée a été obtenue sur :

OS : Ubuntu 18.04.6 LTS
CPU : Intel(R) Xeon(R) CPU @ 2.00GHz

Benchmark 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

L'accélération fonctionne car cette bibliothèque est capable d'éviter les allocations CPU inutiles. À notre connaissance, il est malheureusement impossible de reproduire ce résultat en pytorch pur. La bibliothèque fonctionne en mappant le fichier en mémoire, en créant le tensor vide avec pytorch et en appelant cudaMemcpy directement pour transférer le tensor directement sur le GPU. L'accélération actuellement présentée a été obtenue sur :

OS : Ubuntu 18.04.6 LTS.
GPU : Tesla T4
Version du pilote : 460.32.03
Version de CUDA : 11.2

Update on GitHub