MoE-ERAS: Expert Residency Aware Selection

Authors: abhimanyu  bambhaniya
sashankh  chengavalli kumar
tushar  krishna
Venue: iscaconf_org_ISCA_2024_Workshop_MLArchSys
Published: 2024-05-14 18:01:22

Abstract: Mixture of Experts models have quickly grown in popularity due to their faster inference and training than dense models of similar capability. Parameter compression and offloading allow the users to run these models on smaller GPU memory (leading to cost savings). However, unpredictability in expert activation results in slower inference for offloaded experts. In this work, we profile and study the expert activation patterns when running large MoE models. Based on insights from activation patterns, we propose a new way of expert selection, which takes the expert residency into account. We introduce \textit{MoE-ERAS}, Expert Residency Aware Selection to select the most suitable experts considering \textbf{both performance and accuracy}. We show substantial gains in decoding latency and expert swaps, and present analysis to show pre-fetching opportunities for future work. MoE-ERAS allows users to choose an acceptable point on the speedup-quality trade-off.

Main paper: 
 i. introduction: 
 Numerous recent advancements in natural language processing hinge on large pre-trained language models, exemplified by entities like GPT-3 and 4 [1] , [4] , Palm & Gemini [5] , [7] , among others. However, the swift strides in this domain owe much to openly accessible LLMs like LLaMA [17] , Mixtral [12] , OPT [21] , and many more. The primary boon of these open-access LLMs lies in researchers' ability to deploy and tailor them locally -a feat infeasible with proprietary APIs. With LLM models growing at an unprecedented scale, a Mixture of expert(MoE) models is one of the most promising directions that help us scale models to larger dimensions. The scale of a model is one of the most important axes for better model quality. Given a fixed computing budget, it is generally considered better to train a larger model for fewer steps, than to train a smaller model for more steps.Despite the openly available range of LLMs, their sizes poses a challenge to utilization. Cutting-edge open-access language models demand multiple high-end GPUs even for rudimentary inference tasks. To make these LLMs feasible on more economical hardware setups, practitioners must either reduce model parameters or transfer parameters to less expensive but slower storage mediums, whether RAM or SSD [2] , [14] .When moving the parameters to slower storage, the challenge of unreliable expert activation during the inference stage has emerged as a crucial issue in the development and optimization of MoE models. The gating mechanism determines which expert to activate for a given input just before the expert layer. This leads to sub-optimal model throughput as the appropriate weights would have to be brought from the host to the device, if not already on-chip.We aim to increase the reuseblity of experts that are already present on the HBM rather than bring new experts each time from host memory. Our key insight is that activating good enough experts can help in running the models faster and could also help us gain significant throughput improvements. Building on this insight we introduce, MoE-ERAS, Expert Residency Aware Selection. MoE-ERAS is an expert selection method that takes the expert residency factor into consideration when selecting experts for each token. We propose 2 techniques, thresholding and biasing to tweak the router (gating network) for selecting throughput favourable experts.Our approach is orthogonal to past works which aim to speedup MoE inference through quantization, prefetching and cacheing. To prove this, we implement MoE-ERAS on top of techniques like quantization and caching, and still show significant speedups with minimal accuracy degradation.Our contributions in this work are :• Profile and analyze the expert activation patterns for Mixtral-8x7B and Switch Transformer-32E. We collect over 500k token samples and provide insights for expert activation profiles. Mixture-of-experts models have emerged as a powerful approach for scaling up deep learning models to handle complex tasks with high-dimensional data. By dividing the computational workload among multiple expert sub-networks, each specializing in a different aspect of the input, MoE models can achieve high representational capacity while maintaining computational efficiency during training and inference. However, despite their promising performance, MoE models still face challenges regarding inference time, which can significantly impact their practical deployment, especially in latency-sensitive applications.A. Gating in Mixture-of-Experts Fig. 1 : The gating network of a mixture-of-experts model.In mixture-of-experts models, the gating mechanism is crucial in selectively activating the appropriate expert subnetworks for a given input. The gating network is a separate component that takes the input data and produces a set of soft assignment scores or gating values, one for each expert. An example of gating is shown in 1. These gating values represent each expert's relative importance or suitability in handling the input data. The gating network is typically a single, fully connected network trained jointly with the experts during the model's training phase.The gating values are then used to compute a weighted combination of the outputs from the individual selected experts, effectively creating a mixture of expert predictions. The gating mechanism allows the MoE model to dynamically allocate computational resources to the most relevant experts, enabling efficient processing of diverse input data while maintaining high predictive performance.MoE serving can be divided into 2 baskets: either we can have enough devices (GPUs, TPUs, etc) to store all parameters (weights+KV Cache) inside on-chip HBM memory, or we have fewer devices and offload some of the unused parameters to the slower host memory(CPU DRAM). For the first approach, more devices are required; thus, we would have a higher cost of serving.The second approach and main use case, which our work targets, is to use fewer devices and offload some parameters to the host memory, resulting in cheaper inference. However, offloading the experts to the host device comes at the cost of slower inference than the previous approach. Techniques like parameter compression [3] , [6] , [9] , [11] , [20] , expert activation prediction [18] , [20] , and caching [6] are used to mitigate runtime degradations.While these techniques can help speed up the MoE inference, these alone are insufficient solutions. We profile the time of reading experts from the host CPU vs from the H100 GPU. 2 compares the read time GPU and CPU for different expert sizes. We can clearly see the CPU read time is orders of magnitude greater than the GPU read time. This would mean prefetching experts with 1-2 layers would not result in any meaningful speedup of modern GPUs. Hence, we aim to increase the use of experts already present in the HBM memory. In this section, we first present the analysis of the expert activation patterns in Mixtral-8x7b and Switch Transformer-32 MoE models during inference. By examining the predictability and disparity in the activation patterns in this model, we present an analysis that motivates the design of MoE-ERAS. We then describe two routing schemes proposed and tested in this work -Thresholding and Biasing.Mixtral-8x7b contains 32 hidden layers, with each containing 8 experts. During the generation of each token, k of these experts (default k = 2) are selected by the gating network (or router) immediately prior to the layer. These experts are said to be activated. We hypothesize that:• (a) despite the use of a load balancing loss during training, the activation of experts will be uneven within each layer, creating "hot" and "cold" experts that are often and rarely activated, respectively. • (b) There exist expert-expert activation correlations between the different layers in the model, mainly between the early and late layers. • (c) given the experts activated in earlier layers, it is possible to predict with reasonable accuracy the experts activated in later layers. To test these hypotheses, we profile the text generation task on a corpus [16] . Inherently, the support for these hypotheses shows opportunities for optimizing the inference latency and throughput of MoE models. The support we find for these hypotheses helps us design the interventions described in Section 3.2 to the gating mechanism to speed up inference. While our work focuses on running state-of-the-art MoE models on commodity hardware where batch size = 1 is acceptable, we also examine how expert activation varies with batch size in Switch Transformer-32.The behavior of the gating network in standard Mixtral-8x7b is shown below. The output of self-attention is passed through a dense network, which gives logit values for each expert. Softmax function is applied to these logits to convert them into probabilities. The Top-K method is used to select the experts to activate, where k is a parameter derived from the model configuration. In Mixtral-8x7b, this defaults to k = 2.EQUATIONOur motivation results from section II show that offloading experts to the host can significantly affect inference latency. MoE Offloading [6] selects experts to cache and offload, making it possible to run Mixtral-8x7b on a Tesla T4 with 16GB VRAM. Our examination of the gating network's outputs shows that there are two key ways in which we can make our gating flexible -selecting "good-enough" experts and "biasing" towards those already on-chip.1) Thresholding: Early analysis of the gating network's output logits showed us that there might not always be "clear winners" when selecting experts. Sometimes, an expert is selected because it is marginally better than other close competitors. If it so happens that the top expert is off-chip, this represents an opportunity to loosen the bottleneck in a memory-bound decoding process.The thresholding approach aims to select good-enough experts by boosting the activation probability of on-chip experts artificially by α, a user-defined hyper-parameter. This has the effect of tipping the balance in favor of the experts on-chip in cases where a close competitor to the top expert is on-chip. The equation below describes how the probability is adjusted for an expert E i , in a residency-aware manner.W eights i = W eights i + α if E i is in fast mem W eights i if E i is in slow mem2) Biasing: Using the profiling defined in subsection III-A, we estimate the normalized activation frequencies f req, as shown in Figure 3a . We then define a more expressive penalty for off-chip experts that penalizes the choice of experts by both the frequency of its activation and scales it by the user-defined hyper-parameter β. The key idea in the biasing method is that bringing a rarely used expert on-chip will likely result in it being swapped out again in a later token, creating two swaps between HBM and host memory in an already memory-bound process. Instead, settling for a competitor that is likely to be reused improves latency. This is an accuracy-performance trade-off controlled by the user through β, but it also considers the fact that the frequently used experts are likely to be the top choice for later tokens, which presents a second advantage of biasing over thresholding.The equation below describes the use of the penalty to adjust the logit for expert E i in cases where it is on or offchip. Note that in contrast to thresholding, we adjust the raw logits and then apply the softmax function to obtain the final probabilities.Logits i = Logits i if E i is in fast mem Logits i -β(1 -f req(E i )) if E i is in slow mem IV. EVALUATIONIn this section, we verify our earlier hypotheses about MoE behavior and benchmark the inference latency with different configurations. We focus our quality evaluations on Mixtral-8x7B models since they represent the current state of the art among open-access MoE models. We organize this section as follows: subsection IV-A presents the key insights for expert activations based on the hypothesis presented in subsection III-A. subsection IV-B compares the real system speedups when using MoE-ERAS. Finally, subsection IV-C measures the quality implications of using the resident expert.We examined the hypotheses in subsection III-A by running a large inference workload on a text summarization task using Mixtral-8x7b, and Switch-T-32E on the CNN DailyMail Dataset [16] . It is a summarization dataset that contains long text articles and condensed summaries of the article called highlights. This allows the collection of the activations in both the pre-fill and decode stages, but we focus on the sequential decode phase for analysis. We collect output logits from the gating networks and infer the k selected experts and the softmax distribution in each stage. For Mixtral, given h = 32 hidden layers, each containing E = 8 experts, we obtain a 32 × 8 tensor containing the logits from gating networks. For Switch Transformer, this is instead h = 6 and E = 32, giving a 6 × 32 activation tensor.We collect activation data over 139k tokens for Mixtral and 500k tokens for Switch Transformer. The analysis of the router activations presents support for our hypotheses. Figure 3a and Figure 3b present distributions normalized along each layer. The dark and light spots in these visualizations represent experts that are rarely and frequently activated, respectively, where activation is defined as being in the top k = 2 experts within the layer. This supports our hypothesis (a). For Mixtral, a perfectly equitable distribution of tokens would give 0.125 for all experts, and the visualization confirms that many experts are above that threshold. Likewise, for the Switch Transformer, many experts are activated well in excess of 0.03 (1/32).In order to study expert-expert activations, we obtain the output logits from each gating network (32 for Mixtral, 6 for Switch Transformer), apply the softmax function, and build a correlation matrix. The correlation matrix shows both cells with high positive and negative correlations. Considering that this was collected over 500k tokens, this supports the idea that expert-expert correlations of reasonable strength exist regardless of the input tokens.This leads us to the idea that we might be able to see greater predictability by considering the first few layers together. If it is possible to predict the expert activations of layers deep in the network, then it can support scheduling expert pre-fetching for requests ahead of time. Predicting expert activation based on the first few layers is particularly meaningful since this can lend itself to efficient pre-fetching. We use the activation probabilities of each expert in the first 4 layers as the input features to the regression model. We predict the activated experts for k = 2 in layers 8 -32, and the prediction accuracy scores are shown in Figure 5a . For the Switch Transformer, Figure 5b shows an attempt to predict softmax values instead of activations. The promising results show that we consistently beat random chance significantly. With Mixtral, we are able to predict the activated experts in first and second place with ¿50% accuracy in all experts, with some surpassing 70%. These beat the random chance of 12.5% substantially.There are three important factors to weigh while considering these results. First, this may be an underestimation of the accuracy since we individually predict the first and second experts. If the regression selects the same two experts in inverse order, it diminishes the accuracy but has no effect for the purpose of pre-fetching. Second, this is simple linear regression, and the intention is to show that they are correlated simply. In implementing this in a scheduler, it may be replaced with a few fully connected layers with non-linearity to improve prediction accuracy. Finally, the pattern is borne out over our workload of 500k tokens, and while we believe this to be a general pattern, distribution shifts may be addressed by periodically re-calibrating the regression as done in other works such as MoE-Infinity [?].Similar interpretations are borne out for the Switch Transformer. There, we try to predict the softmax values of the experts directly -a harder problem. We see that we still beat random chance consistently, as shown in Figure 5b .However, key challenges remain to be explored in this direction. This paper focuses on cases where batch size = 1 is applicable, such as commodity and edge settings. However, we also analyze the growth of the number of activated experts with batch size (without thresholding/biasing) to understand the applicability of this work to resource-constrained settings where larger batch sizes are desired. Figure 6 shows this analysis for Switch Transformer-32E, and the number of distinct experts grows to about E/2 at batch size 16. The later layers also consistently show more diversity than the earlier layers, which is an interesting perspective. This supports the idea that pre-fetching the later layers using the predictions presented above can present substantial gains in future improvements of this work.In this section, we profile the speed-ups we are able to achieve with biasing and thresholding. We examine the se- . quential decoding to count the number of expert loads saved, and the overall impact on latency. We compare against the baseline implemented in dvmazur/mixtral-offloading, which includes quantization and expert caching. Our optimizations are orthogonal to these and can be applied with or without other techniques. We consider our top-K routing with quantization and LRU caching as proposed in Moe-offload as our baseline. We generate sequences with l = 100 tokens, n iter = 50 times, and collect the mean latency (wall clock time) and throughput. We see substantial gains using these approaches as seen in Figure Figure 7 .We see two insights from this result: First, the threshold determines the savings. In all offload settings, thresholding requires selecting α, which we test at 0.05, 0.15, 0.25. We find that as the threshold increases, the performance improves owing to saving more expert offloads. Together with the quality metrics in subsection IV-C, a threshold can be selected for the desired performance. These performance metrics should only be compared between approaches in this paper, as latency, throughput, and tokens/second metrics are hardware-dependent. While the ordering should be the same on other hardware, the actual numbers will likely differ.Second, as of f load per layer grows, the savings become more significant. As more experts are offloaded when less VRAM is available, it becomes more likely that an off-chip expert is activated, causing performance degradation while the decoding waits for experts to be brought into memory. This shows that as the environment gets more and more resourceconstrained, our approach becomes more important.In summary, depending on the number of experts offloaded, we find that we can achieve 10% -13% reduction in latency using thresholding at α = 0.15, and 8.0% to 9.7% reduction using biasing with β = 1. At higher α, we achieve even more savings as shown in Figure 7 . Since this work represents a performance-accuracy trade-off, subsection IV-C examines the quality of the generation with these performance gains.Next, we test how different residency-aware routing schemes affect MoE inference quality. We only perform the quality experiments with Mixtral-8x7B as that is the SOTA MoE open-source model. We measure perplexity for Wiki-Text2 [13] and C4 [15] . We also measure 5-shot MMLU [8] accuracy. For WikiText2 and C4, we use the test set and validation sets, respectively. We use a sliding-window strategy with a stride of 512 and a max generation length of 2048. For MMLU, we ran the test over the complete dataset.As shown in Table I , our expert activation technique presents minimal quality degradation at low threshold values. As we increase the threshold β, the quality goes down. This result and speedup seen in the previous section present a quality-speedup trade-off for MoE model inference.Several prior efforts have a similar goal of reducing the inference latency of mixture-of-experts models.EdgeMoE [20] aims to reduce the latency of inference of MoEs on edge systems. It uses quantization and 1-2 layer early [18] performs activationaware prefetching and caching of experts. They use a sample workload (e.g., validation) to form Expert Activation Matrices (EAMs) that they store in a collection. They rely on temporal locality (repeated activation of an expert in a sequence) and sparse activation (only a few activated) assumptions to select the expert to cache and prioritize the prefetch. MoE-offload [6] propose quantization along with LRU caching and hidden state-based expert prediction for MoE inference on commodity hardware. While all these works are focused on expert prefetching and/or quantization, our work focuses on taking expert residency into account. Thus our work is orthogonal to all the related works and can be implemented along with any other proposed quantization or prefetching technique.In this paper, we have shown that Expert Residency Aware Selection (ERAS) shows considerable performance gain for those running Mixtral-8x7b in resource constrained environments requiring expert offloading to host memory. We provide parameters the user can tune navigate the accuracyperformance trade-off, and show that the impact of ERAS on perplexity or accuracy is minimal compared to the performance benefit it offers. This can be applied on top of, or instead of other approaches like parameter compression for performance gains.However, this work comes with limitations. While our profiling and analysis include Switch Transformer as well, our implementation is limited to Mixtral at the moment. While we show both downstream tasks and text generation accuracy, a larger validation on all available benchmarks is required to establish the accuracy retained on other tasks. In addition, since both thresholding and biasing are inference time changes, they may redirect tokens to experts that have seen few such training examples, leading to increased risk of hallucinations.In our next steps, we aim to establish test it on more comprehensive evaluation benchmarks, implement it for other MoE models, and study the effect of biasing/thresholding without aggressive quantization to compare the trade-offs. 
 
 Appendix:  


