There might be a slight discrepancy between the loading and saving process in safetensors. When loading a SD-based model like sd-1.4 packaged into a PyTorch checkpoint, we'll call it sd-v1-4.ckpt. We can package its state_dict and discard the torch format.

Packaging as safe_tensors

    sf_filename = "sd-v1-4.safetensors"
    filename = "sd-v1-4.ckpt"

    loaded = torch.load(filename)
    loaded = loaded['state_dict']

    # appears to pop nothing in this case
    shared = shared_pointers(loaded)
    for shared_weights in shared:
        for name in shared_weights[1:]:
            loaded.pop(name)

    loaded = {k: v.contiguous() for k, v in loaded.items()}

    save_file(loaded, metadata={"format": "pt"})

    check_file_size(local, filename)

Loading the tensors

load_file('sd-v1-4.safetensors', device='cpu')

Results in error:

File "venv\lib\site-packages\safetensors\torch.py", line 99, in load_file result[k] = f.get_tensor(k) RuntimeError: shape '[1280, 1280, 3, 3]' is invalid for input of size 7290352

Expected behaviour: safetensors fails while trying to save unexpected tensor data or creates tensors which can be loaded Affected version: safetensors=2.4.0, torch=1.12.1+cu113

ckpt size: 3.97 GB 4,265,381,888 bytes (4,265,380,512 bytes) safetensor size: 3.97 GB 4,265,148,416 bytes (4,265,146,304 bytes) SHA fe4efff1e174c627256e44ec2991ba279b3816e364b49f9be2abc0b3ff3f8556

Using pruned version of CompVis/stable-diffusion-v-1-4-original

Apologies if this is already fixed with addition of more dtypes. Will try to get more info by running through check output and debug info of this specific tensor

Used

os.environ['SAFETENSORS_FAST_GPU'] = '1'

I observed on the webui that despite setting device='cuda' with this flag and using safetensors.torch.load_file, it was taking almost 45 seconds to load a 4GB .safetensors file. However, when trying to replicate it in a separate program but using the same libraries, the model loads fast and only takes a few seconds. These can be represented by these curves, each a 60s time chunk:

It appears that due to some pollution of memory, webui always falls back to loading from CPU by default. This pollution appears to persist, and if you terminate the webui and then run a separate program which uses safetensors, it also falls back to loading using slow CPU copy.

Steps to replicate:

1. Load a program using loading various large torch files to cpu (regular torch files with torch.load), and loading a large file to GPU (safetensors).
2. Terminate it after loading the torch CPU files and GPU for a few seconds.
3. Within 10 seconds, try to launch an external program using safetensors to GPU with the fast gpu flag. The program resorts to slow copy via CPU. *Does not replicate if you interrupt the second external program and try again.

safetensors tested: 0.2.5

Next steps:

• print how much memory cuda sees as available at runtime

• disable windows 10 Hardware-accelerated GPU scheduling

• figure out how the webui pollutes space (I think it loads a few models with torch to CPU, 5GB to memory)

Much better API:

with PySafeFile(self.local, framework="pt") as f:
      tensor = f["test"][:, :, 1:2]

Design:

