Résumé Technique : PagedWeight

Problématique

Les modèles de langage à grande échelle (LLM) de type Mixture-of-Experts (MoE) offrent une grande efficacité et précision en n'activant qu'un sous-ensemble d'experts par jeton (token). Cependant, dans les scénarios de service (serving), les modèles MoE font face à une tension de mémoire critique : le GPU doit stocker à la fois les poids massifs du modèle et le cache Clé-Valeur (KV cache) croissant. Bien que des techniques comme PagedAttention gèrent la fragmentation du cache KV, elles ne traitent pas l'empreinte mémoire significative des poids MoE, qui peuvent occuper plus de 60 % de la mémoire du GPU.

Les méthodes de quantification existantes sont largement statiques, fixant la précision avant le début de l'inférence. Bien que certaines méthodes d'adaptation de précision au moment de l'exécution existent, elles n'utilisent pas les changements de précision pour réallouer dynamiquement la mémoire GPU entre les poids et le cache KV. Par conséquent, il manque des systèmes capables de libérer dynamiquement de la mémoire des poids d'experts moins critiques à mesure que le cache KV s'étend, sans sacrifier la précision des tâches ni interrompre l'inférence.

Méthodologie : PagedWeight

Les auteurs proposent PagedWeight, un nouveau système de gestion pour le service de LLM MoE qui quantifie dynamiquement les poids du modèle au moment de l'exécution afin d'équilibrer la précision des poids d'experts avec la taille du cache KV. Le système traite les poids d'experts de manière similaire aux blocs de cache KV paginés, permettant un déchargement (offloading) et un rechargement dynamiques.

Composants Clés

1. Représentation des Poids Paginés (Paged Weight Representation) :

   • PagedWeight opère à la granularité des blocs linéaires par expert (identifiés par couche, expert et type de module : gate_up ou down).
   • Il utilise la technologie Any-Precision LLM (APL), représentant les poids dans un format de plan de bits superposé avec des tables de recherche (LUTs) associées.
   • Une Table de Pages de Poids (Weight Page Table) suit la largeur de bits souhaitée ($d_i$) et la largeur de bits engagée ($q_i$) pour chaque bloc linéaire. Les pages peuvent se trouver dans des états résidents sur le GPU, résidents sur le CPU ou en cours de transfert.

2. Planificateur de Qualité au Moment de l'Exécution (Quality-Aware Runtime Planner) :
   Le planificateur détermine quels poids quantifier (réduire la largeur de bits) en se basant sur trois facteurs pour minimiser la perte de précision :

   • Sensibilité Globale Hors Ligne (Offline Global Sensitivity) : Utilise des scores pondérés par la Hessienne pour estimer la sensibilité intrinsèque de chaque bloc linéaire à la quantification.
   • Statistiques de Routage en Ligne (Online Routing Statistics) : Suit la « masse de routage » (fréquence et poids de sélection) des experts. Les experts fréquemment routés (« chauds ») sont protégés par des planchers de largeur de bits plus élevés et des multiplicateurs de dommages.
   • Résidus de Prompt (Prompt Residuals) : Un terme de correction spécifique au prompt prédit via des têtes de régression linéaire. Cela ajuste l'estimation de la sensibilité globale basée sur la séquence d'entrée actuelle, reconnaissant que l'impact de la quantification varie selon le prompt.

   Le planificateur calcule un score de « dommage » pour les réductions potentielles de largeur de bits. Lorsque la pression du cache KV déclenche une cible de mémoire, le planificateur sélectionne de manière gloutonne les réductions de dommage les plus faibles (plus haute qualité par octet économisé) pour atteindre la cible.

3. Exécution Asynchrone et Optimisation de Kernel :

   • Mouvement de Page Asynchrone : Pour masquer la latence, le système décharge les pages de poids vers la RAM du CPU et les recharge de manière aschrone. Les engagements de largeur de bits ne sont mis à jour qu'aux limites sûres (par exemple, après qu'un plan a été pleinement exécuté), garantissant que l'inférence n'est pas interrompue.
   • Kernel de Précision Mixte Fusionné (Fused Mixed-Precision Kernel) : Un kernel CUDA personnalisé lit directement les plans de bits et les LUTs de précision quelconque, supportant différentes largeurs de bits pour chaque bloc linéaire au sein d'une seule opération fusionnée afin de réduire les surcoûts (overhead).

Principales Contributions

1. Conception du Système : La proposition de PagedWeight, un système de poids paginés pour les poids MoE Any-Precision en ligne qui gère dynamiquement les largeurs de bits engagées/souhaitées, les états de pages et la consommation de mémoire.
2. Algorithme de Planification : Un planificateur de runtime sensible à la qualité qui synthétise la sensibilité hors ligne, les statistiques de routage en ligne et les résidus de prompt pour sélectionner des stratégies de réduction de poids à faible dommage sous pression de mémoire.
3. Stratégie d'Exécution : Implémentation du mouvement de page asynchrone pour masquer la latence de déchargement/rechargement avec une perte de débit minimale, couplée à un kernel MoE de précision mixte fusionné.
4. Validation Empirique : Évaluation complète sur trois modèles MoE (Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B) démontrant des performances supérieures aux lignes de base existantes.

Résultats

Les auteurs ont évalué PagedWeight par rapport à FP16, Any-Precision LLM (APL), DP-LLM et MxMoE à travers des tâches de modélisation de langage, de raisonnement et de contexte long.

• Compromis Qualité-Mémoire : PagedWeight atteint une précision équivalente à FP16 avec jusqu'à 72,0 % d'économies de mémoire GPU et une amélioration du débit de 1,94× par rapport aux références.
• Qualité à Budget Fixe : À budgets mémoire similaires, PagedWeight améliore la qualité des tâches par rapport aux méthodes de quantification jusqu'à 39,3 %, avec une perte de débit de 4,1 % maximum.
• Performance en Contexte Long : Sous des budgets de mémoire serrés (ex: 10 Go), PagedWeight maintient une qualité de contexte long (Passage Retrieval, NarrativeQA) comparable à FP16, là où les bases de quantification statique (comme APL) souffrent d'une dégradation significative.
• Débit (Throughput) : PagedWeight égale le débit des bases de quantification uniforme, avec la plus grande baisse observée étant de 4,1 % à un batch size de 4.
• Ablation : Le retrait de composants tels que les statistiques de routage, les résidus de prompt ou le mouvement de page dégrade systématiquement les performances, confirmant la nécessité de la conception complète du système.

Signification

L'article affirme que PagedWeight navigue efficacement dans le compromis complexe entre la précision des tâches du modèle, la consommation de mémoire et le débit/latence. En traitant les poids d'experts comme des unités gérables et paginées plutôt que comme des actifs statiques, PagedWeight débloque un espace de compromis inexploré entre la précision des tâches et l'allocation de la mémoire GPU entre les poids et le cache KV. Le système démontre qu'une quantification dynamique et sensible à la qualité est une stratégie viable et supérieure pour le service MoE par rapport à la quantification statique ou aux méthodes qui ne s'adaptent pas à la pression de la mémoire au moment de l'exécution.