PreMoe : Alléger les MoEs sur mémoire contrainte par élagage et récupération d'experts

PreMoe: Lightening MoEs on Constrained Memory by Expert Pruning and Retrieval

May 23, 2025

Auteurs: Zehua Pei, Ying Zhang, Hui-Ling Zhen, Xianzhi Yu, Wulong Liu, Sinno Jialin Pan, Mingxuan Yuan, Bei Yu

cs.AI

Résumé

Les architectures de type *mixture-of-experts* (MoE) permettent de mettre à l'échelle les grands modèles de langage (LLM) avec un nombre considérable de paramètres sans augmentation proportionnelle des coûts de calcul. Cependant, les exigences mémoire importantes des grands modèles MoE entravent leur déploiement dans divers environnements informatiques, des serveurs cloud aux appareils grand public. Cette étude démontre d'abord une spécialisation marquée des modèles d'activation des experts en fonction des tâches au sein des couches MoE. Sur cette base, nous introduisons PreMoe, un nouveau cadre permettant un déploiement efficace des modèles MoE massifs dans des environnements à mémoire limitée. PreMoe comprend deux composants principaux : l'élagage probabiliste des experts (PEP) et la récupération adaptative des experts en fonction de la tâche (TAER). Le PEP utilise une nouvelle métrique, le *task-conditioned expected selection score* (TCESS), dérivé des logits du routeur pour quantifier l'importance des experts pour des tâches spécifiques, identifiant ainsi un ensemble minimal d'experts critiques. Le TAER exploite ces profils d'importance des experts spécifiques à la tâche pour une inférence efficace. Il précalcule et stocke des modèles d'experts compacts pour diverses tâches. Lorsqu'une requête utilisateur est reçue, le TAER identifie rapidement le modèle de tâche stocké le plus pertinent et reconstruit le modèle en chargeant uniquement le petit sous-ensemble d'experts essentiels pour cette tâche. Cette approche réduit considérablement l'empreinte mémoire dans tous les scénarios de déploiement. DeepSeek-R1 671B maintient une précision de 97,2 % sur MATH500 lorsqu'il est élagué à une configuration 8/128 (réduction de 50 % des experts) et atteint encore 72,0 % avec un élagage agressif 8/32 (réduction de 87,5 % des experts). Pangu-Ultra-MoE 718B atteint 97,15 % sur MATH500 et 81,3 % sur AIME24 avec un élagage 8/128, tandis qu'un élagage encore plus agressif à 4/64 (390 Go de mémoire) préserve une précision de 96,95 % sur MATH500. Nous mettons notre code à disposition à l'adresse suivante : https://github.com/JarvisPei/PreMoe.

English

Mixture-of-experts (MoE) architectures enable scaling large language models (LLMs) to vast parameter counts without a proportional rise in computational costs. However, the significant memory demands of large MoE models hinder their deployment across various computational environments, from cloud servers to consumer devices. This study first demonstrates pronounced task-specific specialization in expert activation patterns within MoE layers. Building on this, we introduce PreMoe, a novel framework that enables efficient deployment of massive MoE models in memory-constrained environments. PreMoe features two main components: probabilistic expert pruning (PEP) and task-adaptive expert retrieval (TAER). PEP employs a new metric, the task-conditioned expected selection score (TCESS), derived from router logits to quantify expert importance for specific tasks, thereby identifying a minimal set of critical experts. TAER leverages these task-specific expert importance profiles for efficient inference. It pre-computes and stores compact expert patterns for diverse tasks. When a user query is received, TAER rapidly identifies the most relevant stored task pattern and reconstructs the model by loading only the small subset of experts crucial for that task. This approach dramatically reduces the memory footprint across all deployment scenarios. DeepSeek-R1 671B maintains 97.2\% accuracy on MATH500 when pruned to 8/128 configuration (50\% expert reduction), and still achieves 72.0\% with aggressive 8/32 pruning (87.5\% expert reduction). Pangu-Ultra-MoE 718B achieves 97.15\% on MATH500 and 81.3\% on AIME24 with 8/128 pruning, while even more aggressive pruning to 4/64 (390GB memory) preserves 96.95\% accuracy on MATH500. We make our code publicly available at https://github.com/JarvisPei/PreMoe.