Speculating Experts Accelerates Inference for Mixture-of-Experts

Ce papier propose une méthode de préchargement d'experts pour les modèles MoE qui utilise les représentations internes du modèle pour anticiper les experts futurs, permettant ainsi de chevaucher les transferts mémoire et le calcul afin de réduire le temps d'inférence jusqu'à 14 % sans sacrifier la précision.

L'essentiel

Imaginez que vous avez un cuisinier génial (c'est votre modèle d'intelligence artificielle) qui doit préparer un repas complexe, bouchée par bouchée. Ce cuisinier est incroyablement doué, mais il a un problème : il n'a pas assez d'espace sur son plan de travail (la mémoire de votre carte graphique) pour ranger tous ses ingrédients.

Il possède une immense bibliothèque de recettes et d'épices (les "experts" du modèle) stockées dans une cave froide (la mémoire de votre processeur ou CPU), loin de la cuisine.

Le Problème : Le Cuisinier qui Attend

Dans la méthode actuelle (appelée "chargement à la demande"), voici ce qui se passe :

1. Le cuisinier prépare une bouchée.
2. Il réalise qu'il a besoin d'un ingrédient spécifique (par exemple, de la vanille).
3. Il doit arrêter de cuisiner, courir jusqu'à la cave froide, chercher la vanille, revenir et la mettre dans le plat.
4. Il recommence pour la prochaine bouchée.

Le problème, c'est que le trajet vers la cave est très long par rapport au temps de cuisson. Le cuisinier passe 85 % de son temps à courir et seulement 15 % à cuisiner. C'est lent et inefficace.

La Solution : La "Prévision Gourmande"

Les auteurs de ce papier ont eu une idée brillante : au lieu d'attendre de savoir quel ingrédient il faudra, le cuisinier devine ce dont il aura besoin pour la prochaine bouchée pendant qu'il cuisine la présente.

C'est ce qu'ils appellent "l'anticipation des experts" (Speculating Experts).

Voici comment cela fonctionne avec une analogie simple :

1. Le "Quasi-État" : L'instinct du Chef

Le cuisinier a un instinct très fort. En regardant ce qu'il est en train de faire (la bouchée actuelle), il peut presque sentir ce dont il aura besoin dans 5 secondes.

• Ils ont créé un outil mathématique (le "quasi-état") qui agit comme cet instinct. Il analyse la bouchée actuelle et dit : "Hé, je parie que la prochaine étape aura besoin de cannelle !"
• Pendant que le cuisinier est occupé à mélanger la cannelle pour la bouchée actuelle, un assistant invisible (le système de prévision) court déjà chercher la cannelle pour la bouchée suivante dans la cave.

2. Le Chevauchement : Cuisiner en courant

Au lieu d'attendre que le trajet soit fini pour recommencer, le cuisinier et l'assistant travaillent en parallèle.

• Avant : Cuisiner -> Arrêter -> Courir -> Revenir -> Cuisiner.
• Maintenant : Cuisiner + (Courir en même temps pour la prochaine étape).

