PreMoe: Ottimizzazione delle MoE su Memoria Limitata mediante Potatura e Recupero degli Esperti

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

May 23, 2025

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

cs.AI

Abstract

Le architetture Mixture-of-Experts (MoE) consentono di scalare modelli linguistici di grandi dimensioni (LLM) a conteggi di parametri elevati senza un aumento proporzionale dei costi computazionali. Tuttavia, le significative esigenze di memoria dei grandi modelli MoE ostacolano il loro dispiegamento in vari ambienti computazionali, dai server cloud ai dispositivi consumer. Questo studio dimostra innanzitutto una marcata specializzazione specifica per task nei modelli di attivazione degli esperti all'interno degli strati MoE. Basandoci su questo, introduciamo PreMoe, un nuovo framework che consente il dispiegamento efficiente di modelli MoE di grandi dimensioni in ambienti con memoria limitata. PreMoe presenta due componenti principali: il pruning probabilistico degli esperti (PEP) e il recupero adattivo degli esperti in base al task (TAER). Il PEP utilizza una nuova metrica, il punteggio di selezione atteso condizionato al task (TCESS), derivato dai logit del router per quantificare l'importanza degli esperti per task specifici, identificando così un insieme minimo di esperti critici. Il TAER sfrutta questi profili di importanza degli esperti specifici per task per un'inferenza efficiente. Pre-calcola e memorizza modelli compatti di esperti per task diversi. Quando viene ricevuta una query dell'utente, il TAER identifica rapidamente il modello di task memorizzato più rilevante e ricostruisce il modello caricando solo il piccolo sottoinsieme di esperti cruciali per quel task. Questo approccio riduce drasticamente l'impronta di memoria in tutti gli scenari di dispiegamento. DeepSeek-R1 671B mantiene il 97,2\% di accuratezza su MATH500 quando viene potato alla configurazione 8/128 (riduzione del 50\% degli esperti), e raggiunge ancora il 72,0\% con un pruning aggressivo 8/32 (riduzione dell'87,5\% degli esperti). Pangu-Ultra-MoE 718B raggiunge il 97,15\% su MATH500 e l'81,3\% su AIME24 con il pruning 8/128, mentre un pruning ancora più aggressivo a 4/64 (390GB di memoria) preserva il 96,95\% di accuratezza su MATH500. Rendiamo il nostro codice disponibile pubblicamente all'indirizzo 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.