PreMoe: Otimizando MoEs em Memória Limitada por meio de Poda e Recuperação de Especialistas

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

Resumo

Arquiteturas de mistura de especialistas (MoE) permitem escalar modelos de linguagem grandes (LLMs) para contagens vastas de parâmetros sem um aumento proporcional nos custos computacionais. No entanto, as demandas significativas de memória de grandes modelos MoE dificultam sua implantação em diversos ambientes computacionais, desde servidores em nuvem até dispositivos de consumo. Este estudo primeiro demonstra uma especialização acentuada e específica de tarefas nos padrões de ativação de especialistas dentro das camadas MoE. Com base nisso, introduzimos o PreMoe, uma nova estrutura que possibilita a implantação eficiente de modelos MoE massivos em ambientes com restrições de memória. O PreMoe apresenta dois componentes principais: poda probabilística de especialistas (PEP) e recuperação de especialistas adaptativa à tarefa (TAER). O PEP emprega uma nova métrica, o escore esperado de seleção condicionado à tarefa (TCESS), derivado dos logits do roteador para quantificar a importância dos especialistas para tarefas específicas, identificando assim um conjunto mínimo de especialistas críticos. O TAER aproveita esses perfis de importância de especialistas específicos da tarefa para inferência eficiente. Ele pré-computa e armazena padrões compactos de especialistas para diversas tarefas. Quando uma consulta do usuário é recebida, o TAER identifica rapidamente o padrão de tarefa armazenado mais relevante e reconstrói o modelo carregando apenas o pequeno subconjunto de especialistas crucial para essa tarefa. Essa abordagem reduz drasticamente a pegada de memória em todos os cenários de implantação. O DeepSeek-R1 671B mantém 97,2% de precisão no MATH500 quando podado para a configuração 8/128 (redução de 50% nos especialistas), e ainda alcança 72,0% com uma poda agressiva de 8/32 (redução de 87,5% nos especialistas). O Pangu-Ultra-MoE 718B atinge 97,15% no MATH500 e 81,3% no AIME24 com poda 8/128, enquanto uma poda ainda mais agressiva para 4/64 (390GB de memória) preserva 96,95% de precisão no MATH500. Disponibilizamos nosso código publicamente em 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.