Decision: Congratulations! We are pleased to inform you that your paper has been accepted for presentation at MLArchSys 2024. We look forward to your participation at the workshop. Further details regarding the schedule and format will be provided soon. See you at the workshop!

Reviews: review: title: title: Interesting activation selection strategy, yet more cost analysis is needed show if it works and what scenario it fits better than the previous work
strength_weakness: strength_weakness: The idea is interesting yet more experiment is needed to verify if it works. The result would be more trustworthy if the author can explain how the prediction accuracy will affect the e2e metrics (i.e. latency and throughput). It would be better to show what is the sweet spot when trading "extra cost of better prediction accuracy w/ more complex regression model" by "better latency & throughput".
questions: questions: The paper describes the overall design of Profiling Expert Activations and Expert Residency Aware Routing. It would be better to share more insight on how to deploy the overall selection strategy into the system. It would be better if the author can address these questions: Is it one-time cost of training and building such regression model for prediction accuracy? Do we need recurring training to make it adaptive to different resources of tokens? What is the extra cost of running such regression model during inference and serving?
title: title: The paper was interesting to read with simple effective approaches, however adding data for more meaningful insights could improve it a lot
strength_weakness: strength_weakness: - The paper needs more clarity between ideas suggested vs evaluations discussed. A prefetching idea was discussed in Section 4.1, predicting experts based on data from initial layers seems to be a good idea, however performance gains from it are not evaluated.  

