PaperlayerView exact revision on arXiv

Source-linked AI summary

Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, Damai Dai

arXiv:2408.15664v1cs.LGcs.CL

TL;DR

MoE training needs balanced expert routing to avoid routing collapse and computational overhead, but auxiliary losses can interfere with the main objective. The paper proposes Loss-Free Balancing, which updates expert-wise routing biases from recent loads without auxiliary-loss gradients. Experiments report better model performance and load balance than traditional auxiliary-loss training, within the evaluated MoE settings.

  • Problem

    MoE load imbalance can cause routing collapse or increased computational overhead, while auxiliary-loss control creates a trade-off between balance and model performance.

  • Method

    Loss-Free Balancing adds expert-wise biases before top-K routing and dynamically updates them according to recent expert loads without introducing auxiliary-loss gradients.

  • Results

    Loss-Free Balancing achieves better model performance and load balance than traditional auxiliary-loss training in experiments on 1B and 3B MoE models.

  • Takeaways & Limitations

    The method provides an auxiliary-loss-free way to control expert load while preserving the primary language-modeling training objective.

  • Takeaways & Limitations

    The paper emphasizes causal language modeling, where alternative Expert Choice routing is unsuitable because it causes future token leakage.

Abstract

from arXiv · show

For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead. Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance. In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing, featured by an auxiliary-loss-free load balancing strategy. To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert. By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load. In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training. We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens. Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.

1 INTRODUCTION

MoE models can suffer load imbalance that causes routing collapse or computational overhead. Loss-Free Balancing controls expert loads without auxiliary-loss interference gradients and improves both performance and balance in experiments.

  • MoE load imbalance can cause routing collapse or increased computational overhead.
  • Loss-Free Balancing is an auxiliary-loss-free strategy designed to control expert load without introducing interference gradients.
  • Expert-wise biases are added before top-K routing and updated from recent expert loads to produce balanced routing results.
  • Experiments train 1B-parameter models on 100B tokens and 3B-parameter models on 200B tokens from scratch.
  • Loss-Free Balancing achieves better validation loss and significantly better global and batch-level load balance than traditional auxiliary-loss-controlled models.

2 BACKGROUND

MoE layers reduce scaling costs by routing tokens to selected experts, but uncontrolled routing can imbalance loads. Auxiliary loss encourages balance while potentially interfering with language-model training, creating a performance–balance dilemma.

  • 2.1 MIXTURE-OF-EXPERTS: MoE architectures replace Transformer MLP layers with MoE layers, where Top-K routing selects experts for each token.
  • 2.2 AUXILIARY LOSS FOR LOAD BALANCE: Uncontrolled routing can cause routing collapse, leaving some experts insufficiently trained, or create computation bottlenecks across devices.
  • 2.2 AUXILIARY LOSS FOR LOAD BALANCE: The auxiliary loss uses expert token fractions and average gating scores, with α controlling its strength.
  • 2.2 AUXILIARY LOSS FOR LOAD BALANCE: A small or absent α can produce poor load balance, whereas a large α can interfere with language-model training and impair performance.

3 AUXILIARY-LOSS-FREE LOAD BALANCING STRATEGY

Loss-Free Balancing adjusts expert gating scores with dynamically updated biases rather than an auxiliary loss. Its design targets balanced routing while avoiding direct interference with the main training gradients.

  • 3 AUXILIARY-LOSS-FREE LOAD BALANCING STRATEGY: Loss-Free Balancing adds an expert-wise bias to gating scores and uses the biased scores for routing decisions.
  • 3 AUXILIARY-LOSS-FREE LOAD BALANCING STRATEGY: The bias is initialized at zero, and each training batch updates it according to the expert’s load violation error.
  • 3 AUXILIARY-LOSS-FREE LOAD BALANCING STRATEGY: The bias affects only top-K expert selection and is not added to the weights used for the selected experts’ final outputs.
  • 3 AUXILIARY-LOSS-FREE LOAD BALANCING STRATEGY: Heavy-load experts have their biases reduced, while lighter-load experts have their biases increased during iterative updates.
  • 3 AUXILIARY-LOSS-FREE LOAD BALANCING STRATEGY: Compared with auxiliary-loss control, Loss-Free Balancing avoids noisy gradients while addressing the trade-off between load balance and model performance.

