safetensors 속도 비교 — pickle보다 얼마나 빠를까
safetensors 속도 비교 — pickle보다 얼마나 빠를까
safetensors는 정말 빠른 게 특징이에요. gpt2 가중치를 PyTorch(pickle)와 비교해서 로드 시간을 직접 맞붙여 보는 벤치마크가 문서에 있어요.
준비
pip install safetensors huggingface_hub torch
import os, datetime
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import 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
start_pt = datetime.datetime.now()
weights = torch.load(pt_filename, map_location="cpu")
load_time_pt = datetime.datetime.now() - start_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
이 속도 향상은 라이브러리가 파일을 직접 매핑하면서 불필요한 복사를 피하기 때문에 나와요.
GPU 벤치마크
os.environ["SAFETENSORS_FAST_GPU"] = "1"
torch.zeros((2, 2)).cuda() # CUDA 시작 비용은 측정 제외
# ... 측정 ...
# on GPU, safetensors is faster than pytorch by: 2.1 X
GPU에서는 파일을 메모리 매핑하고 텐서를 만들고 cudaMemcpy로 GPU에 직접 옮겨 CPU 할당을 건너뛰어요.