Anleitung: Refaktorisierung von LLMs als Router-entkoppelte Mischung von Experten mit System-Co-Design

Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design

October 24, 2024

Autoren: Ruisi Cai, Yeonju Ro, Geon-Woo Kim, Peihao Wang, Babak Ehteshami Bejnordi, Aditya Akella, Zhangyang Wang

cs.AI

Zusammenfassung

Die Verbreitung großer Sprachmodelle (LLMs) hat zur Übernahme von Mixture-of-Experts (MoE)-Architekturen geführt, die spezialisierte Teilnetzwerke dynamisch nutzen, um Effizienz und Leistung zu verbessern. Trotz ihrer Vorteile stehen MoE-Modelle vor erheblichen Herausforderungen während der Inferenz, darunter ineffizientes Speichermanagement und suboptimales Batching aufgrund von nicht übereinstimmenden Designentscheidungen zwischen der Modellarchitektur und den Systemrichtlinien. Darüber hinaus wird der konventionelle Ansatz, MoEs von Grund auf zu trainieren, zunehmend kostspielig. In diesem Papier schlagen wir ein neuartiges Framework namens Read-ME vor, das vortrainierte dichte LLMs in kleinere MoE-Modelle umwandelt (im Gegensatz zur "Aufwertung" von Generalist-MoEs), um die hohen Kosten des Neutraining zu vermeiden. Unser Ansatz nutzt Aktivitätssparsamkeit zur Extraktion von Experten. Um Experten zu komponieren, untersuchen wir das weit verbreitete schichtweise Router-Design und zeigen dessen Redundanz auf. Daher führen wir den vorgeschalteten Router ein, der vom MoE-Grundgerüst entkoppelt ist und die systemfreundliche Vorberechnung und Vorausschauzeitplanung ermöglicht, um Expertenbewusstes Batching und Caching zu verbessern. Unsere Codesign-Strategie adressiert somit kritische Lücken sowohl auf algorithmischer als auch auf systemischer Ebene und etabliert eine skalierbare und effiziente Alternative für die Inferenz von LLMs in ressourcenbeschränkten Umgebungen. Read-ME übertrifft andere beliebte Open-Source-dichte Modelle ähnlicher Größenordnung und erzielt Verbesserungen von bis zu 10,1% bei MMLU und verbessert die mittlere End-to-End-Latenz um bis zu 6,1%. Die Codes sind verfügbar unter: https://github.com/VITA-Group/READ-ME.

English

The proliferation of large language models (LLMs) has led to the adoption of Mixture-of-Experts (MoE) architectures that dynamically leverage specialized subnetworks for improved efficiency and performance. Despite their benefits, MoE models face significant challenges during inference, including inefficient memory management and suboptimal batching, due to misaligned design choices between the model architecture and the system policies. Furthermore, the conventional approach of training MoEs from scratch is increasingly prohibitive in terms of cost. In this paper, we propose a novel framework Read-ME that transforms pre-trained dense LLMs into smaller MoE models (in contrast to "upcycling" generalist MoEs), avoiding the high costs of ground-up training. Our approach employs activation sparsity to extract experts. To compose experts, we examine the widely-adopted layer-wise router design and show its redundancy, and thus we introduce the pre-gating router decoupled from the MoE backbone that facilitates system-friendly pre-computing and lookahead scheduling, enhancing expert-aware batching and caching. Our codesign therefore addresses critical gaps on both the algorithmic and system fronts, establishing a scalable and efficient alternative for LLM inference in resource-constrained settings. Read-ME outperforms other popular open-source dense models of similar scales, achieving improvements of up to 10.1% on MMLU, and improving mean end-to-end latency up to 6.1%. Codes are available at: https://github.com/VITA-Group/READ-ME.