La distillation des IA devient enfin abordable : voici comment

LA DISTILLATION D'IA : COMPRENDRE LE PRINCIPE

La distillation de connaissances est une technique bien connue en apprentissage automatique. Elle consiste à entraîner un petit modèle (modèle étudiant) pour qu'il reproduise les performances d'un grand modèle (modèle enseignant). Imaginez un professeur qui explique un cours à un élève : l'élève essaie de retenir les mêmes réponses que le professeur, mais en plus petit. Cette méthode est devenue essentielle avec l'arrivée des grands modèles de langage open source comme gpt-oss, Qwen ou GLM. Leur déploiement coûte extrêmement cher : le modèle Kimi-K3, par exemple, possède 2,8 billions de paramètres et nécessite environ 3 To de VRAM juste pour être chargé en mémoire.

La compression de ces géants en modèles plus petits tout en conservant leurs capacités originales est devenue une pratique standard. Des entreprises comme Nvidia avec son modèle Nemotron 3 Puzzle 75B ou Multiverse Computing avec Hypernova 60B ont récemment publié des modèles compressés de haute qualité. Pourtant, l'étape de distillation reste la plus coûteuse et la plus gourmande en ressources du processus. Garder les deux modèles (enseignant et étudiant) en mémoire en même temps, et produire une distribution de probabilités sur l'ensemble du vocabulaire pour chaque token, nécessite des quantités astronomiques de VRAM. Cela n'est généralement réalisable qu'avec des centaines de GPU et des stratégies complexes de parallélisme tensoriel.

DEUX INNOVATIONS QUI CHANGENT TOUT

Une nouvelle étude intitulée Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss propose deux changements majeurs pour rendre la distillation abordable à grande échelle. Le premier consiste à cacher les prédictions du modèle enseignant une fois pour toutes, et le second à utiliser une nouvelle fonction de perte plus efficace en mémoire. Ensemble, ces innovations réduisent suffisamment les coûts d'entraînement pour permettre la distillation de modèles à long contexte sur un seul GPU, et rendent les expérimentations à grande échelle pratiques et économiques.

Pour comprendre pourquoi la distillation classique est si coûteuse, examinons son fonctionnement. Dans une distillation en ligne classique, le modèle enseignant et le modèle étudiant sont chargés simultanément. À chaque étape d'entraînement, le modèle enseignant effectue un passage avant complet pour produire sa distribution de sortie, et le modèle étudiant est entraîné pour correspondre à cette distribution. Cette méthode est la plus expressive car la distribution complète du modèle enseignant est disponible, mais elle est aussi la plus gourmande en mémoire et en calcul : deux tenseurs de taille égale à la taille du vocabulaire doivent être maintenus pour chaque position de token, et le modèle enseignant doit être recalculé à chaque étape, même si son comportement ne change pas pendant l'entraînement.

LE PROBLÈME DE LA MÉMOIRE : UNE TENSION QUI EXPLOSE

Prenons un exemple concret. Le modèle gpt-oss-120b possède un vocabulaire de 201 088 tokens. Avec une longueur de séquence de 32 000 tokens et une taille de lot de 4, le tenseur des probabilités du modèle enseignant seul a la forme 4 × 201 088 × 32 768. En format bfloat16, cela représente déjà environ 50 Go de VRAM pour un seul tenseur. Ajoutez les gradients, les activations, les poids du modèle et les états de l'optimiseur, et une seule itération d'entraînement de distillation peut atteindre environ 250 Go de VRAM. C'est plus que ce qu'un GPU H200 ou B200 peut fournir. La nouvelle méthode reformule la fonction de perte KL pour traiter les données par morceaux, réduisant ce coût à presque rien.

Le graphique ci-dessous montre l'évolution de la mémoire utilisée : la méthode classique atteint environ 250 Go, dépassant largement la capacité d'un H200 (141 Go), tandis que la méthode par morceaux ne dépasse jamais 128 Go.

PREMIÈRE INNOVATION : LA DISTILLATION HORS LIGNE

