What Are Safetensors And Why Are They Better Than .ckpt (pickletensor)

Safetensors is a remarkable model format designed for securely storing tensors, boasting impressive speed and performance in comparison to many of the alternatives in the field. Not only is Safetensors incredibly fast, but it’s also been engineered with simplicity and safety at its core, setting it apart from many traditional formats used within machine learning.

Why Choose Safetensors?

With a plethora of formats available for machine learning and data storage, you might wonder what sets Safetensors apart. Well, let’s take a moment to compare some of these formats to understand how Safetensors delivers on key parameters:

• Safe: Will a file randomly downloaded not run arbitrary code?
• Zero-copy: Does reading the file require more memory than the original file?
• Lazy loading: Can the file be inspected without loading everything? And can only some tensors in it be loaded without scanning the whole file (important in a distributed setting)?
• Layout control: Can the location of tensors within the file be controlled to ensure fast access to individual tensors?
• No file size limit: Is there a limit to the file size?
• Flexibility: Can custom code be saved in the format and used later with zero extra code?
• Bfloat16: Does the format support native bfloat16 without the need for weird workarounds?

Each symbol ✔ represents the presence, ❌ the absence, ~ partial presence, and ? unknown presence of a feature.

Format	Safe	Zero-copy	Lazy loading	No file size limit	Layout control	Flexibility	Bfloat16
SafeTensors	✔	✔	✔	✔	✔	❌	✔
pickle (PyTorch)	❌	❌	❌	✔	❌	✔	✔
H5 (Tensorflow)	✔	❌	✔	✔	~	~	❌
SavedModel (Tensorflow)	✔	❌	❌	✔	✔	❌	✔
MsgPack (flax)	✔	✔	❌	✔	❌	❌	✔
Protobuf (ONNX)	✔	❌	❌	❌	❌	❌	✔
Cap’n’Proto	✔	✔	~	✔	✔	~	❌
Numpy (npy,npz)	✔	?	?	❌	✔	❌	❌
pdparams (Paddle)	❌	❌	❌	✔	❌	✔	✔

From this comparison, you can see that Safetensors meets or surpasses the features offered by other model formats. It’s particularly noteworthy for its capability for ‘zero-copy’ data manipulation and lazy loading, which can significantly boost the efficiency of data handling in machine learning tasks.

Some Noteworthy Comparisons

While all the formats have their unique strengths, let’s examine how Safetensors fares against some of the more commonly used ones:

• Pickle (PyTorch): A widely used format, but pickle comes with a significant downside – it’s unsafe as it allows for the execution of arbitrary code.
• H5 (Tensorflow): Though secure, it doesn’t support zero-copy, a feature integral to efficient data manipulation.
• SavedModel (Tensorflow): This format is specific to Tensorflow and lacks flexibility for other applications.
• Numpy (npz): While popular, this format doesn’t support bfloat16, an increasingly crucial feature in the ML world.

Beneath the Hood of Safetensors

Diving a little deeper into Safetensors, we find that it’s not just a data format, but a combination of thoughtful design choices that enhance safety, efficiency, and utility:

• Zero-copy: Safetensors design allows for true zero-copy on the CPU if the file is already in cache. On GPU, although a copy is always necessary due to the absence of a disk cache, the design bypasses the need to allocate all tensors on CPU at any given point.
• Endianness: Safetensors uses little-endian ordering, a popular choice in data storage.
• Order: It adopts ‘C’ or row-major ordering. This choice aligns with the prevalent standard in data storage.
• Stride: Striding is not supported. All tensors need to be packed before being serialized. This design decision was made based on current usage patterns in serialized formats.

The Extra Mile

Beyond simply being an efficient and safe data storage format, Safetensors also brings to the table some added advantages:

• Preventing DOS Attacks: The design of Safetensors mitigates the possibility of malicious files causing DOS attacks. There’s a limit on the size of the header to prevent parsing extremely large JSON, and there’s a guarantee that addresses in the file do not overlap, meaning loading a file should never exceed the file’s size in memory.
• Faster Load: Safetensors bypasses the extra CPU copy in PyTorch, enabling extremely fast CPU loading times compared to pickle. GPU loading times are as fast or faster than their PyTorch equivalent.
• Lazy Loading: In distributed settings (multi-node or multi-gpu), Safetensors enables loading only part of the tensors, significantly speeding up model loading times and improving development feedback loops.

In the Wild

Safetensors is already seeing widespread adoption at leading AI enterprises such as  Hugging Face, EleutherAI, and StabilityAI. Several prominent projects utilizing Safetensors include:

• huggingface/transformers
• AUTOMATIC1111/stable-diffusion-webui
• Llama-cpp
• microsoft/TaskMatrix
• hpcaitech/ColossalAI
• huggingface/pytorch-image-models
• CivitAI
• huggingface/diffusers
• coreylowman/dfdx
• invoke-ai/InvokeAI
• oobabooga/text-generation-webui
• Sanster/lama-cleaner
• PaddlePaddle/PaddleNLP
• AIGC-Audio/AudioGPT
• brycedrennan/imaginAIry
• comfyanonymous/ComfyUI
• LianjiaTech/BELLE

In conclusion, Safetensors, with its unique set of features and benefits, is all set to become a game-changer in the world of machine learning and data storage. So, go ahead, and give it a spin – the future of model formats awaits you!