- Experimental setup details, additional information on how latency is measured. 

- Some sensitivity analysis across all parameters would be useful.
questions: questions: ## Summary

This work aims to slightly tweak the expert selection algorithm in an MOE layer to enable performance improvement by preferential selection of resident experts. To achieve this goal, this work proposes and evaluates speedups and quality across two key techniques (1) Thresholding (2) Biasing which nudge the expert probability activation to prefer resident experts. It claims that such an approach adds minimal degradation to quality.


## Few suggestions/comments:

- Profiling expert activations to show hot and cold experts is very insightful. Section 4.1 shows the data across all tokens in section 4.1. Are there any more insights on correlation between expert activation and input queries, can input query categorization help ?
- Biasing is a good technique, but it has a tendency to downplay or completely avoid a rare expert, as it is reliant on past frequency. Some added functionality to mitigate this corner case may be better. I am hinting at approaches similar to branch prediction, where recent history is accounted for prediction instead of total history. Also input categorization may have its own implications on this effectiveness of biasing.  
- Figure 4 for correlation coefficient is unclear, could you add more labeling to the figure indicating what exact data is correlated ?
- Expert Prefetching: 
  - Does Figure 5b shows probabilities, what is on the X-axis ?
  - The speedups from expert prefetching has not been evaluated right?