Au lieu de recalculer le modèle enseignant à chaque étape, la nouvelle méthode calcule ses prédictions une seule fois, cache les 100 tokens les plus probables par position, et entraîne le modèle étudiant sur ce cache. Le modèle enseignant n'a plus besoin d'être en mémoire pendant l'entraînement et n'a pas besoin d'être relancé une fois le cache créé. Ce cache peut même être réutilisé pour de nombreuses expériences et comparaisons.

DEUXIÈME INNOVATION : UNE FONCTION DE PERTE PLUS ÉCONOME

Pour comprendre pourquoi la fonction de perte classique est si gourmande en mémoire, imaginez ce qu'elle construit en réalité. Pour chaque position de token dans une séquence et chaque mot du vocabulaire, la perte a besoin d'un nombre décrivant à quel point la prédiction de l'étudiant diffère de celle de l'enseignant. Si on représente cela sous forme de grille, chaque ligne correspond à une entrée du vocabulaire et chaque colonne à une position dans la séquence. Avec un vocabulaire de plus de 100 000 mots et une longue séquence, cette grille devient monstrueuse. La méthode classique de calcul d'une perte KL construit toute cette grille avant de pouvoir produire un seul résultat.

Trois méthodes différentes pour calculer cette même perte, toutes mathématiquement équivalentes, ont été comparées. L'animation ci-dessous montre la différence entre les approches dense et par morceaux : l'une construit toute la grille de comparaison et la conserve en mémoire, l'autre construit et rejette une tranche à la fois, de sorte que la mémoire ne dépasse jamais la taille d'un seul morceau.

L'implémentation de la fonction de perte par morceaux a été open sourcée : github.com/CompactifAI/Full-Chunked-KL-Loss

CE QUE ÇA CHANGE EN PRATIQUE : LES RÉSULTATS SONT BLUFFANTS

Le tableau ci-dessous compare quatre configurations différentes : la distillation en ligne classique et les trois implémentations hors ligne de la fonction de perte décrites précédemment. Toutes ont été testées sur un seul GPU H200 avec Llama 3.1 8B Instruct comme modèle enseignant et un modèle Llama 3.2B comme étudiant, avec un contexte de 8 000 tokens. Les quatre méthodes atteignent des pertes d'entraînement quasi identiques, même si les runs hors ligne s'entraînent uniquement sur les 100 logits les plus probables par token mis en cache.

Les courbes de perte se superposent presque parfaitement pour les quatre méthodes, confirmant que la distillation hors ligne avec les 100 logits les plus probables par token est aussi précise que la distillation en ligne classique. À cette longueur de séquence, la fonction de perte par morceaux n'est pas encore la plus rapide : ses coûts supplémentaires de projection lors du passage arrière réduisent un peu la vitesse. Son avantage réel n'apparaît qu'avec l'augmentation de la longueur de contexte, comme le démontre la section suivante.

LA PERFORMANCE SE DÉVOILE À GRANDE ÉCHELLE

La distillation d'un modèle GPT-OSS 20B avec un contexte de 32 768 tokens a permis de réduire l'infrastructure nécessaire de quatre nœuds GPU à un seul GPU. Le temps par étape est passé de 57,0 à 12,23 secondes, soit environ 5 fois plus rapide, et le débit par GPU est passé de 74,2 à 345,7 TFLOP/s.

LE MODÈLE FINAL : COMPACT MAIS PERFORMANT

Cette configuration efficace de distillation hors ligne a rendu possible une campagne de distillation à grande échelle. Le modèle étudiant compact résultant, distillé à partir de Llama 3.1 8B Instruct jusqu'à environ 3,2 milliards de paramètres, conserve la plupart de la précision de l'enseignant sur BoolQ et HellaSwag, et reste dans une marge d'environ neuf points sur MMLU, avec moins de la moitié du nombre de paramètres.

