So sánh tốc độ

Safetensors thực sự rất nhanh. Hãy so sánh nó với PyTorch bằng cách nạp trọng số của gpt2. Để chạy GPU benchmark, hãy đảm bảo máy của bạn có GPU hoặc bạn đã chọn GPU runtime nếu bạn đang sử dụng Google Colab.

Trước khi bắt đầu, hãy đảm bảo bạn đã cài đặt tất cả các thư viện cần thiết:

pip install safetensors huggingface_hub torch

Hãy bắt đầu bằng cách nhập tất cả các gói sẽ được sử dụng:

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

Tải trọng số safetensors & torch cho gpt2:

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

Đo hiệu năng 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

Tốc độ tăng này là do thư viện này tránh các bản sao không cần thiết bằng cách ánh xạ trực tiếp tệp. Thực tế, điều này có thể thực hiện được trên pure pytorch. Tốc độ tăng được hiển thị hiện tại được đo trên:

Hệ điều hành: Ubuntu 18.04.6 LTS
CPU: Intel(R) Xeon(R) CPU @ 2.00GHz

Đo hiệu năng 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

Tốc độ tăng lên hoạt động vì thư viện này có thể bỏ qua các phân bổ CPU không cần thiết. Theo như chúng tôi biết, điều này không thể tái hiện được trong pytorch thuần. Thư viện hoạt động bằng cách ánh xạ bộ nhớ tệp, tạo tensor rỗng bằng pytorch và gọi trực tiếp cudaMemcpy để chuyển tensor trực tiếp lên GPU. Tốc độ tăng lên hiện được hiển thị đã đạt được trên:

Hệ điều hành: Ubuntu 18.04.6 LTS.
GPU: Tesla T4
Phiên bản Driver: 460.32.03
Phiên bản CUDA: 11.2

Update on GitHub