- Batching effects:
  - Could you add discussion on how the proposed ERAS techniques work in the case of increased batching ?
- Please provide information on platform/setup/methodology for profiling in 4.2, Eras speedups.
  - The github repo seems to use rely on the GPU alloted to the google colab, which seems to be a V100.
  - If possible, numbers from A100 would be more insightful considering the current GPU deployments in datacenter. As A100s or H100s have increased memory capacity from 40 GB to 80 GB memory, gains on these chips are relevant.
- Could you expand the Results in Section 4.3
  - Biasing parameter $\beta$ - sensitivity? It will be interesting to see how speedups and quality vary with $\beta$
  - Could you include an experiment with both thresholding and biasing ? Alpha and beta together?
title: title: Expert selection based on heuristics for efficient MoE models
strength_weakness: strength_weakness: More evaluations of proposed techniques on diverse benchmarks are required for validation.
questions: questions: The paper proposes heuristic-based techniques to make MoE models more efficient for inference. The proposed techniques, thresholding and biasing, can be orthogonally applied with other optimizations proposed for MoE models. The authors profile the Mistrail and SwiftTransformer models to gather heuristics on determining which experts get activated based on correlation. That allows reduced movement of expert parameters from HBM to low speed memory. All ideas are well discussed and motivated. The evaluations are very detailed and insightful. The paper's writing could be further improved. 

