Lea-ME: Refactorización de LLMs como Mezcla de Expertos Desacoplada de Enrutador con Diseño de Sistema

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

October 24, 2024

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

cs.AI

Resumen

La proliferación de grandes modelos de lenguaje (LLMs) ha llevado a la adopción de arquitecturas de Mezcla de Expertos (MoE) que aprovechan dinámicamente subredes especializadas para mejorar la eficiencia y el rendimiento. A pesar de sus beneficios, los modelos MoE enfrentan desafíos significativos durante la inferencia, incluida una gestión ineficiente de la memoria y lotificación subóptima, debido a decisiones de diseño desalineadas entre la arquitectura del modelo y las políticas del sistema. Además, el enfoque convencional de entrenar MoEs desde cero es cada vez más prohibitivo en términos de costos. En este documento, proponemos un nuevo marco, Read-ME, que transforma LLMs densos pre-entrenados en modelos MoE más pequeños (en contraste con "reciclar" MoEs generalistas), evitando los altos costos del entrenamiento desde cero. Nuestro enfoque emplea la dispersión de activación para extraer expertos. Para componer expertos, examinamos el diseño de enrutador por capas ampliamente adoptado y mostramos su redundancia, e introducimos el enrutador de pre-gating desacoplado de la estructura central de MoE que facilita la precomputación amigable para el sistema y la programación de anticipación, mejorando la lotificación y el almacenamiento en caché conscientes de los expertos. Nuestro diseño conjunto aborda brechas críticas tanto en los aspectos algorítmicos como en los del sistema, estableciendo una alternativa escalable y eficiente para la inferencia de LLM en entornos con recursos limitados. Read-ME supera a otros modelos densos de código abierto populares de escalas similares, logrando mejoras de hasta un 10.1% en MMLU y mejorando la latencia media de extremo a extremo hasta un 6.1%. Los códigos están disponibles en: 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.