Leia-ME: Refatoração de LLMs como Mistura de Especialistas Desacoplados de Roteador com Co-Design 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

Resumo

A proliferação de grandes modelos de linguagem (LLMs) levou à adoção de arquiteturas Mixture-of-Experts (MoE) que dinamicamente aproveitam sub-redes especializadas para melhorar a eficiência e o desempenho. Apesar de seus benefícios, os modelos MoE enfrentam desafios significativos durante a inferência, incluindo gerenciamento ineficiente de memória e agrupamento subótimo, devido a escolhas de design desalinhadas entre a arquitetura do modelo e as políticas do sistema. Além disso, a abordagem convencional de treinar MoEs do zero está se tornando cada vez mais proibitiva em termos de custo. Neste artigo, propomos um novo framework denominado Read-ME que transforma LLMs densos pré-treinados em modelos MoE menores (em contraste com a "reciclagem" de MoEs generalistas), evitando os altos custos do treinamento do zero. Nossa abordagem emprega esparsidade de ativação para extrair especialistas. Para compor os especialistas, examinamos o design de roteador em camadas amplamente adotado e mostramos sua redundância, introduzindo assim o roteador de pré-gateamento desacoplado do espinha dorsal do MoE que facilita a pré-computação amigável ao sistema e o agendamento de lookahead, aprimorando o agrupamento e o armazenamento em cache conscientes dos especialistas. Nosso codesign aborda lacunas críticas tanto nos aspectos algorítmicos quanto nos sistemas, estabelecendo uma alternativa escalável e eficiente para inferência de LLM em ambientes com recursos limitados. O Read-ME supera outros modelos densos de código aberto populares de escalas semelhantes, alcançando melhorias de até 10,1% no MMLU e melhorando a latência média de ponta a ponta em até 6,1%. Os códigos estão disponíveis em: 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.