Qwen Global-Batch Load Balance for MoE LLM Training

Global-batch load balance improves MoE performance and expert specialization

Qwen has introduced a global-batch load balancing loss method for training Mixture-of-Experts (MoE) Large Language Models (LLMs). By shifting the calculation of load-balancing loss from the micro-batch level to the global-batch level, the model achieves better performance across various tasks and fosters significant domain specialization among experts, addressing a limitation where micro-batch balance often prevents experts from specializing in specific data domains.

The limitation of micro-batch load balancing

Most existing MoE training frameworks, such as Megatron-core, implement load balancing loss ($L_{\text{balance}}$) at the micro-batch level. This means the loss is calculated within every micro-batch and then averaged across the global batch.

This approach is problematic when a micro-batch lacks diverse data. For example, if a micro-batch contains only code data, a micro-batch-level loss forces the router to distribute those code tokens uniformly across all experts. This prevents experts from specializing in specific domains (like code) and can hurt overall model performance. Because data in a single micro-batch is often from the same domain in LLM training, this explains why many open-source MoE models lack notable expert specialization.

Implementing global-batch load balance

To resolve this, Qwen proposes calculating the load-balancing loss at the global-batch level ($L_{\text{global}}$). This is achieved through three steps:

  1. Synchronizing expert selection frequency ($f_{i}$) across all parallel groups.
  2. Calculating the load-balancing loss in each parallel group (e.g., on a single GPU).
  3. Aggregating the loss across all micro-batches.

Because the expert selection frequency is a small vector (with dimensions equal to the number of experts), synchronizing this data across micro-batches is computationally inexpensive, making the implementation "almost free."

Performance gains and expert specialization

Qwen tested the method across three MoE configurations (3.4B with 0.6B activated, 15B with 2.5B activated, and 43B with 6.6B activated) and two data configurations (120B and 400B tokens).

Improved Model Effectiveness

Compared to micro-batch-level loss, the global batch approach achieved better performance in all tested settings, including all model sizes, data volumes, and tasks.

Enhanced Domain Specialization

Global batch balance enables experts to specialize. While micro-batch balance results in experts being uniformly activated regardless of the domain, global batch balance allows specific experts to be frequently activated by specific domains, demonstrating clear specialization.

Impact of Balance Batch Size

Experiments with a 3.4B model (0.6B activated) showed that ptr-training PPL (perplexity) decreases rapidly as the balance batch size increases from 2 to 128, saturating after 128. This highlights the significance of the global approach, as mainstream MoE frameworks typically use a balance batch size between 8 and 16 for larger models.

Balancing efficiency and effectiveness

While global batch balance can improve model quality, it may lead to a degradation in micro-batch balance, which can negatively impact computational efficiency.

To mitigate this, Qwen experimented with adding a micro-batch balance loss on top of the global-batch balance loss (using a constant weight of 0.01 of the global batch loss). This hybrid approach improved the update step speed (from 1.64 to 1.59 seconds per update step) while leaving the model's effectiveness almost unaffected.

Conclusion

By implementing global-batch load balance, Qwen addresses a critical limitation in MoE training that allows for more performant and specialized experts. This method provides a novel perspective for optimizing MoE models, particularly for the ability to scale to more substantial and specialized MoE models across various domains.

Sources