Comparación de velocidad

Safetensors es realmente rápido. Comparemos su rendimiento con PyTorch cargando los pesos de gpt2. Para ejecutar la prueba de rendimiento en GPU, asegúrate de que tu máquina tenga una GPU o de que hayas seleccionado GPU runtime si estás usando Google Colab.

Antes de comenzar, asegúrate de tener instaladas todas las bibliotecas necesarias:

pip install safetensors huggingface_hub torch

Comencemos importando todos los paquetes que se utilizarán:

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

Descarga los pesos de safetensors & torch para gpt2:

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

Benchmark de 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

Esta aceleración se debe a que esta biblioteca evita copias innecesarias al mapear el archivo directamente. De hecho, es posible hacerlo con pure pytorch. La aceleración mostrada actualmente se obtuvo en:

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

Benchmark de 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

La aceleración funciona porque esta biblioteca puede omitir asignaciones de CPU innecesarias. Por desgracia, no es replicable en PyTorch puro, según nuestro conocimiento. La biblioteca funciona mapeando el archivo en memoria, creando el tensor vacío con PyTorch y llamando directamente a cudaMemcpy para mover el tensor directamente a la GPU. La aceleración mostrada actualmente se obtuvo en:

SO: Ubuntu 18.04.6 LTS.
GPU: Tesla T4
Versión del controlador: 460.32.03
Versión de CUDA: 11.2

Update on GitHub