速度比較

Safetensors は非常に高速です。gpt2 の重みを読み込むことで、PyTorch と比較してみましょう。GPU ベンチマークを実行するには、お使いのマシンに GPU が搭載されているか、Google Colab を使用している場合は GPU runtime が選択されていることを確認してください。

開始する前に、必要なライブラリがすべてインストールされていることを確認してください。

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

この高速化は、このライブラリがファイルを直接マッピングすることで不要なコピーを回避するためです。実際には、pure pytorchでも可能です。現在示されている高速化は、以下の環境で得られました：

OS: 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に直接移動することで動作します。現在表示されている高速化は、以下の環境で得られたものです：

OS: Ubuntu 18.04.6 LTS
GPU: Tesla T4
Driver Version: 460.32.03
CUDA バージョン: 11.2

Update on GitHub