SkewAdam réduit la mémoire d'entraînement des modèles Mixture-of-Experts de plus de 60 % et offre un bond notable en précision, permettant aux développeurs de faire tenir un modèle MoE de 6,78 milliards de paramètres sur un seul GPU de 40 Go.

Pourquoi l'entraînement des MoE se heurte à un mur de mémoire

Les architectures Mixture-of-Experts acheminent chaque entrée via un petit sous-ensemble de sous-réseaux « experts » tout en conservant une structure de base (backbone) dense. L'avantage est d'obtenir un modèle capable de passer à des milliards de paramètres sans augmentation proportionnelle du calcul. En pratique, l'optimiseur — plus précisément la variante AdamW utilisée par la plupart des équipes — consomme la majeure partie de la RAM du GPU. Pour un modèle de 6,78 milliards de paramètres, les tenseurs de momentum et de variance de l'optimiseur exigent à eux seuls environ 50 Go. Ajoutez les activations et les poids du modèle, et la mémoire de pointe dépasse les 81,4 Go, imposant des configurations multi-GPU ou des cycles d'entraînement plus lents.

Ce qui différencie SkewAdam

SkewAdam n'invente pas de nouvelle règle d'apprentissage. Il réorganise l'emplacement des moments d'Adam :

  • La structure de base dense conserve un momentum en pleine précision, préservant ainsi les mises à jour fluides dont les couches denses ont besoin.
  • La banque d'experts stocke une approximation factorisée de la variance, réduisant ainsi les données conservées pour chaque expert.
  • Le routeur, qui décide quels experts activer, conserve une estimation exacte du second moment.

En traitant ces trois composants comme des « niveaux » distincts, l'optimiseur supprime la précision redondante là où elle est moins importante et la conserve là où elle est cruciale.

Impact mesurable

L'exécution du même modèle MoE de 6,78 milliards de paramètres avec SkewAdam produit les résultats suivants :

  • Mémoire GPU de pointe : 31,3 Go (contre 81,4 Go auparavant)
  • Empreinte de l'état de l'optimiseur : 1,29 Go (contre 50 Go auparavant)

L'état réduit tient confortablement dans un seul accélérateur de 40 Go.

Gains de précision, pas seulement d'économies

Les réductions de mémoire nuisent souvent à la qualité du modèle, mais SkewAdam améliore la perplexité — une métrique standard des modèles de langage — passant de 126,8 (AdamW) à 108,4. Il surpasse également Muon (120,2) et Lion (393,7) sur la même tâche. Les auteurs affirment que ce gain provient de la préservation du momentum à travers les trois niveaux ; les méthodes qui abandonnent totalement le momentum, comme Adafactor, sont à la traîne.

Questions en suspens

Les résultats de SkewAdam sont démontrés sur un modèle de 6,78 milliards de paramètres. Il reste à déterminer si les mêmes ratios d'état de mémoire s'appliquent à des modèles d'un ordre de grandeur supérieur ou à des tâches dépassant la modélisation du langage.

À surveiller

  • Les benchmarks sur des configurations MoE plus larges (des dizaines de milliards de paramètres).
  • L'adoption par les frameworks open-source et l'inclusion dans les scripts d'entraînement populaires.
  • Les retours de la communauté sur les compromis cachés, tels que la vitesse de convergence ou la stabilité sous différents programmes de taux d'apprentissage.

Source : article de développeur détaillant la conception de l'optimiseur et ses résultats empiriques.