Cette étude s'inscrit dans le cadre des recherches continues de Multiverse Computing pour rendre la distillation et la récupération de modèles pratiques à grande échelle. Ce n'est pas seulement une recette ponctuelle, mais une méthode que les équipes peuvent itérer de manière économique. L'article couvre également des expériences supplémentaires, comme l'impact du choix de la fonction de perte et du regroupement des séquences sur la qualité de récupération.

POUR ALLER PLUS LOIN : LES DÉTAILS TECHNIQUES

Si vous souhaitez les détails techniques complets, y compris le gradient en forme fermée derrière la fonction de perte par morceaux et la configuration d'entraînement complète, vous pouvez consulter l'article complet ou contacter l'équipe pour discuter de l'application de cette méthode à vos propres pipelines de distillation. L'implémentation de la fonction de perte par morceaux est également open source : github.com/CompactifAI/Full-Chunked-KL-Loss

LES MODÈLES ET ARTICLES MENTIONNÉS

Modèles mentionnés dans cet article :
- gpt-oss
- Qwen
- GLM
- Kimi
- Llama 3.1 8B Instruct
- Llama 3.2B
- GPT-OSS 20B
- Nemotron 3 Puzzle 75B
- Hypernova 60B

Articles mentionnés dans cet article :
- Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss

Collections mentionnées dans cet article :
- CompactifAI

POURQUOI C'EST UNE RÉVOLUTION POUR L'IA

Cette avancée est majeure car elle rend enfin accessible la distillation des grands modèles de langage à toute équipe, même avec des ressources limitées. Avant, seules les grandes entreprises avec des centaines de GPU pouvaient se permettre d'entraîner des modèles compressés de qualité. Aujourd'hui, avec ces deux innovations, une seule personne avec un GPU standard peut distiller un modèle comme Llama 3.1 8B Instruct en un modèle 3,2B presque aussi performant. Les économies de mémoire et de calcul sont telles que les expérimentations deviennent possibles pour des budgets modestes. Les implications sont immenses : plus d'équipes peuvent innover, tester des idées, et contribuer à l'avancement de l'IA sans se ruiner.

LES LIMITES ET PERSPECTIVES D'AVENIR

Malgré ces progrès impressionnants, il reste des défis à relever. La méthode par morceaux, bien que très efficace, ajoute une légère complexité lors du passage arrière (backward pass), ce qui peut réduire légèrement la vitesse d'entraînement à des longueurs de contexte modérées. Son avantage maximal ne se révèle qu'avec des contextes très longs, au-delà de 64 000 tokens. De plus, la qualité de la distillation dépend fortement de la qualité du cache des logits du modèle enseignant : si ce cache est de mauvaise qualité, la performance du modèle étudiant en pâtira. Enfin, cette méthode a été testée principalement avec des modèles de langage, et son efficacité sur d'autres types de modèles (vision par ordinateur, multimodal) reste à démontrer.

Les chercheurs envisagent déjà d'étendre cette méthode à d'autres architectures et d'explorer des optimisations supplémentaires. L'objectif est clair : rendre la distillation encore plus efficace, plus rapide et encore moins gourmande en ressources. À terme, ces innovations pourraient permettre de distiller des modèles encore plus grands, avec des contextes encore plus longs, tout en gardant des coûts maîtrisés.

CONCLUSION : L'IA COMPRIMÉE DEVIENT LA NORME

La distillation des modèles d'IA n'est plus un luxe réservé aux entreprises capables d'investir des millions dans des infrastructures. Grâce à ces deux innovations — la distillation hors ligne et la fonction de perte par morceaux — elle devient enfin accessible à tous. Les économies réalisées sont telles que même les petits budgets peuvent désormais entraîner des modèles compressés de haute qualité. Cette avancée ouvre la porte à une nouvelle ère où l'innovation en IA ne sera plus limitée par les coûts, mais par l'imagination des chercheurs et des développeurs. Le futur de l'IA n'est pas seulement dans les modèles toujours plus grands, mais aussi dans leur capacité à être comprimés, partagés et utilisés par tous, partout.