$\textit{Read-ME}$: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-06 23:51:02

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 $\textit{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.
$\textit{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: \url{https://github.com/VITA-Group/READ-ME}.

Main paper: 
  
 Appendix:  


Decision: The manuscript has been reviewed by four reviewers. After the rebuttal, three reviewers gave positive ratings, while one gave borderline rejection.

Overall, the reviewers praised the submission for the motivation, the presentation, and experimental validation. 

The AC agrees with consensus from the reviewers, and believes that the manuscript, despite not flawless, should be presented to a large audience.

Please, however, do account for the comments in the final version.

Congrats!

Reviews: review: paper_summary: paper_summary: The paper proposes an inference-aware method to convert a pretrained dense model into a Mixture-of-Experts architecture, where each expert is smaller than the original dense model. To extract the different experts, they use a dataset from given subdomain to identify the top activated channels. To route among the experts, the authors first break down the limitations of current per-layer routing schemes (waiting for all tokens at the layer to finish, redundancy across layers). They propose to decouple the router from the base model, and instead train a 1-layer transformer block as the router, which predicts token expert assignment autoregressively. 

Overall, these changes allow for inference friendly deployment, by enabling careful batching of examples that share similar experts, and with better caching of experts in resource constrained settings. The authors show that their approach strikes a good trade-off between efficiency and performance.
strength_weakness: strength_weakness: 1.  I am a bit surprised that the gains in latency compared to the seed llama models are somewhat small (19%) given that the number of active parameters is reduced by 30%, and that there is also a decent performance drop from the seed llama model. For resource constrained settings, how would the approach of loading llama layer-by-layer fair against Read-ME ? Like Read-ME, this approach can know in advance what layers to load, and could hopefully retain the performance of the full lama model. 
2. Some parts of the analysis require further clarification (please see questions)
questions: questions: 1. If you keep the top activated neurons, should equation 3 be an argmax ?
2. What is the impact of the routing distillation loss ? Can you provide an ablation experiment where it is not used to access its impact ?
3. How would figure 3 look with READ-ME instead of standard MoEs ? This would be a good visualization to compare the methods. 
4. How were the baseline numbers in Figure 4 obtained ? Did the authors rerun the baselines ? Can you confirm that all these baselines start from the same base model ?
5. The "cost" column of table 1 is misleading; given that your model starts from llama-2, you must either include the compute to create llama-2 in your analysis, or only monitor the additional compute starting from llama 2 (in which case the cost for llama 2 would be 0).
limitations: limitations: N/A
paper_summary: paper_summary: This paper proposes Read-ME, a novel framework for pruning large LLMs into smaller MoEs with minimal training cost. Read-ME separates the gating routers from the critical paths of the inference process and trains an individual expert subnetwork to perform offline pre-gating. With this refactorization of MoEs, the paper further demonstrates the effectiveness of Read-ME by designing optimal expert pre-fetching and caching and expert-aware batching for low-latency and high-throughput serving. Extensive experiments show that Read-ME produces small MoEs that outperform existing small dense models with significantly higher inference performance, lower latency, and higher memory efficiency.
strength_weakness: strength_weakness: - Overall, this is a very interesting work. My biggest question is: what's the motivation behind pruning LLMs into smaller MoEs? If I need a smaller model, why don't I just prune an LLM and have a small but dense model, if the small dense model has approximately the same (or even less) number of activated parameters?

- The observation of Figure 2 is interesting. However, what's the dataset (or input tokens) you use for plotting this figure? Does this observation still hold if the dataset (or input tokens) changes significantly?

- Section 2.3, "The above observations suggest that among many routing paths, only a few of them are in use during the inference..." with only mutual information between adjacent layers may not justify this assumption. It would be better to visualize the end-to-end routing paths (from the first to the last layer) and explicitly show that only a few paths are used frequently.

- The idea of separating routing logic from the inference process is not new. For example, [1] tries to distill the knowledge of gating networks and performs offline routing. How do you compare your expert subnetwork with [1]? Does training the expert subnetwork cost more than KD-based methods?

- To follow up on the previous comment, the evaluation does not compare Read-ME with any existing offline routing approaches like [1]. If the major benefits of Read-ME come from separating routing logics, then the paper's contributions would decrease.

- In evaluation, comparing Read-ME 4.7B with those baselines with fewer parameters may be unfair. Since Read-ME is pruned from much larger LLMs and still holds more activated parameters than baselines, it wouldn't be surprising to see that Read-ME has the best performance. This performance increase comes at the price of requiring more GPU memory. Once again, this question goes back to the motivation: why would someone need to prune LLMs into smaller MoEs instead of smaller dense models? Comparing Read-ME MoEs with smaller dense models pruned from the same LLM may help address this question.

- This pruning process may need a theoretical justification on the performance guarantee, i.e., smaller MoEs are guaranteed to not suffer large performance degradations.


[1] SiDA: Sparsity-Inspired Data-Aware Serving for Efficient and Scalable Large Mixture-of-Experts Model. MLSys'24

Minor issues:

- Figure 2(c) is never mentioned.

- Figure 3 does not have indices for sub-figures.
questions: questions: Pleease refer to the Weaknesses.
limitations: limitations: Pleease refer to the Weaknesses.
paper_summary: paper_summary: The paper proposes a novel framework to enhance the efficiency of pre-trained LLMs by transforming them into post-hoc Mixture-of-Experts (MoE) models. The key innovation lies in decoupling the router from the MoE backbone, which facilitates pre-gating and lookahead scheduling, thereby improving memory management and batching during inference. The proposed method, Read-ME, demonstrates significant improvements in both latency and task performance compared to existing models.
strength_weakness: strength_weakness: 1. The experimental subject is limited. While the paper demonstrates impressive results on the LLaMA-2 model, the scalability and generalization of the proposed method to other model types and larger scales are not extensively explored. It would be essential to include experiments on more diverse datasets and larger models.
2. Although the paper claims minimal overhead for the auto-regressive router, the detailed analysis of its computational costs relative to traditional routers could be expanded. Can you provide a more detailed breakdown of the computational overhead introduced by the auto-regressive router compared to traditional routers? How does this impact overall inference latency, especially in high-throughput scenarios?
3. The paper could benefit from a more detailed explanation of the pre-gating and batching algorithms, potentially with pseudocode or flow diagrams to aid reproducibility. Sometimes I cannot fully follow the authors’ texts.
questions: questions: See weakness. Overall, this seems to be reasonably solid LLM co-design work, but my main concern is its limited evaluation to Llama2 only. Reporting addition experiments on Gemma or Mistral would be strongly encouraged.
limitations: limitations: Limitations are adequately addressed.
paper_summary: paper_summary: This paper proposes to perform data-specific pruning from a regular LLM and, with a small amount of continual training, builds a set of smaller LLMs. Then they are used as experts, and a router network is trained to route requests to them. Note that the routing decision is one per query and not per token or per layer, and hence the proposal is very different from regular MoE. Empirical results are presented with Llama-7B as the starting point and show that the resulting ensemble has good accuracy-latency tradeoff.
strength_weakness: strength_weakness: The main results of Table 1 does not support that the proposed method is at a better accuracy vs inference-cost trade-off. The average accuracy of the proposed method is only marginally better than Sheared-Llama yet its effective inference size of 4.7B is much larger than 2.7B. In the other direction and comparing against the Llama-7B starting point, the overall accuracy cost of pruning and ensemble is quite substantial from 61.4% to 55.5%. 

In additional to the tasks in Table 1, it would help to add perplexity comparison.
questions: questions: Please see above.
limitations: limitations: Yes

Comments: title: Thanks for the rebuttal
Re: rebuttal
Many thanks for raising the score
Response to Authors
comment: Thank you very much for your thorough review and for the thoughtful comments you provided on our work. 

We have made the necessary revisions in response to your comments, and we hope that our revisions meet your expectations. 
If there are any remaining concerns or if you have any further suggestions that could help us improve the quality of our work,  please do not hesitate to let us know. We would greatly appreciate any further guidance you can provide. 

Once again, we are truly grateful for your support.
Authors
I thank the authors for the discussion and has raised my rating from 3 to 4.
Many thanks for your thoughtful comments and suggestions. We value your support and will be incorporating all your feedback in our revised version
Thank you once again for your insightful comments and suggestions. We greatly appreciate your support of our work and will be incorporating all the feedback into our revised version.
The authors have provided a very clear rebuttal and detailed answers to my initial set of questions. I will change my score accordingly
I appreciate the thorough answers from the authors. The rebuttal has addressed all my questions, provided interesting new results, and revealed the novelty of this paper.

In response to the rebuttal, I would like to raise my rating to help advance the possibility of acceptance for this paper.
Thank you very much for your insightful suggestions, which have been greatly enlightening and are crucial for enhancing the quality of our paper! We will adhere to these suggestions in the final version and also revise the paper according to all other comments.
I thank the authors for their rebuttal and supplemented experiments. I appreciate the reply, and I found my concerns adequately addressed. I will raise the score to 7 in response to the authors' rebuttal.