Replicate and Quantize: A Plug-and-Play Strategy for Load Balancing in Sparse Mixture-of-Experts LLMs

We thank the reviewers for the comments and suggestions to improve our paper. Please see the following clarification.

[W1: ablation study - Sure]

We provide an ablation study of replication and quantization as follows.

[W2: Does this phenomenon also exist in other tasks and other MoE blocks? - Yeah]

In Figure 2, our goal is to demonstrate the existence of two distinct dimensions among the experts, which represent our observation framework rather than a phenomenon inherent to the experts themselves. These dimensions provide a structured way to quantify and analyze the functionality of the experts. While PIQA serves as a representative example, we observed that this two-dimensional framework can be applied consistently across other tasks and MoE blocks.

[W3: Related Work in load balance - Sure]

The current load balanced method focuses on the training stage. For the current SMoE model, they have tried to solve their load imbalance issue by introducing the imbalance loss or other methods during training the model. Moreover, there are some current works focus on the gpu utilization rather than the load balance for each experts.

[W4: Tuning strategies - we have show their results in Table 1]

[W5: The details of adopting the proposed method on DeepSeekMoE should be given - Sure]

The Deepseek MOE total has 64 experts in each layer: 63 isolated experts and 1 shared expert.  In each routing, the system will select three experts from the 63 isolated experts, always using the shared expert. Therefore, our strategy is exclusively applied to the 63 isolated experts.

[W6: Do all comparisons share the same activated parameters in inference? - Yeah, they shared the same activated parameters during the inference]

[W7: Format - Thank you for reminder. We have updated the paper. We will post a revised version with a summary of revision.]

[W8: The scalability of the proposed method is encouraged to be evaluated or discussed. - Sure]

In Figure 3, Panels (a) and (b) illustrate the load balance scores across multiple timesteps under two distinct routing strategies. Panel (a) demonstrates the case where the system utilizes cumulative information aggregated from all prior timesteps, providing a holistic approach to decision-making. Conversely, Panel (b) focuses on a scenario where only information from the immediately preceding timestep is leveraged, showcasing a more localized decision-making approach.

The purpose of this setup is to identify the most important "heavy-hitter" experts from the less important ones. These decisions depend directly on the input data at each timestep. To mimic how data arrives in a real-world streaming scenario, we divided the MMLU dataset into 10 parts and set the timestep count to 10. This approach allows us to simulate the flow of data over time and study how it affects the routing strategy in a manageable and realistic way.

We thank the reviewers for the comments and suggestions to improve our paper. Please see the following clarification.

