PreMoe: Entlastung von MoEs bei begrenztem Speicher durch Expertenausdünnung und -abruf

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

May 23, 2025

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

cs.AI

Zusammenfassung

Mixture-of-Experts (MoE)-Architekturen ermöglichen die Skalierung großer Sprachmodelle (LLMs) auf enorme Parameterzahlen ohne einen proportionalen Anstieg der Rechenkosten. Die erheblichen Speicheranforderungen großer MoE-Modelle behindern jedoch deren Einsatz in verschiedenen Rechenumgebungen, von Cloud-Servern bis hin zu Endgeräten. Diese Studie zeigt zunächst eine ausgeprägte aufgabenbezogene Spezialisierung in den Aktivierungsmustern der Experten innerhalb der MoE-Schichten. Aufbauend darauf stellen wir PreMoe vor, ein neuartiges Framework, das den effizienten Einsatz massiver MoE-Modelle in speicherbeschränkten Umgebungen ermöglicht. PreMoe umfasst zwei Hauptkomponenten: probabilistisches Experten-Pruning (PEP) und aufgabenadaptives Experten-Retrieval (TAER). PEP verwendet eine neue Metrik, den aufgabenbedingten erwarteten Auswahlwert (TCESS), der aus den Router-Logits abgeleitet wird, um die Bedeutung der Experten für spezifische Aufgaben zu quantifizieren und so eine minimale Menge kritischer Experten zu identifizieren. TAER nutzt diese aufgabenspezifischen Expertenbedeutungsprofile für effizientes Inferenzverfahren. Es berechnet und speichert kompakte Expertenmuster für verschiedene Aufgaben im Voraus. Wenn eine Benutzeranfrage eingeht, identifiziert TAER schnell das relevanteste gespeicherte Aufgabenmuster und rekonstruiert das Modell, indem nur die kleine Teilmenge der Experten geladen wird, die für diese Aufgabe entscheidend sind. Dieser Ansatz reduziert den Speicherbedarf in allen Einsatzszenarien drastisch. DeepSeek-R1 671B behält 97,2 % Genauigkeit auf MATH500 bei einer Pruning-Konfiguration von 8/128 (50 % Expertenreduktion) und erreicht sogar 72,0 % bei aggressivem 8/32-Pruning (87,5 % Expertenreduktion). Pangu-Ultra-MoE 718B erreicht 97,15 % auf MATH500 und 81,3 % auf AIME24 mit 8/128-Pruning, während noch aggressiveres Pruning auf 4/64 (390 GB Speicher) eine Genauigkeit von 96,95 % auf MATH500 bewahrt. Unser Code ist öffentlich verfügbar unter 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.