Pourquoi un modèle peut-il entrer sur votre machine sans qu’il soit possible de l’entraîner?

La règle habituelle vous dit si un modèle va entrer : paramètres multipliés par les octets par poids, comme l’explique cet article sur le calcul de la mémoire. Ce chiffre répond à la question « puis-je l’exécuter? ». L’entraînement est une tout autre facture. Bien des gens s’en rendent compte quand une exécution qui semblait sûre plante à sa première étape.

Ce projet mesure cette facture avec un petit modèle, SmolLM2-135M-Instruct, qui compte 134 515 008 paramètres et une licence Apache 2.0. Chaque ligne ci-dessous a été exécutée dans un nouveau processus sur le processeur, en nombres de 32 bits sauf indication contraire, et a lu le pic de mémoire donné par le système d’exploitation. C’est donc un vrai pic et non une estimation. Les pics incluent un peu de surcoût de chargement, et les 244 Mo qu’occupaient Python et les bibliothèques avant tout modèle ont déjà été soustraits. Le script est 07_memory_budget.py.

Où va la mémoire

L’exécution exige la première ligne et une plus petite part de la dernière. L’entraînement exige les quatre. Avant même de traiter un seul exemple, l’entraînement coûte donc déjà quatre fois les poids. Ces trois lignes fixes totalisent 2 152 Mo.

Ce qui a été mesuré

La longueur des séquences est de 256 jetons, sauf indication contraire. Mesuré ici sur un Apple M5 Max avec 64 Go, PyTorch 2.15 et Transformers 5.22.

Voici comment le lire :

• Entrer ne dit pas grand-chose. Avec un lot de 8, l’exécution exige 1 269 Mo et l’entraînement 6 484 Mo, soit 5,1 fois plus. Par rapport aux seuls 538 Mo des poids, l’entraînement en prend 12 fois autant.
• Avec un lot de 1, vous payez les coûts fixes. Un seul exemple coûte quand même 3 093 Mo. Réduire la taille du lot a un plancher, qui se situe à environ quatre fois les poids.
• La longueur et le lot pèsent autant. Un lot de 4 à longueur 512 (2 048 jetons) exige à peu près autant qu’un lot de 8 à longueur 256 (2 048 jetons) : 6 489 et 6 484 Mo. Les longs documents sont une décision de mémoire, et les tronquer ou les découper économise autant que réduire le lot.
• Les correctifs agissent sur des lignes différentes. Les points de contrôle de gradient recalculent les activations au lieu de les garder, ce qui a économisé 37 % ici, au prix de calculs supplémentaires. Les nombres de 16 bits réduisent tout de moitié en principe; mesurés, ils ont économisé 34 %, puisque le pic comprend aussi du surcoût. LoRA gèle les poids et n’entraîne que de petites matrices ajoutées, si bien que les gradients et l’état de l’optimiseur disparaissent presque : 460 800 paramètres entraînables, 0,34 % du modèle.

Ces chiffres viennent d’un petit modèle sur un processeur. Les cartes graphiques gardent d’autres mémoires tampons et n’ont pas été testées ici : voyez-y un plancher et gardez une marge.

Le mesurer vous-même

Voici le cœur d’une mesure. Elle doit s’exécuter dans un nouveau processus, parce que le pic ne redescend jamais.

import resource, sys, torch
from transformers import AutoModelForCausalLM

def peak_mb():
    v = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return v / (1024 ** 2 if sys.platform == "darwin" else 1024)

model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct", dtype=torch.float32)
model.gradient_checkpointing_enable()        # retirez cette ligne pour comparer
model.config.use_cache = False
model.train()
opt = torch.optim.AdamW(model.parameters(), 1e-4)
x = torch.randint(0, 1000, (8, 256))
for _ in range(2):                           # l'état d'Adam existe à partir de la deuxième étape
    model(input_ids=x, labels=x).loss.backward()
    opt.step()
    opt.zero_grad(set_to_none=False)
print("pic en Mo :", round(peak_mb()))

La deuxième boucle n’est pas un détail. Adam alloue son état à la première mise à jour : un test d’une seule étape donne donc un chiffre trop bas et vous envoie vers un plantage à l’étape deux.

Comment estimer votre budget avant de lancer une exécution?

Additionnez les lignes, puis gardez une marge. Voici une formule rapide, où P est le nombre de paramètres en milliards et O, le nombre d’octets par poids :

Calculé pour un modèle de 7 milliards de paramètres à 2 octets par poids, avant toute activation :

Une machine dotée de 24 Go de mémoire graphique peut donc facilement contenir ce modèle pour l’exécuter, ne peut pas l’entraîner en entier et peut souvent y entraîner un adaptateur. Ce dernier cas peut exiger le chargement en précision réduite expliqué dans l’article sur la quantification, car 14 Go laissent peu de place aux activations. Aucune de ces trois lignes n’a été exécutée ici; c’est de l’arithmétique fondée sur le comportement du petit modèle, et aucun GPU de cette taille n’a été testé.

Pour connaître la part des activations, faites comme ce projet : exécutez deux vraies étapes d’entraînement avec votre plus long exemple et la taille de lot prévue, lisez le pic, puis doublez le lot et regardez comment il augmente. Deux minutes de ce test en disent plus que n’importe quel tableau. Vous saurez aussi quel correctif essayer d’abord : si les coûts fixes dominent, utilisez un adaptateur ou une précision réduite; si les activations dominent, raccourcissez les séquences, réduisez le lot ou activez les points de contrôle. Le projet 13 couvre ce qu’il faut regarder avant de décider que la mémoire est bien le vrai problème.