Lees-ME: Het herstructureren van LLM's als een Router-Ontkoppelde Mix van Experts met Systeem Co-Design

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

Samenvatting

De proliferatie van grote taalmodellen (LLM's) heeft geleid tot de adoptie van Mixture-of-Experts (MoE) architecturen die dynamisch gespecialiseerde subnetswerken benutten voor verbeterde efficiëntie en prestaties. Ondanks hun voordelen ondervinden MoE-modellen aanzienlijke uitdagingen tijdens inferentie, waaronder inefficiënt geheugenbeheer en suboptimale batchverwerking, als gevolg van niet-uitgelijnde ontwerpkeuzes tussen de modelarchitectuur en de systeembeleidsregels. Bovendien wordt de conventionele aanpak van het trainen van MoE's vanaf nul steeds meer belemmerd door de kosten. In dit artikel stellen we een nieuw framework, Read-ME, voor dat vooraf getrainde dichte LLM's transformeert naar kleinere MoE-modellen (in tegenstelling tot het "upcyclen" van generalistische MoE's), waarbij de hoge kosten vanaf de basis trainen worden vermeden. Onze aanpak maakt gebruik van activatieschaarste om experts te extraheren. Om experts samen te stellen, onderzoeken we het veelgebruikte laagsgewijze routerontwerp en tonen we de redundantie ervan aan, en introduceren we daarom de voorpoortrouter die losgekoppeld is van de MoE-ruggengraat en die systeemvriendelijk voorberekenen en lookahead-planning vergemakkelijkt, waardoor expert-bewuste batchverwerking en caching worden verbeterd. Onze codesign pakt daarom kritieke hiaten aan op zowel het algoritmische als het systeemfront, en biedt een schaalbaar en efficiënt alternatief voor LLM-inferentie in omgevingen met beperkte middelen. Read-ME presteert beter dan andere populaire open-source dichte modellen van vergelijkbare omvang, met verbeteringen tot wel 10,1% op MMLU en een verbetering van de gemiddelde end-to-end latentie tot wel 6,1%. De codes zijn beschikbaar op: 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.