PreMoe: Het verlichten van Mixture of Experts-modellen op beperkt geheugen door expert pruning en retrieval

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

Samenvatting

Mixture-of-experts (MoE)-architecturen maken het mogelijk om grote taalmodel(len) (LLMs) op te schalen naar enorme parameteraantallen zonder een proportionele stijging in rekencosten. De aanzienlijke geheugeneisen van grote MoE-modellen belemmeren echter hun implementatie in diverse rekenomgevingen, van cloudservers tot consumentenapparaten. Deze studie toont eerst een duidelijke taakspecifieke specialisatie aan in expertactiveringspatronen binnen MoE-lagen. Hierop voortbouwend introduceren we PreMoe, een nieuw framework dat efficiënte implementatie van massieve MoE-modellen in geheugenbeperkte omgevingen mogelijk maakt. PreMoe bevat twee hoofdcomponenten: probabilistische expertpruning (PEP) en taakadaptieve expertretrieval (TAER). PEP gebruikt een nieuwe metriek, de taakgeconditioneerde verwachte selectiescore (TCESS), afgeleid van router-logits om het belang van experts voor specifieke taken te kwantificeren, waardoor een minimale set kritieke experts wordt geïdentificeerd. TAER benut deze taakspecifieke expertbelangprofielen voor efficiënte inferentie. Het berekent en slaat compacte expertpatronen voor diverse taken vooraf op. Wanneer een gebruikersquery wordt ontvangen, identificeert TAER snel het meest relevante opgeslagen taakpatroon en reconstrueert het model door alleen de kleine subset van experts te laden die cruciaal zijn voor die taak. Deze aanpak vermindert het geheugengebruik aanzienlijk in alle implementatiescenario's. DeepSeek-R1 671B behoudt 97,2% nauwkeurigheid op MATH500 wanneer het wordt geprunt naar een 8/128-configuratie (50% expertreductie), en behaalt nog steeds 72,0% met agressieve 8/32-pruning (87,5% expertreductie). Pangu-Ultra-MoE 718B behaalt 97,15% op MATH500 en 81,3% op AIME24 met 8/128-pruning, terwijl zelfs nog agressievere pruning naar 4/64 (390GB geheugen) 96,95% nauwkeurigheid op MATH500 behoudt. We maken onze code publiekelijk beschikbaar op 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.