• Context manager This is meant to open a file, we need to close it, and this is the Python way (lazy load on a real buffer instead of a file doesn't have real advantages since torch is already able to not copy a single thing with striding in the first place)
• Framework arg: I played with other ideas, as if the code was framework agnostic it would be easier on users on any kind of failure and to adapt directly within Python. However, memoryview doesn't handle dtype, and bfloat16 is only handled by pytorch, so using it as intermediary is not really as viable as I thought.
• f["test"] This object is ~equivalent to TensorView<'data> except we cannot have a lifetime on it, so it owns the Mmap instead. This is not as torch.Tensor at this point, only when you slice into does it become torch.Tensor.
• [:, :, 1:2] This is where we can calculate how to slice into the data, creating the buffer iteratively (so only once) as PyByteArray so we do actually avoid allocating the entire buffer here (hopefully). In order to avoid shape shenanigans, we simply return the torch.Tensor directly.

Pros:

• handles a lot of boilerplate for caller.
• Feels more pythonesque

Cons:

• Unsure about soundness of keeping Mmap struct (in PySafeFile), it's meant to be used as context manager, so it should clean up after itself.
• Lots more boilerplate in the binding code.

Loading roughly 2x as fast (depends on hardware) models on GPU.

IT works by not allocating on CPU, and directly allocating + memsetting the memory on the GPU. Afaik, there's no way to do that in Pytorch. All the storage methods have intermediary CPU alloc.

@davidhewitt Sorry to ping you out of the blue, but you have been a huge help on tokenizers.

The present library aims to prevent arbitrary code execution when loading weights (in particular Pytorch). https://www.youtube.com/watch?v=2ethDz9KnLk

At HF, we haven't yet fully committed to actually switch format, but multiple other nice features could come out of it, namely this handy 2x speedup when loading tensors on GPU (because we can remove entirely the CPU allocations). However the present solution uses a good dose of unsafe.

• Do you have ideas on removing this unsafe ?
• Or validating it ?

Currently there's is no way (afaik) to access an equivalent of cudaMemcpy directly from torch. That indirection would help put the safety back in pytorch and not in this crate. However after a healthy dose of looking I couldn't find anything. Still since pytorch is in Python world, I'm guessing it's always going to require passing a PyBufferArray which we have to reallocate (because of trailing '\0' at least).

The current PR does the following:

• It figures out the libcudart being used by pytorch itself
• Loads it
• When creating tensors on GPU, it will allocate an empty buffer (managed by Pytorch) through torch.empty(shape, dtype, device).
• Then lookup cudaMemcpyHostToDevice.
• Call it directly to set the GPU RAM with the actual tensors on disk.

Extremely fast (I'm guessing it's just pointer arithmetic and actual load time would it when actually using the tensors.)

benches/test_pt.py::test_pt_pt_load_cpu PASSED
benches/test_pt.py::test_pt_sf_load_cpu PASSED


------------------------------------------------------------------------------------- benchmark: 2 tests ------------------------------------------------------------------------------------
Name (time in ms)            Min                 Max                Mean            StdDev              Median                IQR            Outliers       OPS            Rounds  Iterations
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_pt_sf_load_cpu       4.4699 (1.0)        9.8395 (1.0)        5.5089 (1.0)      0.9217 (1.0)        5.2233 (1.0)       0.7899 (1.0)         24;13  181.5238 (1.0)         175           1
test_pt_pt_load_cpu     162.7417 (36.41)    183.1321 (18.61)    173.9028 (31.57)    8.4244 (9.14)     174.7930 (33.46)    13.4904 (17.08)         2;0    5.7503 (0.03)          6           1
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

Legend:
  Outliers: 1 Standard Deviation from Mean; 1.5 IQR (InterQuartile Range) from 1st Quartile and 3rd Quartile.
  OPS: Operations Per Second, computed as 1 / Mean
=============================================================================================================== 2 passed, 4 deselected in 8.96s ================================================================================================================

I checked against transformers and I get:

============== DUMMY ====================
Loaded in 0:00:03.180369
Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
[{'generated_text': 'test.\n\nThe first thing to note is that the "new" version of the app is not'}]
Ran in 0:00:04.339221
============== NO ALLOC ====================
Lookup in 0:00:00.000292
model from config in 0:00:00.581647
tokenizer in 0:00:00.641293
Loaded in 0:00:00.648781
Loaded weights in 0:00:00.654212
Loaded on model in 0:00:00.658909
Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
Ran in 0:00:01.828281
[{'generated_text': 'test.\n\nThe first thing to note is that the "new" version of the app is not'}]
Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
Ran in 0:00:02.982728
[{'generated_text': 'test.\n\nThe first thing to note is that the "new" version of the app is not'}]

And without safetensors

============== DUMMY (PT) ====================
Loaded in 0:00:04.272222
Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
[{'generated_text': 'test.\n\nThe first thing to note is that the "new" version of the app is not'}]
Ran in 0:00:05.432200

https://gist.github.com/Narsil/57ddec904ee795cc6e78eb13562f0c36

I am not sure where the 500-600ms come from when loading the model with no_init_weights from accelerate, I'm guessing it's torch internals but I'm not sure.

the NO_ALLOC case is more of a showcase but maybe at some point we can manage to integrate into transformers.

TLDR: add benchmarking for serialization & deserialziation

Not sure if we will do parallelization https://github.com/huggingface/safetensors/issues/15 But either way, I thought it is probably a good idea to add some benchmarking.

I'm using criterion, which is a popular rust benchmarking framework (also used on tokenizers).

Can be run as:

cargo bench

Produces output as:

Serlialize 10_MB
time:   [427.15 µs 436.21 µs 446.56 µs]                             
change: [-0.3956% +0.7567% +1.9834%] (p = 0.32 > 0.05)
No change in performance detected.

Deserlialize 10_MB                     
time:   [5.5303 µs 5.5634 µs 5.5997 µs]                                
change: [-0.4513% +0.5147% +1.5265%] (p = 0.32 > 0.05)
No change in performance detected.

Read more about cargo bench here. Note: since criterion is a dev-dependency, it does not get added in the prod build

Very cool project!

I was wondering if there's possibility to use this package as the "backend" for torch.save to enable easier integration with downstream projects that are already using torch.save/load...

seems like it might be possible since torch.save and torch.load have a pickle_module= argument. Something like:

import safetensors
torch.save(my_tensor, "my_tensor.pt", pickle_module=safetensors)   # or safetensors.torch

import numpy as np
import safetensors.numpy

a = safetensors.numpy.save({"a": np.array(range(6), dtype='>u4')})
b = safetensors.numpy.save({"a":np.array(range(6), dtype='<u4')})
print(a)
print(b)

b'7\x00\x00\x00\x00\x00\x00\x00{"a":{"dtype":"U32","shape":[6],"data_offsets":[0,24]}}\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x02\x00\x00\x00\x03\x00\x00\x00\x04\x00\x00\x00\x05'
b'7\x00\x00\x00\x00\x00\x00\x00{"a":{"dtype":"U32","shape":[6],"data_offsets":[0,24]}}\x00\x00\x00\x00\x01\x00\x00\x00\x02\x00\x00\x00\x03\x00\x00\x00\x04\x00\x00\x00\x05\x00\x00\x00'

Getting rid of pyO3 and using ctypes instead will allow to

• get rid of the requirement of having and executing Rust compiler in order to build a wheel;
• reduce bloat by reusing the shared library

For doc-builds & in general, pip install .[dev] is preferred over python setup.py develop

1. Previously, pip install .[dev] was erring due to lines below missing: https://github.com/huggingface/safetensors/blob/109f54c1f390e2906098f3e5144eff4272aac384/bindings/python/pyproject.toml#L1-L3
2. For the extras management, I've copied it from transformers

Getting error below on Intel CPU Macbook:

cargo check inside safetensors/bindings/python

(base) Mishigs-MacBook-Pro:python mishig$ cargo check
    Checking safetensors-python v0.2.3 (/Users/mishig/Desktop/safetensors/bindings/python)
error[E0412]: cannot find type `c_uint` in module `std::ffi`
  --> src/lib.rs:33:16
   |
33 | ) -> std::ffi::c_uint;
   |                ^^^^^^ not found in `std::ffi`
   |
help: consider importing this type alias
   |
3  | use std::os::raw::c_uint;

However, the error goes away when std::ffi::c_uint is replaced with std::os::raw::c_uint

I don't know what's the difference between the two: their docs look identical

Since it looks like both of them are data store library. But flatbuffer not only can store Tensor, but also structures && many customized data types. Does safetensors also able to?

Furthermore, does safetensor able to using cross platform with a very tiny dependency lib? such as use it on Android or iOS

Ideally we want to support the following API:

with safe_open(filename, mode="w", metadata=metadata) as f:
    f.get_tensor("embedding")[::2] = my_tensor 

https://www.reddit.com/r/StableDiffusion/comments/z8mnak/comment/iyzztaf/?utm_source=share&utm_medium=web2x&context=3

Currently SAFETENSORS_FAST_GPU=1 expected cuda cards. In case of hybridization, the code will likely find the cuda part, yet if we're allocating on the AMD card we need to not use the fast path.

We need to check :

• that this is indeed a failure mode
• Fix it

Currently an option to save model with pytorch is to use (https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model)

*** Save/Load Entire Model

#Save:
torch.save(model, PATH)

Load:
# Model class must be defined somewhere
model = torch.load(PATH)
model.eval()

Would it be possible to use safetensors to do the same?