CRAFT: Cost-aware Expert Replica Allocation with Fine-Grained Layerwise Estimations

Le papier présente CRAFT, un cadre d'allocation de répliques d'experts à coût conscient qui optimise l'équilibrage de charge et le débit de service des grands modèles de langage en effectuant une réplication fine par couche basée sur une estimation des bénéfices, le tout sans nécessiter de réentraînement ni de modifications du modèle.

L'essentiel

Imaginez que vous gérez un immense restaurant appelé MoE (Mélange d'Experts). Ce restaurant est célèbre pour sa capacité à servir des millions de clients (les données) très rapidement.

Pour y parvenir, le restaurant ne compte pas sur un seul chef, mais sur une armée de 300 à 400 chefs spécialisés (les "experts"). Chaque client qui entre est dirigé vers seulement quelques-uns de ces chefs (généralement 8) selon ce qu'il commande. C'est très efficace : le restaurant peut grandir sans avoir besoin de plus de cuisines, car chaque client ne sollicite qu'une petite équipe.

Le problème : L'embouteillage
Le hic, c'est que les clients ne sont pas répartis équitablement. Certains plats sont très à la mode (comme les "tacos" ou les "burgers"), et soudainement, 90 % des clients commandent ces plats.

• Résultat : Les chefs spécialisés dans les tacos sont submergés, en sueur, et ne peuvent plus rien faire.
• Pendant ce temps, les chefs spécialisés dans les plats exotiques (que personne ne commande) sont assis à ne rien faire, en train de jouer aux cartes.
• Le restaurant entier ralentit parce que tout le monde attend que les chefs "tacos" finissent leur travail. C'est ce qu'on appelle le déséquilibre de charge.

La solution habituelle (EPLB) : Le copier-coller coûteux
Pour résoudre ce problème, la méthode actuelle (appelée EPLB) consiste à dire : "Bon, on va copier le chef 'tacos' 64 fois !"
On installe donc 64 chefs "tacos" dans la cuisine, répartis sur toutes les tables.

• Avantage : Plus personne n'est débordé, tout le monde travaille.
• Inconvénient majeur : La cuisine est pleine ! Chaque chef a besoin d'un espace de travail (de la mémoire GPU). En ayant 64 chefs "tacos", on n'a plus assez de place pour les ingrédients (la mémoire pour les conversations en cours, le "KV Cache"). Le restaurant devient lent non pas parce que les chefs sont lents, mais parce qu'ils n'ont plus de place pour travailler. C'est comme essayer de faire entrer 64 cuisiniers dans une cuisine de taille moyenne : on s'écrase les coudes et on perd du temps.

La nouvelle solution : CRAFT (Le chef intelligent)
C'est là qu'intervient CRAFT. Au lieu de copier bêtement tous les chefs, CRAFT agit comme un manager très fin et économe.

Voici comment CRAFT fonctionne, étape par étape :

1. L'observation (La carte de chaleur) :
   Avant d'ouvrir, CRAFT regarde les commandes des 3 000 derniers jours. Il se rend compte que :

   • Le chef "tacos" est fou de travail (il faut le copier).
   • Le chef "sushi" est un peu occupé (peut-être le copier une fois).
   • Le chef "dessert" est tranquille (pas besoin de le copier du tout).

2. L'allocation intelligente (Le budget mémoire) :
   CRAFT a un budget limité d'espace dans la cuisine. Il ne veut pas gaspiller de place.

   • Il dit : "On va mettre 4 chefs 'tacos' ici, 2 chefs 'sushi' là, et aucun chef 'dessert' ailleurs."
   • Il ne copie que ceux qui en ont vraiment besoin, et seulement autant qu'il faut.

3. Le résultat :

   • Moins de gaspillage : On n'occupe pas la cuisine avec des chefs inutiles. Il reste beaucoup de place pour les ingrédients (le KV Cache).
   • Plus de rapidité : Comme il y a de la place, on peut servir plus de clients en même temps. Le restaurant tourne à plein régime.

L'analogie du bus
Imaginez un bus qui doit transporter des passagers vers différentes destinations.

• La méthode ancienne : Le bus a 100 sièges. Il remplit 80 sièges avec des gens qui vont à la même destination (les "tacos"), et laisse 20 sièges vides pour les autres destinations. Le bus est lourd, lent, et ne peut pas prendre beaucoup de passagers.
• La méthode CRAFT : Le manager regarde la liste. Il voit que 80 % des gens vont à "Tacos". Il décide de mettre juste assez de sièges supplémentaires pour les "Tacos" (par exemple, 10 sièges de plus), et laisse le reste pour les autres. Le bus est plus léger, il peut prendre plus de passagers au total, et il arrive plus vite à destination.

En résumé
CRAFT est un système qui dit : "Arrêtons de copier tout le monde bêtement. Regardons qui travaille vraiment dur, et copions-les intelligemment, couche par couche, pour ne pas gaspiller l'espace de la cuisine."

Grâce à cela, les géants de l'intelligence artificielle (comme les modèles de 1000 milliards de paramètres) peuvent fonctionner plus vite et moins cher, car ils utilisent mieux leurs ressources informatiques. C'est comme passer d'un restaurant où l'on jette la moitié des ingrédients à la poubelle, à un restaurant où chaque gramme compte.

Résumé technique

1. Le Problème : Déséquilibre de Charge dans les MoE à Grande Échelle

Les modèles de langage (LLM) basés sur l'architecture Mixture-of-Experts (MoE) permettent de scaler le nombre de paramètres tout en maintenant un coût de calcul par token constant. Cependant, leur déploiement via le parallélisme d'experts (Expert Parallelism - EP) introduit un goulot d'étranglement majeur : le déséquilibre de charge (load imbalance).

• Cause : Le routeur MoE distribue les tokens vers un sous-ensemble d'experts. En raison de la distribution de Zipf des tokens naturels, certains experts deviennent "chauds" (très sollicités) tandis que d'autres restent "froids".
• Conséquences : Cela crée des goulots d'étranglement sur les GPU hébergeant les experts chauds, entraînant de l'inactivité sur les autres GPU et une congestion réseau lors des opérations all-to-all.
• Limites des solutions existantes :
  • Le placement d'experts (co-localiser experts chauds et froids) ne suffit pas lorsque la charge est extrêmement déséquilibrée.
  • La réplication d'experts (copier les experts chauds sur plusieurs GPU) est efficace pour équilibrer la charge, mais elle consomme énormément de mémoire GPU.
  • Les techniques actuelles (comme EPLB) utilisent une réplication uniforme (une copie par couche par GPU), ce qui entraîne un gaspillage de mémoire : de nombreuses copies apportent un gain de performance marginal, réduisant la taille du cache KV disponible et dégradant le débit global.

2. Méthodologie : CRAFT

Les auteurs proposent CRAFT, un cadre d'allocation de répliques d'experts conscient des coûts et basé sur des estimations fines par couche. L'objectif est de maximiser l'équilibre de charge sous une contrainte de mémoire donnée.

Observations Clés

L'analyse des auteurs révèle deux phénomènes cruciaux :

1. Rendements décroissants : L'ajout de répliques améliore l'équilibre de charge, mais les gains diminuent rapidement après un certain seuil.
2. Hétérogénéité par couche : Toutes les couches du modèle ne bénéficient pas de la même manière de la réplication. Certaines couches sont déjà bien équilibrées (faible skew), tandis que d'autres sont fortement déséquilibrées (fort skew). Une allocation uniforme gaspille de la mémoire sur les couches équilibrées et en manque sur les couches critiques.

Flux de Travail de CRAFT

CRAFT fonctionne en trois étapes principales :

1. Estimation du Bénéfice par Couche (Offline) :

   • Le système analyse la distribution de charge des experts hors ligne (via un jeu de données de profilage).
   • Il estime le gain d'équilibre (bénéfice) pour chaque couche en fonction du nombre de répliques (de 1 à un nombre maximal). Cela génère une matrice de gains potentiels.

2. Optimisation de l'Allocation (Dynamic Programming) :

   • CRAFT formule le problème d'allocation comme un Problème du Sac à Dos Multi-Choix (MCKP).
   • L'objectif est de sélectionner le nombre de répliques optimal pour chaque couche individuellement, de manière à maximiser le gain total d'équilibre tout en respectant le budget mémoire global (défini par l'utilisateur ou automatiquement).
   • Contrairement à EPLB qui alloue le même nombre de répliques partout, CRAFT alloue plus de répliques aux couches "fortement skewées" et zéro réplique aux couches déjà équilibrées.

3. Assignation et Placement (Capacity-Aware) :

   • Pour gérer le fait que le nombre de répliques varie par couche, CRAFT utilise un algorithme d'assignation entrelacé (interleaved).
   • Il attribue les experts et les répliques aux GPU de manière à minimiser l'imbalance de capacité mémoire entre les appareils, garantissant que la taille du cache KV reste uniforme sur tous les GPU (évitant ainsi la fragmentation mémoire).

3. Contributions Clés

• Caractérisation inédite : C'est la première étude détaillée montrant comment le débit et l'équilibre de charge évoluent avec la réplication, mettant en évidence le compromis mémoire/performance et la nécessité d'une granularité par couche.
• Framework CRAFT : Un système de bout en bout qui s'intègre directement dans les frameworks de service existants (comme SGLang) sans nécessiter de réentraînement du modèle ni de modifications architecturales.
• Efficacité des ressources : CRAFT démontre qu'il est possible d'atteindre un équilibre de charge quasi-optimal avec une fraction de la mémoire requise par les méthodes uniformes.

4. Résultats Expérimentaux

Les évaluations ont été menées sur un cluster de 8 nœuds (AWS p4de.24xlarge) avec des modèles de grande envergure (DeepSeek-R1-671B et Kimi-K2-1000B) et divers jeux de données.

• Gain de Débit (Throughput) : CRAFT améliore le débit de bout en bout de 1,14 fois en moyenne (jusqu'à 1,2 fois) par rapport aux techniques de réplication existantes (EPLB).
• Réduction de la Mémoire : CRAFT utilise 7,25 à 7,5 fois moins de répliques que EPLB pour atteindre des gains d'équilibre similaires, libérant ainsi de la mémoire pour le cache KV et permettant de traiter des lots de séquences plus grands.
• Latence (TTFT) : CRAFT réduit le Time-To-First-Token de 29 % en moyenne par rapport à la base (sans réplication), rivalisant avec les performances d'EPLB tout en étant beaucoup plus économe en mémoire.
• Robustesse : Les gains sont observés sur des jeux de données très déséquilibrés (où EPLB échoue souvent à cause de la sur-réplication) et sur des clusters de tailles variées (6, 8 et 12 nœuds).
• Surcoût : L'overhead d'initialisation (estimation des bénéfices) est négligeable (~10 secondes) et n'affecte pas la latence d'inférence en temps réel.

5. Signification et Impact

Ce travail est significatif car il résout un problème critique de déploiement des LLMs MoE à grande échelle : l'efficacité opérationnelle.

• Optimisation Coût-Performance : En évitant la sur-réplication, CRAFT permet de réduire les coûts d'infrastructure (moins de GPU nécessaires pour la même charge de travail) ou d'augmenter la capacité de traitement (plus d'utilisateurs simultanés) sur une infrastructure fixe.
• Adoption Pratique : Le fait que CRAFT ne nécessite aucune modification du modèle ni de réentraînement le rend immédiatement applicable aux infrastructures de production actuelles.
• Nouveau Paradigme : Il déplace la stratégie de réplication d'une approche "uniforme et statique" vers une approche "adaptative, fine-granularité et orientée bénéfice", ouvrant la voie à des optimisations plus intelligentes pour les systèmes distribués futurs.

En résumé, CRAFT prouve que l'allocation intelligente et économique des ressources de réplication est la clé pour débloquer le plein potentiel des modèles MoE géants dans des environnements de production réels.