4 EXPERIMENTS

The experiments compare Loss-Free Balancing with auxiliary-loss-controlled routing and evaluate perplexity, global and batch-level load balance, and bias-update variants. Loss-Free Balancing achieves better perplexity and load balance, while update-rate and bias-form choices affect training behavior and performance.

  • Experimental Setup: The experiments use DeepSeekMoE-based 1B and 3B models trained on 100B and 200B tokens, respectively.The training corpus is multilingual, and sigmoid gating is used in the main experiments.
  • Experimental Setup: Loss-Free Balancing is compared with an auxiliary-loss-controlled baseline using α = 0.001.The baseline coefficient is selected as a reasonable trade-off between model performance and load balance.
  • Metrics: Perplexity and MaxVioglobal are evaluated on a validation set, while MaxViobatch measures load balance on each training batch.MaxVioglobal reflects expert utilization and efficiency at large batch sizes; MaxViobatch is more related to training efficiency.
  • Main Results: Loss-Free Balancing achieves lower perplexity and much better global load balance than auxiliary-loss-controlled routing for both 1B and 3B models.Its load-balance advantage persists across training steps, as shown by MaxViobatch curves.
  • Update Rate: An update rate of u = 0.001 provides good training balance and validation perplexity, whereas u = 0.0001 converges slowly and u = 0.01 deteriorates later-stage balance.The update rate controls how quickly expert biases converge to suitable values.
  • Update Rule: The error-scaled update rule slightly improves load balance without improving model performance, so the sign-based rule is retained.The tested alternative changes bi = bi + u ∗sign(ei) to bi = bi + u ∗ei.
  • Multiplicative Bias: Multiplicative biases produce similar load balance but slightly worse model performance than additive biases, supporting additive biases as the preferred choice.The multiplicative variant applies expert-wise scaling rather than adding biases to gating scores.

5 DISCUSSION

Loss-Free Balancing improves computation-batch-level load balance as batch size increases and avoids the future-token leakage associated with Expert Choice. Its compatibility with expert parallelism makes the balance advantage more relevant at larger scales.

  • 5.1 LOSS-FREE BALANCING IS COMPATIBLE WITH EXPERT PARALLELISM: Expert parallelism increases the computation batch by ep_data_parallel_size times, making per-step load balance crucial for efficiency.A computation batch contains micro_batch_size * ep_data_parallel_size samples.
  • 5.1 LOSS-FREE BALANCING IS COMPATIBLE WITH EXPERT PARALLELISM: Loss-Free Balancing approaches nearly optimal global balance, with per-step balance moving closer to global balance as computation-batch size increases.Its load balance keeps improving with larger computation batches, unlike the approximately constant level of auxiliary-loss control at large batch sizes.
  • 5.2 LOAD BALANCING AND FUTURE TOKEN LEAKAGE: Expert Choice violates causal language-modeling constraints because future tokens can influence earlier tokens’ expert assignments.The assignment mechanism can allow earlier tokens to infer successor identities.
  • 5.2 LOAD BALANCING AND FUTURE TOKEN LEAKAGE: A nine-layer model with 16 experts and two activated experts per token can leak 50 bits per token through routing allocation.The passage states this amount is sufficient for each token to determine its successor’s identity.
  • 5.1 LOSS-FREE BALANCING IS COMPATIBLE WITH EXPERT PARALLELISM: Figure 5 reports better balance than auxiliary-loss training as computation-batch size increases, especially with a moderately sized computation batch.
  • 5.2 LOAD BALANCING AND FUTURE TOKEN LEAKAGE: Reducing top-K selection chunks from 8192 to 512 tokens produced an abnormal loss drop of about 10%, confirming future-token leakage.Shuffling tokens across chunks was used to make leakage more difficult.
  • 5.2 LOAD BALANCING AND FUTURE TOKEN LEAKAGE: Future-token leakage destroys generalization and prevents reliable evaluation, making Loss-Free Balancing safer than Expert Choice for scaling MoE models.