1. There are some grammatical issues in Section 2.2 and Section 4.1. 

2. While the insights and existing evaluation are detailed, more evaluation on diverse benchmarks is required to see if similar insights and heuristics still hold.


3. Are both thresholding and biasing applied together? Are there any benefits to doing that? Is biasing + thresholding included in the evaluation in Figure 7?

4. The author suggests that experts for next layers can be prefeteched based on the prediction. What are the differences and similarities with other existing prefetching techniques? Are there any tradeoffs?
title: title: This work seems like a good start to efficient MoE Execution
strength_weakness: strength_weakness: Some components about accuracy versus performance need further exploration.
questions: questions: The authors are presenting their technique to improve the performance of MoEs by predicting routing and by factoring the expected residency of the experts.  This work is finding an interesting way to deal with resource scarcity of MoE execution.  While overall the work seems to be moving in the right direction, there are some areas where it can be improved.  

The primary issue I would like to see the authors address is looking at the performance versus accuracy and the cost or timing of determining the expert routes.  I believe there is still some techniques to be explored here and the costs evaluated in a way that is more efficient besides just predicting using another network.  There may be some training changes that make it easier to allocated to multiple experts at low cost.  

Overall, I would like to see how doing this may impact the hardware.  Is there something we can add to the hardware that will improve over performance when using the MoE allocation scheme proposed here?  Basically, I would like more hardware and the impact on hardware from this work.