Nouveau standard pour le stockage des modèles AI : concepts et avantages

En parcourant des hubs de modèles AI tels que Hugging Face, on commence à voir de plus en plus de fichiers avec l'extension .safetensors, au lieu des fichiers habituels comme .bin ou .pth.

Dans cet article, nous examinerons en détail ce qu'est .safetensors, pourquoi il a été introduit et quels avantages puissants il offre par rapport aux méthodes existantes, du point de vue technique.

1. Qu'est-ce que .safetensors ?

.safetensors est un nouveau format de stockage de tenseurs développé par Hugging Face.

Les modèles de deep learning se composent de milliards de paramètres (poids), et ce format a été créé pour stocker et charger ces énormes quantités de données. Il a été conçu pour résoudre les problèmes critiques (sécurité et vitesse) du format de stockage basé sur le module pickle de Python, qui était jusqu'ici de facto standard.

En résumé, c'est un "fichier de stockage de modèle plus sûr et plus rapide".

2. Pourquoi a-t-il été introduit : les problèmes de l'ancienne méthode (Pickle)

Les modèles PyTorch existants (.bin, .pth) utilisent le module pickle de Python pour sérialiser les données. Cependant, pickle présente un problème critique.

Vulnérabilité de sécurité (Exécution de code arbitraire)

pickle ne se contente pas de stocker des données, il stocke également des objets Python eux-mêmes. Ce processus peut inclure du code Python, permettant à un hacker malveillant d’insérer un code capable de détruire des systèmes ou de voler des informations personnelles à l'intérieur d'un fichier modèle. Dès que l'utilisateur charge le modèle avec load(), ce code malveillant s'exécute.

Exemple de situation problématique :

L'utilisateur charge model.bin téléchargé sur Internet -> Le code caché à l'intérieur du fichier s'exécute -> La clé SSH ou le mot de passe de l'utilisateur est envoyé au serveur de l'attaquant.

.safetensors a été créé pour contrer ces menaces de sécurité.

3. Caractéristiques clés de .safetensors

3.1. Sécurité

.safetensors est, comme son nom l'indique, sécurisé. Ce format ne stocke que des données de tenseurs purs et des métadonnées (sous forme JSON). Il n'y a absolument aucune possibilité d'inclure du code exécutable, ce qui permet de charger en toute sécurité des fichiers même s'ils proviennent de sources non fiables.

3.2. Zero-Copy et rapidité

Le chargement de LLM (modèles de langage large) ou de modèles Stable Diffusion est considérablement plus rapide.

• Méthode existante : le fichier est copié en mémoire CPU -> désérialisation -> reconversion en forme de tenseur -> déplacement vers le GPU. (Processus de copie inutile)

• safetensors : utilise la technologie de mappage mémoire (Memory Mapping, mmap). Le système d'exploitation mappe directement le fichier à une adresse mémoire, ce qui permet d'accéder immédiatement aux données du disque sans processus de copie inutile. C'est ce qu'on appelle le Zero-Copy.

3.3. Chargement paresseux (Lazy Loading)

Il est possible de lire rapidement uniquement les parties nécessaires du modèle sans charger l'ensemble du modèle en mémoire.

Par exemple, si vous souhaitez vérifier le poids d'une couche spécifique dans un fichier modèle de 100 Go, avec l'ancienne méthode, vous deviez lire les 100 Go, mais avec .safetensors, vous pouvez lire uniquement cette partie. C'est particulièrement avantageux dans un environnement d'apprentissage distribué ou pour l'optimisation des inférences.

3.4. Compatibilité avec les frameworks

Il n'est pas dépendant d'un framework de deep learning spécifique (PyTorch, etc.).

• PyTorch

• TensorFlow

• JAX

• PaddlePaddle

Il a été conçu pour être facilement lisible et écrivible dans divers frameworks comme ceux-ci.

4. Structure des fichiers

Les fichiers .safetensors ont une structure très simple.

1. En-tête (Header) : situé au début du fichier, il suit le format JSON. Il contient des informations sur le nom de chaque tenseur, le type de données (dtype), la forme (shape), et la position des données (offset).

2. Données (Data) : un bloc de données binaires qui suit l'en-tête. Les valeurs de tenseurs purs y sont compactement remplies.

Cette structure permet de comprendre la structure du modèle simplement en lisant l'en-tête sans avoir à lire tout le fichier.

4-1. Conseils pratiques : vérifier l'en-tête et les métadonnées dans le terminal

Charger des modèles de plusieurs gigaoctets pour vérifier s'ils sont quantifiés ou quelles couches elles contiennent n'est pas efficace.

La bibliothèque safetensors propose une fonctionnalité qui permet de scanner rapidement les informations d'en-tête sans lire le fichier complet. Il suffit d'entrer la commande Python One-liner suivante dans le terminal pour vérifier immédiatement.

Bien sûr, vous pouvez également vérifier en cliquant sur file info en haut à droite de la page où vous avez téléchargé le modèle sur Hugging Face. Mais ici, je vais expliquer comment le vérifier directement dans mon terminal.

Préparation

Vous devez d'abord installer la bibliothèque.

pip install safetensors

Commande (entrée dans le terminal)

Entrez le chemin réel du fichier à la place de model.safetensors.

python -c "from safetensors import safe_open; \
with safe_open('model.safetensors', framework='pt', device='cpu') as f: \
    print('--- Metadata ---'); \
    print(f.metadata()); \
    print('\n--- Tensor Keys (Layers) ---'); \
    print(list(f.keys())[:5])" # trop long donc affiche seulement les 5 premiers

Interprétation des résultats

En exécutant cette commande, vous obtiendrez deux informations importantes.

1. Métadonnées : informations insérées par le créateur du modèle. S'il est indiqué ici format: gptq ou quantization: int4, même si cela n'est pas écrit dans le nom du fichier, vous saurez qu'il s'agit d'un modèle quantifié. (Notez que si le créateur a laissé les métadonnées vides, None peut apparaître.)

2. Clés : noms des couches qui composent le modèle. Cela peut vous aider à comprendre la structure du modèle.

5. Résumé des comparaisons : .bin (Pickle) vs .safetensors

6. Exemples d'utilisation (Python)

Voici un exemple simple de stockage et de chargement de tenseurs utilisant la bibliothèque safetensors.

import torch
from safetensors.torch import save_file, load_file

# 1. Création et stockage d'un tenseur
tensors = {
    "embedding": torch.zeros((1024, 512)),
    "attention": torch.rand((512, 512))
}

# Stockage du tenseur sous forme de dictionnaire dans un fichier
save_file(tensors, "model.safetensors")

# 2. Chargement du fichier
loaded = load_file("model.safetensors")
print(loaded["embedding"].shape) 
# Affichage : torch.Size([1024, 512])

7. Conclusion

Le format .safetensors n'est pas qu'un simple changement d’extension de fichier, c'est une évolution essentielle pour la s sécurité et l' efficacité des modèles AI. Déjà établi comme format de base par des communautés clés, y compris Hugging Face.

À l'avenir, il est recommandé d'utiliser .safetensors au lieu de .bin au moment de télécharger ou de distribuer un modèle. Cela assurera une protection contre les menaces de sécurité et réduira considérablement le temps de chargement des modèles.