Chargement des modèles

Transformers fournit de nombreux modèles préentraînés prêts à l’emploi en une seule ligne de code. Il est nécessaire d’utiliser une classe de modèle et la méthode from_pretrained().

Appel de from_pretrained() pour télécharger et charger les poids et la configuration d’un modèle stockés sur le Hub Hugging Face.

La méthode from_pretrained() charge les poids stockés au format de fichier safetensors s’ils sont disponibles. Traditionnellement, les poids des modèles PyTorch sont sérialisés à l’aide de l’utilitaire pickle, réputé pour son manque de sécurité. Les fichiers Safetensor sont plus sûrs et se chargent plus rapidement.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="auto")

Ce guide explique comment les modèles sont chargés, les différentes façons de charger un modèle, comment surmonter les problèmes de mémoire pour les modèles très volumineux, et comment charger des modèles personnalisés.

Modèles et configurations

Tous les modèles disposent d’un fichier configuration.py contenant des attributs spécifiques tels que le nombre de couches cachées, la taille du vocabulaire, la fonction d’activation, etc. Vous y trouverez également un fichier modeling.py qui définit les couches et les opérations mathématiques effectuées à l’intérieur de chaque couche. Le fichier modeling.py prend les attributs du modèle définis dans configuration.py et construit le modèle en conséquence. À ce stade, vous disposez d’un modèle aux poids aléatoires qui doit être entraîné pour produire des résultats significatifs.

Une architecture désigne le squelette du modèle, tandis qu’un checkpoint désigne les poids du modèle pour une architecture donnée. Par exemple, BERT est une architecture, alors que google-bert/bert-base-uncased est un checkpoint. Vous verrez le terme modèle utilisé de manière interchangeable pour désigner l’architecture et le checkpoint.

Il existe deux types généraux de modèles que vous pouvez charger :

Un modèle minimal, comme AutoModel ou LlamaModel, qui sort des états cachés.
Un modèle auquel est rattachée une tête spécifique, comme AutoModelForCausalLM ou LlamaForCausalLM, afin d’effectuer des tâches précises.

Classes de modèles

Pour obtenir un modèle pré-entraîné, vous devez charger les poids dans le modèle. Cela se fait en appelant from_pretrained(), qui accepte les poids depuis le Hugging Face Hub ou un répertoire local.

Il existe deux classes de modèle, la classe AutoModel et une classe spécifique au modèle.

AutoModel

classe spécifique au modèle

Grands modèles

Les grands modèles pré-entraînés nécessitent beaucoup de mémoire pour être chargés. Le processus de chargement implique :

créer un modèle avec des poids aléatoires
le chargement des poids pré-entraînés
placer les poids pré-entraînés sur le modèle

Il faut suffisamment de mémoire pour stocker deux copies des poids du modèle (aléatoires et pré-entraînées), ce qui peut s’avérer impossible selon votre matériel. Dans les environnements d’entraînement distribué, la tâche est encore plus ardue, car chaque processus charge un modèle pré-entraîné.

Transformers atténue certains de ces problèmes liés à la mémoire grâce à une initialisation rapide, des points de contrôle fractionnés, la fonctionnalité Big Model Inference d’Accelerate et la prise en charge de types de données à faible nombre de bits.

Points de contrôle fractionnés

save_pretrained() fragmente automatiquement les points de contrôle dont la taille dépasse 50 Go. Cela permet de limiter le nombre de fragments pour les grands modèles et de simplifier la gestion des fichiers.

Les paramètres sont chargés en parallèle et la mémoire de pointe ne dépend que de la taille du modèle. Utilisez max_shard_size dans save_pretrained() pour définir la taille maximale du point de contrôle avant le fractionnement.

La consommation mémoire pour les modèles nécessitant une conversion dynamique des poids dépend de la taille du modèle et de la taille des plus grands paramètres lors d'une seule conversion. Cela s'applique généralement aux modèles de type mixture-of-experts (MoE), où la consommation mémoire correspond à la taille du modèle plus le nombre d'experts sur une couche. Consultez le guide du chargeur de poids dynamique pour en savoir plus sur la manière dont les modèles sont chargés.

save_pretrained() crée également un fichier d’index qui associe les noms des paramètres à leurs fichiers de shards. L’index contient deux clés, metadata et weight_map.

import json
import os
import tempfile

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="50GB")
    with open(os.path.join(tmp_dir, "model.safetensors.index.json"), "r") as f:
        index = json.load(f)

print(index.keys())

metadata stocke la taille totale du modèle.

index["metadata"]
{'total_size': 28966928384}

weight_map associe chaque paramètre à son fichier de fragment.

index["weight_map"]
{'lm_head.weight': 'model-00006-of-00006.safetensors',
 'model.embed_tokens.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.input_layernorm.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.mlp.down_proj.weight': 'model-00001-of-00006.safetensors',
 ...
}

Inférence de grands modèles

Assurez-vous d'avoir installé Accelerate v0.9.0 et PyTorch v1.9.0 ou une version ultérieure pour utiliser cette fonctionnalité !

