Load Balancing

An auxiliary loss is a secondary term added to the primary training objective explicitly to encourage balanced routing. It penalizes the variance in the number of tokens assigned to each expert across a batch.

• Purpose: Prevents router collapse, where a few experts receive all tokens while others are ignored.
• Mechanism: Often implemented as the squared coefficient of variation of the expert assignment counts. A perfect balance results in a loss of zero.
• Impact: This loss is weighted by a hyperparameter (e.g., 0.01) and is crucial for training stable, high-performance MoE models.

Noise Top-k Gating is a routing strategy that adds tunable, random noise to the router's logits before selecting the top-k experts.

• Function: The noise acts as an exploration mechanism during training, encouraging the router to occasionally try underutilized experts.
• Process: Gaussian noise is added to the expert scores: noisy_logits = logits + std * torch.randn_like(logits). The top-k experts are then selected based on these noisy scores.
• Outcome: This technique, introduced in the Switch Transformer, promotes more uniform expert utilization without sacrificing routing quality, as the noise magnitude can be annealed.

The capacity factor is a critical hyperparameter that defines a soft limit on the number of tokens an expert can process in a forward pass.

• Calculation: Expert Capacity = (batch_size * sequence_length * k / num_experts) * capacity_factor.
• Purpose: It allocates a fixed buffer in memory for each expert's computations. A factor >1.0 (e.g., 1.25) provides slack to handle token assignment variance and minimize dropped tokens.
• Trade-off: A higher factor increases GPU memory usage but reduces the chance of tokens being skipped or passed through; a lower factor saves memory but risks degraded model performance.

Expert Parallelism is a model parallelism strategy where different experts are placed on different devices (GPUs). Load balancing is enforced by the All-to-All communication pattern.

• Process: 1) Local routing decides token assignments. 2) An All-to-All operation scatters tokens from all devices to the specific devices hosting their assigned experts. 3) Experts process their received tokens. 4) A second All-to-All gathers the outputs back.
• Implication: This communication pattern inherently balances load across devices. If one expert is overloaded, its host device becomes a bottleneck, making balanced routing a systems-level necessity for performance.

These are two policies for handling tokens when an expert reaches its capacity limit.

• Hard Capacity: Strictly limits the number of tokens per expert. Excess tokens beyond the expert's capacity are dropped (passed to the next layer unchanged) or routed to a fallback expert. This enforces load balance but can lose information.
• Soft Capacity: Allows experts to process slightly more tokens than the nominal capacity by temporarily expanding buffer sizes, often at a performance cost. This prioritizes model accuracy over strict balance.
• System Choice: Most production systems (e.g., vLLM) implement hard capacity for deterministic memory allocation and latency.

Router Z-Loss is an auxiliary regularization term applied to the router's logits to stabilize training and indirectly aid load balancing.

• Formula: Loss_z = (logsumexp(router_logits))^2. It penalizes the router logits from growing too large.
• Effect: By controlling the magnitude of the router logits, it prevents the router logits from entering regions of extreme softmax values, which can lead to a winner-takes-all scenario and poor load balancing.
• Usage: Often used in conjunction with the auxiliary load loss. It improves numerical stability and encourages healthier gradient flow through the gating network.

An additional term added to the primary training loss function specifically designed to penalize imbalanced routing. Its purpose is to encourage the gating network to distribute tokens more evenly across experts, preventing expert underutilization. Common formulations include:

• Importance Loss: Minimizes the squared coefficient of variation of the total routing probability assigned to each expert across a batch.
• Load Loss: A more direct, batch-aware loss that minimizes the squared difference between the fraction of tokens routed to an expert and a perfectly uniform target (1/num_experts). Without this loss, routers often collapse, favoring a small subset of experts.

A critical hyperparameter that defines a soft limit on the number of tokens any single expert can process in a forward pass. It is calculated as: expert_capacity = (batch_size * sequence_length * k) / num_experts * capacity_factor

• A factor >1.0 (e.g., 1.25) provides a buffer to accommodate statistical fluctuations in routing, reducing the chance of dropped tokens.
• A factor closer to 1.0 increases computational efficiency but risks more dropped tokens if routing is uneven. It directly implements a form of hard capacity limiting during inference, forcing load balancing by capping each expert's workload.

A routing strategy that adds tunable, multiplicative Gaussian noise to the expert logits before applying the top-k selection. This technique serves two key load-balancing functions:

1. Encourages Exploration During Training: The noise destabilizes early, arbitrary routing preferences, giving more experts a chance to receive gradients and improve, preventing permanent router collapse.
2. Promotes Uniformity: By making the gating scores noisier, it reduces the model's overconfidence in consistently routing specific token types to the same few experts. It is a foundational technique in models like the Switch Transformer to ensure all experts are trained effectively.

A model parallelism strategy where different expert networks are placed on different physical devices (e.g., GPUs). This architecture makes load balancing a system-critical requirement.

• Imbalanced routing leads to severe device underutilization, as some GPUs sit idle while others are overloaded.
• It relies on efficient all-to-all communication to scatter tokens to and gather outputs from the experts' respective devices. The auxiliary load loss is therefore essential not just for model quality, but for achieving high hardware efficiency and throughput in distributed MoE systems.

A direct consequence of failed load balancing when an expert reaches its hard capacity limit. Tokens assigned to a full expert are 'dropped' and not processed.

• Handling Strategies: Dropped tokens may be passed to the next layer unchanged, routed to a fallback expert, or masked. All strategies degrade model performance.
• Relationship to Load Balancing: The primary goal of load balancing techniques is to minimize token dropping. The capacity factor is tuned to find a trade-off between computational waste (high buffer) and performance loss from dropping (low buffer). It represents the failure mode that load balancing aims to prevent.

A training failure mode where the gating network learns to route all, or a vast majority, of tokens to the same small subset of experts (e.g., 1 or 2 out of 8).

• Result: The model effectively behaves like a much smaller dense network, wasting the parameters and representational capacity of the unused experts.
• Causes: Can occur without explicit load balancing incentives due to initial conditions or reinforcement loops during training.
• Prevention: Directly addressed by auxiliary load losses and noisy gating, which provide the necessary gradient signals to maintain balanced expert utilization throughout training.