Hugging Face Safetensors Security Audit and Default Adoption

Safetensors transition to default model format

Hugging Face, in collaboration with EleutherAI and Stability AI, is moving toward making the safetensors library the default format for saved models. This decision follows an external security audit performed by Trail of Bits, which confirmed that the library is safe and ready for widespread adoption.

The security vulnerability of PyTorch pickle

safetensors was developed to replace the use of pickle in PyTorch, which is inherently unsafe. Because pickle allows for the execution of arbitrary code, a malicious file posing as a model can grant an attacker full control of a user's computer without their knowledge.

While this vulnerability is well-documented in computer security circles and acknowledged in PyTorch documentation, it is not common knowledge within the broader machine learning community. This poses a significant risk for platforms like the Hugging Face Hub, where any user can upload and share models.

Technical capabilities of safetensors

Beyond security, safetensors provides several technical advantages over existing formats:

  • Framework Agnostic: The library supports saving and loading tensors across multiple frameworks, including PyTorch, TensorFlow, JAX, PaddlePaddle, and NumPy.
  • Performance: Loading is generally faster, with speeds up to 100x faster on CPU compared to other formats.
  • Lazy Loading: safetensors enables efficient lazy loading, allowing only specific parts of a tensor to be loaded. This is critical for efficient inference libraries, such as text-generation-inference, to load Large Language Models (LLMs) like LLaMA and StarCoder across various hardware configurations via arbitrary sharding.

Results of the Trail of Bits security audit

To validate the safety guarantees of the library, an external audit was conducted by Trail of Bits. The key findings include:

  • No Critical Flaws: No critical security flaws leading to arbitrary code execution were discovered.
  • Specification Fixes: The audit identified and fixed imprecisions in the specification format.
  • Validation Improvements: Issues allowing "polyglot files" were identified and resolved through improved validation.
  • Test Suite Enhancements: The audit resulted in numerous improvements to the library's test suite.

The library's implementation in Rust provides an additional layer of security inherent to the language itself.

Implementation roadmap

The transition to safetensors as the default format will be executed through the following steps:

Integration in Transformers

Within the transformers library, Hugging Face is implementing the following sequence:

  1. Creation and Verification: Ensuring the format delivers on promises regarding lazy loading, cross-framework compatibility, and load speeds.
  2. Security Validation: Confirming safety via the external audit.
  3. Core Dependency: Making safetensors a core dependency of the transformers library.
  4. Default Saving Format: Transitioning to safetensors as the default saving format within a few months to minimize disruption.

Ecosystem Adoption

EleutherAI has already added support for safetensors in their LM Evaluation Harness and is working on integrating it into the GPT-NeoX distributed training library. The format is already utilized by other ecosystem tools including Civitai, Stable Diffusion Web UI, dfdx, and LLaMA.cpp.

Sources