Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with
  System Co-Design
Authors: name: Ruisi Cai; name: Yeonju Ro; name: Geon-Woo Kim; name: Peihao Wang; name: Babak Ehteshami Bejnordi; name: Aditya Akella; name: Zhangyang Wang
Published: 2024-10-24 19:48:51

Abstract: 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.