6 CONCLUSION

The paper concludes that Loss-Free Balancing controls MoE load without auxiliary-loss gradients and achieves better performance and load balance than traditional auxiliary-loss training.

  • Loss-Free Balancing controls MoE load without introducing auxiliary-loss gradients.
  • Experiments on 1B and 3B MoE models trained on 100B and 300B tokens show better model performance and load balance than traditional auxiliary-loss training.

A MODEL ARCHITECTURE

The experiments use DeepSeekMoE as the backbone, combining routed and shared experts while replacing most FFN layers with MoE layers.

  • The backbone is DeepSeekMoE, which introduces shared experts to mitigate redundancy among routed experts.
  • DeepSeekMoE replaces all FFN layers except the dense FFN layer immediately after input embedding with MoE layers.
  • Detailed architecture hyperparameters are provided in Table 5.

B TRAINING SETTINGS

Training uses specified initialization, sequence-length, optimization, batch-size, and validation settings, with sigmoid and softmax gate comparisons also reported.

  • All learnable parameters use standard deviation 0.006, and the maximum training sequence length is 2048.
  • The 1B model uses a cosine scheduler with warmup, learning rates from 1e-3 to 1e-4, batch size 1152, and 40000 steps over 100B tokens.
  • The 3B model uses a multistep scheduler, warmup of 2000 steps, batch size 1728, and 56514 steps over 200B tokens.
  • Validation retains approximately 70M tokens from the training corpus.
  • The softmax gate has higher perplexity and greater sensitivity to imbalance than the sigmoid gate under similar balance conditions.
  • For the softmax gate, Loss-Free Balancing achieves slightly lower perplexity and significantly better load balance than auxiliary-loss training.

C EXPERIMENTS WITH SOFTMAX GATE

With a softmax gate, Loss-Free Balancing achieves slightly lower perplexity and substantially better load balance than auxiliary-loss training, maintaining this balance through most of training.

  • The softmax gate has higher perplexity than the sigmoid gate under similar load-balance conditions.Its performance is also more sensitive to load imbalance than the sigmoid gate.
  • Loss-Free Balancing achieves slightly lower perplexity while maintaining significantly better load balance than auxiliary-loss training.The baseline uses α = 0.0003, selected as the lowest-perplexity auxiliary-loss setting for the softmax gate.
  • Loss-Free Balancing maintains superior load balance throughout most of the training process with a softmax gate.

D.1 PROOF FOR THEORETICAL LEAKAGE AMOUNT

The section defines the routing-information leakage measure for Expert Choice and reports that it can exceed 50 bits per token in a representative sparse MoE configuration.

  • K is the average number of experts activated per token, while N is the total number of experts and defines MoE sparsity.
  • Expert Choice’s routing-allocation combinations can carry a maximum information leakage I measured in bits per token.
  • More than 50 bits per token of total leakage occurs with sparse ratio R = 2^16 = 0.125 across 9 MoE layers.

D.2 EXPERIMENTAL EVIDENCE

Experiments vary Expert Choice’s chunk size and test whether shuffling removes an abnormal loss decrease. Small chunks produce the drop, while shuffling mitigates it, supporting future-token leakage as its source.

  • The study trains a 2B MoE model on 100B tokens while varying expert top-k-selection chunk sizes from 512 to 8192 tokens.The tested sizes include 512, 2048, and 8192 tokens.
  • A chunk size of 512 produces an abnormal loss drop attributed to significant future-token leakage during training.Smaller chunks make future tokens within the same chunk easier to exploit.
  • With chunk size 512, Expert Choice shows a significant loss drop relative to chunk sizes 8192 or 2048, whereas shuffling eliminates the drop.The figure interprets this pattern as evidence of future-token leakage.
  • Shuffling tokens within a batch before chunking and selection mitigates the loss drop by disrupting future-token context.This supports the hypothesis that the loss decrease comes from accessing and exploiting future-token information.
Loading 2408.15664v1…