Speed Comparison
Safetensors is really fast. Let's compare it against PyTorch by loading gpt2 weights. To run the GPU benchmark, make sure your machine has GPU or you have selected GPU runtime if you are using Google Colab.
Before you begin, make sure you have all the necessary libraries installed:
!pip install safetensors huggingface_hub torch
Let's start by importing all the packages that will be used:
Download safetensors & torch weights for gpt2:
CPU benchmark
Loaded safetensors 0:00:00.004015 ,Loaded pytorch 0:00:00.307460 ,on CPU, safetensors is faster than pytorch by: 76.6 X
This speedup is due to the fact that this library avoids unnecessary copies by mapping the file directly. It is actually possible to do on pure pytorch. The currently shown speedup was gotten on:
- OS: Ubuntu 18.04.6 LTS
- CPU: Intel(R) Xeon(R) CPU @ 2.00GHz
GPU benchmark
Loaded safetensors 0:00:00.165206 ,Loaded pytorch 0:00:00.353889 ,on GPU, safetensors is faster than pytorch by: 2.1 X
The speedup works because this library is able to skip unecessary CPU allocations. It is unfortunately not replicable in pure pytorch as far as we know. The library works by memory mapping the file, creating the tensor empty with pytorch and calling cudaMemcpy directly to move the tensor directly on the GPU.
The currently shown speedup was gotten on:
- OS: Ubuntu 18.04.6 LTS.
- GPU: Tesla T4
- Driver Version: 460.32.03
- CUDA Version: 11.2