Multiverse Computing propose une distillation hors ligne et une perte KL découpée pour alléger fortement la mémoire nécessaire à l’entraînement des LLM.
La distillation de connaissances est une réponse directe à un problème très concret : les grands modèles de langage ouverts deviennent difficiles à déployer, mais aussi coûteux à entraîner ou à adapter. Le principe est connu : un modèle student, plus petit, apprend à reproduire le comportement d’un modèle teacher, plus volumineux. Dans la pratique, l’étape de récupération des capacités du modèle enseignant peut elle-même exiger une infrastructure GPU considérable.
Dans un article technique publié sur Hugging Face, l’équipe de Multiverse Computing décrit deux optimisations système destinées à faire baisser ce coût : mettre en cache une version réduite des sorties du teacher, puis calculer la perte de distillation sans matérialiser de gigantesques tenseurs intermédiaires. Les résultats annoncés sont prometteurs, notamment pour les contextes longs. Mais ils doivent être lus dans leur cadre expérimental précis : un matériel, des modèles et des métriques donnés.
La proposition ne change pas l’objectif de la distillation : elle évite surtout de conserver ou de reconstruire en mémoire des distributions de probabilités complètes sur tout le vocabulaire. Selon les mesures rapportées, cette approche réduit fortement le pic de VRAM tout en conservant une perte d’entraînement proche de celle de la distillation en ligne.
Pourquoi la distillation de LLM consomme autant de mémoire
Dans la distillation dite online, le teacher et le student sont chargés simultanément. À chaque étape d’entraînement, le premier produit une distribution de probabilités pour le prochain token ; le second est optimisé pour s’en rapprocher, généralement avec une divergence de Kullback-Leibler (KL).
Cette méthode est expressive puisqu’elle exploite toute la distribution de sortie du teacher. Son coût est toutefois élevé pour deux raisons. D’abord, le teacher doit exécuter une passe avant à chaque itération, alors que ses prédictions ne changent pas au cours d’un même entraînement. Ensuite, le calcul standard de la perte manipule des représentations denses dont la taille dépend à la fois du vocabulaire et de la longueur de séquence.
L’exemple donné par les auteurs illustre l’ordre de grandeur. Pour gpt-oss-120b, dont le vocabulaire compte 201 088 tokens, une séquence de 32K tokens avec une taille de batch de 4 conduit à un tenseur de probabilités teacher de forme 4 × 201088 × 32768. En bfloat16, ce seul tenseur représente environ 50 Go de VRAM, selon la source. En ajoutant les activations, gradients, poids et états de l’optimiseur, le pic d’une itération de distillation peut atteindre environ 250 Go dans leur exemple.
Ce n’est donc pas uniquement la taille des modèles qui bloque. C’est aussi la manière dont les sorties sont calculées, stockées et comparées pendant l’entraînement.
Mettre le teacher hors de la boucle grâce aux logits Top-K
La première optimisation consiste à passer à une distillation offline. Au lieu de recalculer les sorties du teacher à chaque pas, l’équipe les calcule une fois et conserve, pour chaque position, les 100 tokens les plus probables avec leurs logits associés. Le teacher peut ensuite être retiré de la mémoire durant l’entraînement du student.
Ce cache Top-K répond à deux besoins opérationnels :
- il évite d’allouer le modèle teacher pendant la phase d’optimisation du student ;
- il permet de réutiliser les mêmes sorties pour plusieurs essais, ablations ou réglages d’hyperparamètres.
Le compromis est clair : le student n’observe plus l’intégralité de la distribution du teacher, mais seulement ses 100 sorties les plus probables. Dans le benchmark présenté à 8K tokens, les courbes de perte rapportées par les auteurs se superposent presque entre distillation en ligne et distillation hors ligne avec Top-100. Cela soutient leur affirmation d’une récupération sans perte sur cette mesure et dans ce protocole. En revanche, l’extrait fourni ne donne pas de résultats détaillés sur des évaluations fonctionnelles en aval — raisonnement, génération, benchmarks métier ou robustesse — : on ne peut donc pas généraliser cette équivalence à toutes les capacités d’un modèle.
Trois calculs de perte KL, une même cible mathématique
Le cache ne suffit pas à résoudre tout le problème. Même sans teacher chargé, une implémentation naïve de la divergence KL peut reconstruire des grilles denses très coûteuses. Les auteurs comparent trois variantes de calcul hors ligne, qu’ils présentent comme mathématiquement équivalentes.
Dense KL : le point de comparaison
La méthode dense reconstruit une distribution teacher complète à partir du cache Top-100, puis la compare aux log-probabilités denses du student. Elle sert de référence de correction, mais conserve les gros tenseurs liés au vocabulaire complet et aux positions de séquence.
Forward-chunked KL : conserver le teacher sparse
La deuxième méthode traite la séquence par tranches et garde les sorties teacher sous une forme sparse, limitée aux logits mis en cache. Elle évite donc de densifier la distribution teacher. Toutefois, les logits complets du student restent produits et conservés pour la rétropropagation : la consommation mémoire continue alors de croître fortement avec la longueur de contexte.
Fused chunked KL : ne jamais produire tous les logits student
La troisième variante, la contribution principale annoncée, fusionne la projection de sortie du modèle avec le calcul de la perte. Plutôt que de générer la matrice complète des logits student, elle traite une tranche de positions : projection des états cachés, contribution à la perte, libération de la tranche, puis passage à la suivante.
Lors de la rétropropagation, chaque tranche est recalculée au lieu d’avoir été conservée. Cette stratégie échange donc une partie du temps de calcul contre une baisse du pic mémoire. Les auteurs précisent que la projection est effectuée deux fois — à l’aller puis au retour —, ce qui explique que cette variante ne soit pas systématiquement la plus rapide sur les séquences plus courtes.
Les chiffres annoncés sur un H200 à 8K tokens
Le benchmark présenté compare quatre configurations sur un seul GPU H200, avec Llama 3.1 8B Instruct comme teacher et un modèle Llama 3.2B comme student, pour un contexte de 8K tokens. Voici les valeurs rapportées :
- Distillation online : 102,8 Go de mémoire au pic, 25,9 secondes par itération et 237 TFLOP/s.
- Offline avec KL dense : 78,3 Go, 18,5 secondes et 331 TFLOP/s.
- Offline avec KL découpée en forward : 61,8 Go, 18,4 secondes et 335 TFLOP/s.
- Offline avec KL fusionnée et découpée : 58,3 Go, 20,2 secondes et 304 TFLOP/s.
À 8K, la variante forward-chunked est à la fois moins gourmande en mémoire que la KL dense et légèrement plus rapide dans ce tableau. La variante fusionnée réduit encore le pic de VRAM, mais avec un temps par itération supérieur. Son intérêt, selon les auteurs, se manifeste surtout lorsque le contexte s’allonge.
L’article source indique également qu’un exemple à 32K tokens fait passer le pic d’une KL dense vers environ 250 Go, contre environ 128 Go pour la perte fusionnée et découpée. Cette comparaison aide à comprendre le mécanisme mémoire, mais elle ne constitue pas à elle seule une mesure universelle : la consommation réelle dépend du modèle, du vocabulaire, du batch, de la précision numérique, des activations et de la stratégie de parallélisation.
Ce que tu peux en tirer pour tes pipelines d’entraînement
Si tu entraînes ou compresses des modèles ouverts, ce travail pointe une règle simple : avant d’ajouter des GPU, vérifie si ton pipeline matérialise des sorties denses inutilement. La distillation est souvent pensée comme un problème de modèle ; ici, l’optimisation vient surtout de l’architecture du calcul et de la gestion des données intermédiaires.
Concrètement, tu peux procéder dans cet ordre :
- mesure séparément la mémoire des poids, activations, logits et états d’optimisation ;
- évalue la pertinence d’un cache teacher Top-K pour ton jeu de données et tes tâches ;
- compare une perte dense et une perte traitée par chunks, à qualité d’entraînement égale ;
- teste explicitement les longues séquences, car les bénéfices mémoire n’apparaissent pas forcément sur un contexte court ;
- valide le modèle obtenu sur tes évaluations produit, et pas seulement sur la courbe de loss.
L’implémentation de la perte KL découpée est annoncée comme open source par les auteurs. Pour les équipes qui travaillent sur des modèles aux longues fenêtres de contexte, l’approche mérite un test reproductible dans leur propre stack. Les données fournies suggèrent un gain matériel tangible ; elles ne dispensent pas de vérifier le coût de préparation du cache, les performances finales du student et la compatibilité avec tes outils d’entraînement.