[Q1: The paper's claim on its novelty, that the load imbalance of MoE models has not been studied for inference scenarios, is wrong. - We will discuss those papers as follows]

Towards MoE Deployment: Mitigating Inefficiencies in Mixture-of-Expert (MoE) Inference.

Focus: load balance for each GPU, not the experts, and they do not release their after-accept by Nips 2024, and the description in the paper is too vague for us to replicate their methods.

MEGABLOCKS: EFFICIENT SPARSE TRAINING WITH MIXTURE-OF-EXPERTS

Focus: Sparse training, not inference.

SwapMoE: Serving Off-the-shelf MoE-based Large Language Models with Tunable Memory Budget

Focus: Inference on edge devices with memory constraints. Dynamically load, activate, and swap experts during inference. From the system aspects

Accelerating Distributed MoE Training and Inference with Lina

Focus: Both training and inference.
Load Balancing: If an expert is overloaded, another expert from the top-k set (on the same device) is selected.

Pre-gated MoE: An Algorithm-System Co-Design for Fast and Scalable Mixture-of-Expert Inference

Focus: Add a pre-gating mechanism, which modifies the model architecture.

Mixture of Quantized Experts (MoQE): Complementary Effect of Low-bit Quantization and Robustness

Focus: Memory optimization via quantization. They mainly conduct thorough experiments to show the experts weight is more robust than we expect.

[W2: The Load Imbalance Score defined at Sec 3.1 is an ill-defined metric. - It's good ]

We measure the load distribution across experts for the entire dataset because a batch is typically a uniform sample from the dataset. Therefore, using the dataset-level measurement provides a reasonable approximation.

[W4 & Q1: Fig. 2 shows that the most important expert is the expert that receives the most tokens. - The symptom showing the experts inner function does not influence our strategy.]

In Figure 2, our goal is to demonstrate the existence of two distinct dimensions among the experts, highlighting that their functionality can be analyzed from these two different perspectives. While it is true that in some cases, the most important expert may also be the heaviest one, this does not hinder our method to quantize the less important experts and replicate the heavier ones. While certain experts may excel in multiple aspects, this does not negate the presence of distinct functional dimensions that allow us to manage and optimize the experts accordingly.

[W5 & W6: Thanks for your suggestion, we will modify it!]

[Q2: How does the router adapt to the case where the most important expert is replicated? Will it evenly distribute its tokens to each GPU device? - Randomly and Not necessary]

When the router adapts to the case where the most important expert is replicated, it randomly assigns tokens to one of the replicated experts for processing.

Based on our experiments, splitting the workload of a heavy-hitter expert between just two replicas significantly alleviates the load imbalance. It is not necessary to distribute the workload across all GPUs to achieve this improvement.

[Q3: How are the expert loaded on the GPU? Are the other experts completely unaffected? - Not affected]

In the quantized version, the compressed weights are loaded into the GPU as low-bit representations. This allows the quantized experts to occupy significantly less memory, enabling their efficient deployment on the GPU.

Other experts are unaffected by this process as their weights remain unchanged and continue to operate in their original precision.

We thank the reviewers for the comments and suggestions to improve our paper. Please see the following clarification.

[W1: Weak presentation of algorithms and figures. Some figures lack any caption or explanation. - Thanks! We will complement them! ]

Algorithm 1: Identify Heavy-Hitter Experts

Purpose: This algorithm identifies the most frequently selected "heavy-hitter" experts for each Mixture-of-Experts (MoE) layer based on input tokens and their routing.

Inputs:

• Tokens: input_tokens - A list of input tokens.
• Experts: num_experts - The number of experts in each MoE layer.
• Layers: num_layers - The number of MoE layers.
• Token Count: num_tokens - The number of tokens to be processed.

Outputs:

• Heavy Experts: heavy_hitters - A list where each entry corresponds to the most selected expert in each layer.

Algorithm:

1. Initialize: heavy_hitters ← list of size num_layers

2. For each layer layer_index in range num_layers do:

   a. Initialize: expert_selection ← empty list

   b. For each token token in input_tokens do:

   i. selected_expert ← route_token_to_expert(token, layer_index)

   ii. Append selected_expert to expert_selection

   c. Compute expert_frequencies ← count_frequency(expert_selection)

   d. heavy_hitter ← find_expert_with_max_frequency(expert_frequencies)

   e. Store heavy_hitter in heavy_hitters[layer_index]

3. Return heavy_hitters

[Q1: Error on Page 6, line 288: Are the X-axis and Y-axis labels inverted? - Thanks]

Change to “As shown in Figure 2, the Y-axis represents the workload of an expert by the number of allocated tokens, while the X-axis displays the inverse importance score of an expert (see Definition 3.3). “

[Q2: Should R&Q identify heavy-hitter and important experts for each individual task, or can the identified experts be reused across tasks? - Of course]

They should be identified for each task; as you mentioned, each expert may have a different performance. Moreover, we have evidence that for each task, we can only utilize 0.1% of the data to identify the heavy-hitter experts and less important experts. We then apply the R&Q strategy to the remaining data, focusing on those experts who continue to perform well.

[W2: Weak baseline. - The baseline you mentioned is the one for choosing the less-important experts, not the one for load balance; this one should refer to Table 1.]

The current load balanced method focuses on the training stage. For the current SMoE model, they have tried to solve their load imbalance issue by introducing the imbalance loss or other methods during training the model. So in our Table 3, the raw load balance scores are the model already equipped with their load balanced method.

[W3 & Q3: Weak empirical analysis on computational efficiency gain. - The quantization did slow down inference]

Theoretically, our strategy will not increase the memory because when we duplicate the heavy-hitter experts, we also quantize the duplicate one to a half precision, then we quantize the less important one to a half precision.

[W4 & Q4: Updated the more challenging tasks - Sure!]

Results on other generation task:

TruthfulQA Results

Load Balance Results

We thank the reviewers for the comments and suggestions to improve our paper. Please see the following clarification.

[W1&Q1: Implementation Details - Sure!]

Hyperparameters

Quantization techniques: For Switch Transformer they load in float32 so we use the half quant with the float16 For Llama and Deepseek which load in float16, so we used 8 bit quantization refer (https://arxiv.org/abs/2208.07339)

[W2&Q2: Additional ablation studies to demonstrate the individual contributions of the replication and quantization - Sure!]

We provide an ablation study of replication and quantization as follows.

The replication is for load balancing and it maintains the model performance. The quantization of less important expert ensures us to fit the total memory budget while maintaing the overall model performance.

[Q3: Studies on different levels of model sparsity- We include SMoEs with various sparsity.]

We have conducted experiment in the paper with the following models with different sparsity.

Switch Transformer 8 experts (Sparsity 1/8): total have 8 experts in each layer, only choose 1 expert for each token

Switch Transformer 16 experts (Sparsity 1/16): total have 16 experts in each layer, only choose 1 expert for each token.

Llama MoE 8 experts (Sparsity 2/8): total have 8 experts in each layer, choose 2 experts for each token

Deepseek MoE (Sparsity 1+3/63): total has 64 experts in each layer: 63 isolated experts and 1 shared expert.  In each routing, the system will select three experts from the 63 isolated experts, always using the shared expert. Therefore, our strategy is exclusively applied to the 63 isolated experts.

Thanks for the detailed response! I have adjusted my score given the reviews and the replies to reflect these improvements.

Our Response to Concerns

On Scalability:

• 1. Add more challenging task:

  TruthfulQA generation results and other challenging Tasks:

• 2. Clarify the previous experiments setting We have conducted experiment in the paper with the following models with different sparsity. Switch Transformer 8 experts (Sparsity 1/8): total have 8 experts in each layer, only choose 1 expert for each token Switch Transformer 16 experts (Sparsity 1/16): total have 16 experts in each layer, only choose 1 expert for each token. Llama MoE 8 experts (Sparsity 2/8): total have 8 experts in each layer, choose 2 experts for each token Deepseek MoE (Sparsity 1+3/63): total has 64 experts in each layer: 63 isolated experts and 1 shared expert. In each routing, the system will select three experts from the 63 isolated experts, always using the shared expert. Therefore, our strategy is exclusively applied to the 63 isolated experts.

• 3. Scalability in streaming data In Figure 3, Panels (a) and (b) illustrate the load balance scores across multiple timesteps under two distinct routing strategies. Panel (a) demonstrates the case where the system utilizes cumulative information aggregated from all prior timesteps, providing a holistic approach to decision-making. Conversely, Panel (b) focuses on a scenario where only information from the immediately preceding timestep is leveraged, showcasing a more localized decision-making approach.

  The purpose of this setup is to identify the most important "heavy-hitter" experts from the less important ones. These decisions depend directly on the input data at each timestep. To mimic how data arrives in a real-world streaming scenario, we divided the MMLU dataset into 10 parts and set the timestep count to 10. This approach allows us to simulate the flow of data over time and study how it affects the routing strategy in a manageable and realistic way.

On Inference Efficiency:

Our method ensures memory usage does not increase even when duplicating new experts. After duplication, we quantize the new experts and further quantize less-important ones to half precision.

Regarding inference speed, in each MoE layer, the inference time is determined by the slowest expert, which handles the largest number of tokens. Below, we provide an example to illustrate this behavior.

### Example for Each MoE Layer

Premise:

- Number of Experts: 4  
- Number of Tokens: 80  
- Processing Capacity per Expert: Each expert processes (80 / 4 = 20) samples per unit time (T).  

---

Original Routing Results:

- Experts' Token Allocation: [18, 30, 25, 7]  
- Inference Time (Slowest Expert): 
  Expert 1 handles 30 tokens:  
  30 / 20 = 1.5 T

---

After Rerouting & Quantization (R & Q):

- Experts' Token Allocation: [18, [17, 13], 25, 7]  
- Inference Time (Slowest Expert): 
  Expert 2 handles 25 tokens:  
  25 / 20 = 1.25 T

Result:
By applying our method, the inference time is reduced, resulting in a speed-up of 0.25 T.