DisaggKD déporte le teacher sur des nœuds dédiés et accélère la distillation de Qwen3.5-35B-A3B vers 9B

L'équipe derrière DisaggKD propose de lever une contrainte habituelle de la distillation de connaissances : l'enseignant et l'élève ne sont plus obligés de vivre sur les mêmes GPU. Sur la paire Qwen3.5-35B-A3B (teacher) et 9B (student), le schéma disaggrégé a obtenu le même LM loss que le lancement colocated avec le même loss, tout en s'exécutant plus vite. Sur des paires plus grandes, les lancements colocated butent sur un manque de mémoire, tandis que les lancements disaggrégés continuent de s'entraîner.

Pourquoi le mode colocated se heurte à la mémoire

Dans la plupart des frameworks de KD, le teacher et le student partagent les mêmes accélérateurs. Chaque carte doit donc stocker les poids du teacher, les poids du student, les gradients et les états de l'optimiseur. Plus le teacher est puissant, moins il reste de place pour le reste : la taille du teacher est bornée par la mémoire d'un seul GPU.

Un second problème se cache dans la loss elle-même. La KD classique compare les distributions du teacher et du student sur l'ensemble du vocabulaire. Cela exige des logits de taille « tokens × vocabulaire », et les modèles actuels ont de grands vocabulaires. Cette matrice devient l'un des principaux postes de consommation de mémoire au pic.

Comment fonctionne DisaggKD

Les auteurs décomposent leur schéma en trois parties.

Première partie : la séparation par nœuds. Les répliques du teacher tournent à part et transmettent au student leurs états cachés (hidden states). Les logits ne transitent pas par le réseau.

Deuxième partie : la fused loss. La KD-loss est calculée par fenêtres de lignes et par shards du vocabulaire. En arithmétique exacte, le résultat est égal à la loss sur le vocabulaire complet, mais la matrice de logits complets n'est jamais formée. L'équivalence est annoncée précisément pour l'arithmétique exacte. En pratique, l'ordre de sommation en float peut produire de petits écarts ; les auteurs comparent donc séparément les courbes d'entraînement réelles.

Troisième partie : le prefetch ring. Il s'agit d'un tampon circulaire qui cale le rythme du producteur (le teacher) sur celui du consommateur (le trainer). Pendant que le student calcule un pas, les états cachés suivants sont déjà en route vers lui, et aucun des deux côtés n'attend.

Les auteurs fournissent en outre des modèles analytiques fermés de la mémoire et du trafic réseau. Ils permettent d'estimer à l'avance, pour chaque configuration (colocated ou disaggrégée), où se situe la limite de mémoire.

Ce que montrent les mesures

• Paire Qwen3.5-35B-A3B → 9B : le schéma disaggrégé atteint le même LM loss que le schéma colocated avec le même loss, et s'exécute plus vite que lui.
• La fused loss est utile à elle seule. Un lancement colocated qui l'utilise peut tourner sur moins de nœuds et plus rapidement.
• Si le student occupe plusieurs nœuds, le schéma disaggrégé est plus économe en mémoire.
• Sur des paires plus grandes, les lancements colocated tombent en OOM même avec la fused loss : les poids du teacher et les états du student consomment tout le budget. Le schéma disaggrégé, lui, s'entraîne sur ces paires.

L'abstract ne donne ni chiffres absolus d'accélération, ni nombre de GPU, ni tailles des grandes paires. L'ampleur du gain est à chercher dans l'article lui-même, puis à vérifier sur son propre matériel.

À qui cela servira

Les conclusions des auteurs dessinent un arbitrage clair. Si le teacher est petit et que le student tient sur un ou deux nœuds, le schéma colocated avec fused loss suffit : il économise déjà de la mémoire et accélère l'entraînement. Il faut envisager la configuration disaggrégée lorsque le teacher atteint des dizaines de milliards de paramètres (comme 35B-A3B avec ses experts) ou que le student est réparti sur plusieurs nœuds. Dans ce cas, l'équilibre mémoire décide si le lancement est tout simplement possible.

Il y a aussi un coût. Le schéma disaggrégé exige une liaison rapide entre les nœuds : les états cachés doivent être streamés sans pause. C'est pourquoi les modèles de trafic de l'article sont utiles en pratique : avant de louer un cluster, on peut les injecter dans son calcul et vérifier que le réseau ne deviendra pas le goulot d'étranglement.

L'idée de la fused loss mérite une mention à part. On peut la transposer dans des pipelines de KD existants sans modifier la topologie du cluster, et supprimer immédiatement le pic de mémoire dû aux logits. Pour les équipes qui distillent de grands modèles ouverts vers des modèles compacts, c'est le moyen le moins coûteux de tester l'approche sur leurs propres données.