from_pretrained() est renforcé par la fonctionnalité Big Model Inference d’Accelerate.

Big Model Inference crée un squelette de modèle sur le périphérique meta de PyTorch. Le périphérique meta ne stocke aucune donnée réelle, uniquement les métadonnées.

Les poids initialisés aléatoirement ne sont créés que lorsque les poids pré-entraînés sont chargés, afin d'éviter de maintenir deux copies du modèle en mémoire simultanément. L'utilisation mémoire maximale est donc limitée à la taille du modèle.

En savoir plus sur le placement des appareils dans Concevoir une carte d'appareils.

La deuxième fonctionnalité de Big Model Inference concerne la manière dont les poids sont chargés et distribués dans le squelette du modèle. Les poids du modèle sont répartis sur tous les appareils disponibles, en commençant par l’appareil le plus rapide (généralement le GPU), puis en déplaçant les poids restants vers des appareils plus lents (CPU et disque dur).

La combinaison de ces deux fonctionnalités réduit la consommation de mémoire et les temps de chargement pour les grands modèles pré-entraînés.

Définissez device_map sur "auto" pour activer Big Model Inference.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto")

Vous pouvez également affecter manuellement des couches à un appareil dans device_map. Il doit associer tous les paramètres du modèle à un appareil, mais vous n’avez pas besoin de préciser l’emplacement de tous les sous-modules d’une couche si l’ensemble de la couche se trouve sur le même appareil.

Accédez à l'attribut hf_device_map pour voir comment un modèle est réparti sur les différents appareils.

device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}
model.hf_device_map

Déchargement sur disque

Lorsqu’un modèle est trop volumineux pour la mémoire combinée de vos GPU et de votre RAM CPU, Big Model Inference décharge les poids restants sur le disque. Définissez offload_folder sur un répertoire chaque fois qu’un poids est mappé sur "disk", car c’est là que Transformers stocke les poids déchargés.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    offload_folder="./offload",
)

Utilisez max_memory pour limiter la quantité de mémoire que chaque appareil peut utiliser. Transformers remplit les appareils dans l'ordre de leur vitesse et transfère sur le disque ce qui ne tient pas. L'exemple ci-dessous limite la GPU à 10 Go et le CPU à 30 Go, de sorte que les poids restants sont stockés dans offload_folder.

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={0: "10GB", "cpu": "30GB"},
    offload_folder="./offload",
)

Définissez max_memory={} pour indiquer qu’aucune mémoire GPU ou CPU n’est disponible et pour forcer le déchargement complet du modèle sur le disque.

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={},
    offload_folder="./offload",
)

Le déchargement sur disque échange la mémoire contre la vitesse. Lire les poids depuis le disque à chaque passe avant est plus lent que de les lire depuis la mémoire CPU ou GPU. Le déchargement entièrement sur disque nécessite accelerate>1.14.0

Type de données du modèle

L’argument dtype contrôle le dtype PyTorch utilisé pour instancier les poids du modèle. Par défaut, Transformers charge les poids avec la valeur dtype ou torch_dtype (ancienne version) de config.json. Si config.json ne contient aucune de ces valeurs, Transformers utilise le dtype du premier poids à virgule flottante du point de contrôle.

Écraser la valeur par défaut en transmettant un dtype spécifique.

import torch
from transformers import AutoModelForCausalLM

# specific dtype
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-it", dtype=torch.float16)

AutoConfig accepte également dtype pour les modèles instanciés à partir de zéro.

import torch
from transformers import AutoConfig, AutoModel

my_config = AutoConfig.from_pretrained("google/gemma-2b", dtype=torch.float16)
model = AutoModel.from_config(my_config)

Modèles personnalisés

Les modèles personnalisés s’appuient sur les classes de configuration et de modélisation de Transformers, prennent en charge l’API AutoClass et sont chargés avec from_pretrained(). La différence est que le code de modélisation ne provient pas de Transformers.

Prenez des précautions supplémentaires lors du chargement d'un modèle personnalisé. Bien que le Hub inclue un balayage des logiciels malveillants pour chaque dépôt, vous devez rester vigilant afin d'éviter d'exécuter accidentellement du code malveillant.

Définissez trust_remote_code=True dans from_pretrained() pour charger un modèle personnalisé.

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("sgugger/custom-resnet50d", trust_remote_code=True)

En guise de mesure de sécurité supplémentaire, chargez un modèle personnalisé à partir d'une révision spécifique afin d'éviter de charger du code de modèle qui a pu être modifié. L'identifiant de commit peut être copié depuis l'historique des commits des modèles.

commit_hash = "ed94a7c6247d8aedce4647f00f20de6875b5b292"
model = AutoModelForImageClassification.from_pretrained(
    "sgugger/custom-resnet50d", trust_remote_code=True, revision=commit_hash
)

Reportez-vous au guide Personnaliser les modèles pour plus d'informations.

Update on GitHub