PreMoe: Aligeramiento de MoEs en memoria restringida mediante poda y recuperación de expertos

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

May 23, 2025

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

cs.AI

Resumen

Las arquitecturas de mezcla de expertos (MoE, por sus siglas en inglés) permiten escalar modelos de lenguaje grandes (LLMs) a recuentos de parámetros masivos sin un aumento proporcional en los costos computacionales. Sin embargo, las demandas significativas de memoria de los modelos MoE grandes dificultan su implementación en diversos entornos computacionales, desde servidores en la nube hasta dispositivos de consumo. Este estudio demuestra primero una especialización pronunciada en los patrones de activación de expertos dentro de las capas MoE, específica para cada tarea. Basándonos en esto, presentamos PreMoe, un marco novedoso que permite la implementación eficiente de modelos MoE masivos en entornos con limitaciones de memoria. PreMoe incluye dos componentes principales: poda probabilística de expertos (PEP) y recuperación de expertos adaptativa a la tarea (TAER). PEP emplea una nueva métrica, la puntuación de selección esperada condicionada a la tarea (TCESS), derivada de los logits del enrutador para cuantificar la importancia de los expertos en tareas específicas, identificando así un conjunto mínimo de expertos críticos. TAER aprovecha estos perfiles de importancia de expertos específicos para la tarea para realizar inferencias eficientes. Precalcula y almacena patrones compactos de expertos para diversas tareas. Cuando se recibe una consulta del usuario, TAER identifica rápidamente el patrón de tarea almacenado más relevante y reconstruye el modelo cargando solo el pequeño subconjunto de expertos cruciales para esa tarea. Este enfoque reduce drásticamente la huella de memoria en todos los escenarios de implementación. DeepSeek-R1 671B mantiene un 97.2\% de precisión en MATH500 cuando se poda a una configuración de 8/128 (reducción del 50\% de expertos), y aún logra un 72.0\% con una poda agresiva de 8/32 (reducción del 87.5\% de expertos). Pangu-Ultra-MoE 718B alcanza un 97.15\% en MATH500 y un 81.3\% en AIME24 con una poda de 8/128, mientras que una poda aún más agresiva a 4/64 (390GB de memoria) preserva un 96.95\% de precisión en MATH500. Hacemos nuestro código disponible públicamente en 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.