Shouldn't safetensors allow the storage of e.g. a flax.frozen_dict.FrozenDict?

[alvarobartt]

AFAIU the main idea behind safetensors is to define a more general and safe/efficient format to store tensors. But one question comes to my mind, shouldn't it be useful not to just store plain tensors but also more complex structures containing tensors such as flax.frozen_dict.FrozenDict?

Just asking because it's unclear to me in which direction is the library going or whether a more general serialization format is the end goal or not.

Thanks! 🤗

[Narsil]

Hey thanks for your question.

Currently there's no plan to support anything more than tensors.

The initiation for this lib was simply sparked to try and avoid using pickle. As the README suggests other libraries where considered but all had issues that prevented them from being used (H5 was the biggest contendant, but since TF was deprecating its use and it doesn't support bfloat16 which is kind of important for models like bloom we didn´t go that way).

The goal of this lib is to be trivial (so you can actually load safetensors files without even using this lib : https://gist.github.com/Narsil/3edeec2669a5e94e4707aa0f901d2282

And since we're creating something new (https://xkcd.com/927/) we might as well add some features which were missing from everything else:

• Slice support (you can load only part of a tensor and doing so efficiently, this is important for very LLMs like bloom where you need to load only parts of the models on each GPU).
• Slice support for writing (this is for training LLMs and is in the works)
• Faster loading for both GPU and CPU. Since we can have a very flat layout we can optimize quite a bit the loading of the tensors which wouldn't be the case with more complex data structures.
• In addition to arbitrary code execution (ACE) prevention, I decided to go for DOS (denial of service) prevention, so this library should not DOS you even with maliciously crafted files. (The file will take at most its size in memory)

Does that answer your question ?

[alvarobartt]

Thanks for this https://gist.github.com/Narsil/d5b0d747e5c8c299eb6d82709e480e3d! I had something similar that I was using for simple structures, didn't test it extensively anyway, I'll probably build something simple around that in case any user is looking for similar answers. I was also having a look at what you did at https://huggingface.co/spaces/safetensors/convert/blob/main/convert.py, in case I develop something I'll just tag it here so that anyone can use it. Great work and thanks again! 🤗

[Narsil]

Thanks a lot !

convert.py stemmed from a need to convert pytorch checkpoints, but it's starting to become quite big, I think having 2 files convert_pt.py and convert_flax.py might be better for readability and usability (dependencies).

[alvarobartt]

I'm already working on something, I think that this week I'll be able to release the first dev version of the package that I created to use safetensors to both save and load flax.core.frozen_dict.FrozenDict. I'll probably post it here in the replies if it's useful for anyone! Do you also want me to collaborate in that HF Space with a convert_flax.py file or something? Thanks for your work, really enjoying safetensors usage and efficiency so far!

[Narsil]

Sure let us know when you're done for the script and/or the space !

[alvarobartt]

Thanks for your help @Narsil, project is already available at https://github.com/alvarobartt/safejax, install it as pip install safejax 👍🏻