Lisez-MOI : Refactorisation des LLM en Mélange d'Experts Découplés du Routeur avec Conception Système

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

October 24, 2024

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

cs.AI

Résumé

La prolifération des grands modèles de langage (LLM) a conduit à l'adoption d'architectures Mixture-of-Experts (MoE) qui exploitent dynamiquement des sous-réseaux spécialisés pour une efficacité et des performances améliorées. Malgré leurs avantages, les modèles MoE rencontrent des défis significatifs lors de l'inférence, notamment une gestion inefficace de la mémoire et un regroupement sous-optimal, en raison de choix de conception mal alignés entre l'architecture du modèle et les politiques du système. De plus, l'approche conventionnelle consistant à entraîner les MoE à partir de zéro devient de plus en plus prohibitive en termes de coût. Dans cet article, nous proposons un nouveau cadre, Read-ME, qui transforme les LLM denses pré-entraînés en modèles MoE plus petits (par opposition à la "valorisation" des MoE généralistes), évitant les coûts élevés de l'entraînement à partir de zéro. Notre approche utilise la sparsité de l'activation pour extraire des experts. Pour composer les experts, nous examinons la conception de routeur couche par couche largement adoptée et montrons sa redondance, et introduisons donc le routeur de pré-porte découplé de l'épine dorsale MoE qui facilite la pré-computation conviviale du système et la planification anticipée, améliorant le regroupement et la mise en cache conscients des experts. Notre coconception aborde donc les lacunes critiques à la fois sur les plans algorithmique et système, établissant une alternative évolutive et efficace pour l'inférence LLM dans des environnements contraints en ressources. Read-ME surpasse d'autres modèles denses open-source populaires de tailles similaires, obtenant des améliorations allant jusqu'à 10,1% sur MMLU, et améliorant la latence moyenne de bout en bout jusqu'à 6,1%. Les codes sont disponibles sur : 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.