Le temps perdu à courir disparaît presque totalement car il est "caché" par le temps de cuisson. Résultat : le repas est servi beaucoup plus vite (jusqu'à 14 % plus vite !).

Et si le cuisinier se trompe ?

C'est la partie la plus astucieuse. Parfois, l'instinct du cuisinier se trompe. Il pense qu'il aura besoin de cannelle, mais en réalité, il aurait dû prendre du sel.

• L'ancienne méthode : Si on se trompait, on devait s'arrêter, courir chercher le sel, et recommencer. C'était catastrophique.
• La nouvelle méthode : Les chercheurs ont découvert que même si le cuisinier se trompe sur l'ingrédient exact, utiliser l'ingrédient qu'il a apporté (la cannelle) ne gâche pas le plat. Le goût reste presque identique.
  • Analogie : Si vous prévoyez de mettre de la cannelle mais que le plat aurait besoin de sel, ajouter un peu de cannelle ne va pas transformer votre gâteau en poison. Le résultat final reste très bon.

Quand l'instinct ne suffit pas : Le "Petit Assistant"

Pour certains modèles très complexes (comme le Qwen-30B), l'instinct du cuisinier est parfois trop confus au début du repas (les premières couches du modèle). Là, il se trompe souvent.

Pour régler ça, les auteurs ont entraîné un petit assistant IA (un "estimator").

• Ce petit assistant est très rapide et n'a besoin que de quelques leçons (peu de données) pour apprendre à prédire les ingrédients avec une précision de 90 %.
• Il aide le cuisinier à faire les bons choix pour les étapes difficiles, garantissant que le plat final est parfait, même si le trajet vers la cave est toujours en cours.

En Résumé

Ce papier explique comment rendre les intelligences artificielles géantes beaucoup plus rapides sur des ordinateurs personnels (qui ont peu de mémoire).

Au lieu de faire attendre l'ordinateur pour aller chercher des données, on lui apprend à deviner l'avenir et à préparer les données pendant qu'il travaille. C'est comme si votre ordinateur devenait un chef d'orchestre qui fait jouer les musiciens pendant qu'il tourne la page de la partition, au lieu d'attendre que la page soit tournée avant de faire un bruit.

Le résultat ? Des IA plus rapides, plus fluides et accessibles sur des machines moins puissantes, sans perdre en qualité de réponse.

Résumé technique

1. Le Problème : Goulot d'étranglement I/O dans l'inférence MoE

Les modèles de langage (LLM) basés sur l'architecture Mixture-of-Experts (MoE) permettent de scaler la capacité des modèles tout en maintenant une activation parcimonieuse (seuls quelques "experts" sont activés par token). Cependant, dans des environnements contraints en mémoire (ex. : GPU grand public), les poids des experts doivent être déchargés sur la RAM du CPU.

• Le défi : Lors de l'inférence (phase de décodage), le transfert des poids des experts du CPU vers le GPU devient le goulot d'étranglement principal.
• Conséquence : Pour des modèles comme Qwen3-30B-A3B sur un GPU A6000, les transferts CPU-GPU représentent 84 à 88 % du temps par token de sortie (TPOT), rendant l'inférence limitée par les E/S (I/O-bound) plutôt que par le calcul.
• Limites des approches existantes : Les méthodes de "prefetching" (préchargement) existantes traitent souvent les erreurs de prédiction comme des "cache misses", obligeant à recharger les vrais experts sélectionnés par le routeur, ce qui annule les gains de performance potentiels.

2. Méthodologie : Préchargement et Exécution Spéculative

Les auteurs proposent un schéma de préchargement d'experts qui utilise les représentations internes du modèle pour prédire les experts futurs, permettant de chevaucher (overlap) les transferts de mémoire avec le calcul GPU.

A. Signal de Prédiction : L'État Quasi-Caché

Au lieu d'attendre le résultat du routeur de la couche actuelle pour prédire la suivante, le système utilise un état quasi-caché ($q_l$) :

• Il est construit à partir du flux résiduel normalisé ($r_l$) et d'un vecteur par défaut ($d_l$) qui capture la contribution typique des experts sélectionnés.
• $q_l = LN_{l+1}(d_l + r_l)$.
• Ce vecteur sert d'entrée pour prédire les experts de la couche $l+1$ avant même que le calcul de la couche $l$ ne soit terminé.

B. Exécution Spéculative (Speculative Execution)

Contrairement aux approches précédentes qui ignorent les experts prédits en cas d'erreur, cette méthode exécute les experts prédits.

• Si la prédiction est correcte : Le calcul se fait immédiatement, masquant le temps de transfert.
• Si la prédiction est incorrecte : Le modèle utilise quand même l'expert prédit. Les auteurs démontrent que cela préserve souvent la précision de la tâche en aval, évitant ainsi le coût de rechargement des vrais experts.

C. Estimateurs Neuronaux Légers

Pour les architectures où la dérive représentationnelle (drift) entre les couches est trop importante (rendant la prédiction basée sur l'état quasi-caché imprécise), les auteurs entraînent un estimateur neuronal léger.

• Cet estimateur apprend à mapper l'état quasi-caché directement vers les logits du routeur de la couche suivante.
• Il est entraîné par distillation avec un nombre minimal de tokens (quelques millions) et peut remplacer le routeur standard pour certaines couches critiques.

3. Contributions Clés

1. Préchargement sans paramètres (Parameter-free) : Identification de représentations internes (état quasi-caché) capables de prédire les décisions de routage futures sur des architectures MoE modernes sans entraînement supplémentaire.
2. Exécution spéculative préservant la précision : Démonstration que l'exécution des experts préchargés, même en cas d'erreur de prédiction, maintient la précision des tâches en aval, éliminant le besoin de rechargement d'urgence.
3. Implémentation optimisée : Intégration du schéma dans le moteur d'inférence open-source YALIS, utilisant des flux CUDA asynchrones pour chevaucher copie et calcul.
4. Estimateurs neuronaux : Introduction d'un correcteur léger pour améliorer les taux de réussite (hit rates) sur les couches à forte dérive, réduisant la dégradation de performance.

4. Résultats Expérimentaux

Les expériences ont été menées sur plusieurs modèles (Qwen3-30B-A3B, GPT-OSS-20B/120B, GLM-4.7-Flash) et configurations matérielles (A6000, A100, GH200).

• Performance (TPOT) :
  • Réduction du temps par token de sortie de 5 % à 14 % par rapport au chargement à la demande (on-demand).
  • Les gains sont plus importants sur les GPU moins puissants (A6000) et pour des séquences plus longues, où le temps de calcul est plus élevé, permettant un meilleur chevauchement avec le transfert de données.
  • L'approche approche la limite théorique de vitesse de calcul définie par le temps de copie CPU-GPU.
• Précision :
  • Pour les modèles GPT-OSS, l'exécution spéculative pure (Router-PF) préserve la précision sur des benchmarks de raisonnement (Math, Code, Common Sense).
  • Pour Qwen3-30B-A3B, la précision chute légèrement avec la méthode pure (surtout sur les tâches mathématiques) en raison d'une forte dérive dans les premières couches.
  • L'utilisation de l'estimateur neuronal (Est-PF) ou d'une stratégie Hybride (estimateur sur les couches critiques, routeur ailleurs) récupère la majorité de la perte de précision, ramenant les performances très proches du modèle de base.
• Taux de réussite (Hit Rate) :
  • L'état quasi-caché atteint un taux de réussite d'environ 90 % sur les couches stables.
  • L'estimateur neuronal améliore significativement la précision de prédiction sur les premières couches de Qwen3 (jusqu'à +25 % d'amélioration).

5. Signification et Impact

Ce travail est significatif car il rend le déploiement local de modèles MoE massifs (centaines de milliards de paramètres) plus viable sur du matériel grand public (GPU avec mémoire limitée).

• Efficacité I/O : En transformant un goulot d'étranglement de transfert mémoire en un problème de calcul masqué, l'approche améliore considérablement le débit.
• Accessibilité : Elle permet d'exécuter des modèles open-source de pointe (comme Qwen ou GPT-OSS) sur des cartes graphiques de bureau (ex. A6000) avec une latence acceptable, sans nécessiter de clusters de serveurs coûteux.
• Futur : L'approche ouvre la voie à une inférence sans routeur (router-free) et à des stratégies de préchargement multi-couches, élargissant les possibilités d'optimisation pour l'inférence sur dispositifs embarqués et mobiles.

En résumé, cette recherche propose une solution élégante au problème de la latence mémoire dans les MoE en combinant une prédiction intelligente des experts et une exécution spéculative tolérante aux erreurs, validée par des gains de performance tangibles et une préservation de la qualité du modèle.