PreMoe: Облегчение моделей с экспертами при ограниченной памяти за счет обрезки и извлечения экспертов

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

May 23, 2025

Авторы: Zehua Pei, Ying Zhang, Hui-Ling Zhen, Xianzhi Yu, Wulong Liu, Sinno Jialin Pan, Mingxuan Yuan, Bei Yu

cs.AI

Аннотация

Архитектуры типа "смесь экспертов" (Mixture-of-Experts, MoE) позволяют масштабировать большие языковые модели (LLM) до огромного количества параметров без пропорционального увеличения вычислительных затрат. Однако значительные требования к памяти крупных моделей MoE затрудняют их развертывание в различных вычислительных средах, от облачных серверов до потребительских устройств. В данном исследовании сначала демонстрируется выраженная специализация активации экспертов в слоях MoE в зависимости от конкретных задач. На основе этого мы представляем PreMoe — новый фреймворк, который обеспечивает эффективное развертывание крупных моделей MoE в средах с ограниченной памятью. PreMoe включает два основных компонента: вероятностное сокращение экспертов (Probabilistic Expert Pruning, PEP) и адаптивное извлечение экспертов для задач (Task-Adaptive Expert Retrieval, TAER). PEP использует новый показатель — ожидаемый балл выбора, обусловленный задачей (Task-Conditioned Expected Selection Score, TCESS), который вычисляется на основе логитов маршрутизатора для количественной оценки важности экспертов для конкретных задач, что позволяет определить минимальный набор критически важных экспертов. TAER использует эти профили важности экспертов для эффективного вывода. Он предварительно вычисляет и сохраняет компактные шаблоны экспертов для различных задач. Когда поступает запрос пользователя, TAER быстро идентифицирует наиболее подходящий сохраненный шаблон задачи и восстанавливает модель, загружая только небольшое подмножество экспертов, критически важных для этой задачи. Такой подход значительно сокращает объем используемой памяти во всех сценариях развертывания. Модель DeepSeek-R1 671B сохраняет точность 97,2% на тесте MATH500 при сокращении до конфигурации 8/128 (сокращение экспертов на 50%) и достигает 72,0% при агрессивном сокращении до 8/32 (сокращение экспертов на 87,5%). Модель Pangu-Ultra-MoE 718B показывает точность 97,15% на MATH500 и 81,3% на AIME24 при сокращении до 8/128, а еще более агрессивное сокращение до 4/64 (390 ГБ памяти) сохраняет точность 96,95% на MATH500. Мы делаем наш код общедоступным по адресу 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.