LRMoE élague la passe arrière des MoE et accélère l'entraînement de 21–45%

Une équipe de chercheurs du monde académique propose LRMoE, un système d'entraînement de modèles MoE sur des nœuds où la mémoire GPU fait défaut. MoE (Mixture-of-Experts, « mélange d'experts ») désigne une architecture dans laquelle une couche se compose de dizaines, voire de centaines de petits sous-réseaux, les experts, et où un routeur n'en active que quelques-uns pour chaque token. Les auteurs appellent « paire » un tel choix « token–expert ».

Le chiffre phare de l'étude : sur les modèles Qwen3-30B-A3B et GLM4.5-106B-A12B, le temps cumulé d'entraînement des couches MoE jusqu'à une valeur de fonction de perte (loss) fixée à l'avance a diminué de 21,2–44,8% par rapport à SYMI. Ce dernier est le système précédent : il dissocie le stockage des paramètres du modèle de celui des états de l'optimiseur afin de placer les experts plus souplement.

Pourquoi la copie des experts ne résout pas le déséquilibre de charge

Les experts sont répartis sur les GPU, et les tokens rejoignent les bonnes cartes via all-to-all, un échange « de chacun avec chacun ». Les experts populaires reçoivent davantage de tokens : leurs GPU deviennent le goulot d'étranglement pendant que les autres attendent. Une étape synchrone avance à la vitesse du participant le plus chargé.

Les réponses habituelles consistent à déplacer les experts ou à copier les plus sollicités. C'est l'approche de FasterMoE, FlexMoE, SmartMoE, LAER-MoE, UltraEP et TopoEP. Mais les copies ont un coût :

• une réplique occupe de la mémoire, dont on a besoin pour les tampons de tokens et les activations sauvegardées ;
• les copies doivent être synchronisées, et leurs gradients additionnés ;
• le volume total de travail ne diminue pas : il est simplement déplacé vers d'autres cartes.

Les auteurs montrent que, sous une limite stricte de mémoire GPU, même une répartition parfaitement uniforme peut laisser une charge globale supérieure à ce que la mémoire et le réseau peuvent absorber. Le déséquilibre et la surcharge au sens large sont deux problèmes distincts, et l'équilibrage ne traite que le premier.

Ce que fait LRMoE : passe avant complète, passe arrière élaguée

L'idée est de ne plus considérer la charge de la passe arrière comme une conséquence inévitable du routage, mais comme une grandeur réglable. Le fonctionnement se déroule en deux étapes.

• La passe avant reste complète. Toutes les paires token–expert sont calculées comme d'habitude : à ce stade, les gradients n'existent pas encore et il n'y a rien pour évaluer la valeur des paires.
• Des quotas avant le calcul. Un contrôleur voit combien de tokens sont arrivés chez chaque expert et, en tenant compte de la mémoire et du réseau, décide combien de paires de chaque expert conserver dans la passe arrière. Il planifie aussi le placement des répliques. Le quota est fixé pour toute la durée d'un appel de la couche.
• Sélection par la métrique LRM. Quand les gradients atteignent la couche, la contribution de chaque paire à la mise à jour des poids de l'expert est évaluée. Cette contribution est le produit extérieur du gradient entrant et de l'entrée sauvegardée de la projection. Le carré de sa norme se décompose donc en un produit de trois facteurs : le carré du poids de routage, le carré de la norme du gradient et le carré de la norme de l'entrée. Il n'est pas nécessaire de calculer les matrices de gradients complètes. Dans la limite du quota, les paires ayant le score le plus élevé sont conservées.
• Redistribution. Un algorithme en nombres entiers répartit les paires restantes entre les copies admissibles de leur expert, avec pour objectif de minimiser la charge de pointe sur un seul GPU. Aucune nouvelle réplique n'apparaît pendant la passe arrière.
• Compactage physique. Les lignes restantes de gradients et d'activations sont regroupées en blocs denses par destinataire et par expert. L'échange all-to-all devient plus compact, et les multiplications matricielles s'exécutent en grouped GEMM de taille variable, sans remplissage par des zéros. Le résultat est ensuite renvoyé vers la disposition d'origine.

Un détail important : se contenter de masquer les lignes superflues dans des tenseurs de taille complète n'apporte rien, ni en temps ni en mémoire. Le gain n'apparaît que lorsque les tailles réelles des transferts et des matrices diminuent.

Ce que montrent les expériences, et où se situe le piège

En plus du chiffre principal de 21,2–44,8%, les auteurs ont mené une reproduction distincte sur GPU et obtenu une baisse de 57,8% de la mémoire totale consacrée au placement des experts, aux répliques et aux tampons de tokens, par rapport à une attribution fixe des experts. Il s'agit d'une mesure de mémoire, pas du temps d'entraînement.

Le piège : la passe arrière devient approximative, puisqu'une partie des contributions au gradient est tout simplement écartée. La métrique LRM est une estimation au niveau d'une seule projection et n'offre aucune garantie de convergence. C'est pourquoi les auteurs comparent les systèmes sur le temps nécessaire pour atteindre une loss fixée, et non sur la vitesse par étape : l'accélération n'est comptabilisée que si le modèle atteint la même qualité. Les surcoûts liés au calcul des scores, à l'empaquetage et au routage sont inclus dans les mesures.

Pour les ingénieurs, la conclusion est la suivante. Si l'entraînement d'un MoE est limité par la mémoire GPU et le réseau, et qu'il n'y a plus de place pour ajouter des répliques, il vaut la peine d'envisager de réduire le travail de la passe arrière elle-même, et pas seulement de le déplacer. Les réserves comptent aussi. Les résultats portent sur deux modèles et ont été comparés à un seul système de référence : le transfert à d'autres architectures et à de longues exécutions doit donc être vérifié séparément. Le coût des gradients approximatifs pourrait apparaître face à des objectifs de qualité plus exigeants.