Load Balancing › Measuring Load Balance in Practice
Language AI HandbookAll chapters Load Balancing › Measuring Load Balance in PracticeOne section · Sections in this chapter6 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement functions to compute load balance metrics and visualize expert utilization across training. The goal is to build the diagnostic toolkit you'll use when training and debugging your own MoE models. We'll start with the foundational measurements and then build up to simulations that show how imbalance evolves during training.
import numpy as np
import torch
# Set random seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)We'll start by creating a simple gating network and generating routing decisions for a batch of tokens.
import torch
import torch.nn as nn
import torch.nn.functional as F
class SimpleGatingNetwork(nn.Module):
"""A basic gating network for routing tokens to experts."""
def __init__(self, input_dim, num_experts):
super().__init__()
self.gate = nn.Linear(input_dim, num_experts)
self.num_experts = num_experts
def forward(self, x):
# x shape: (batch_size, seq_len, input_dim)
logits = self.gate(x) # (batch_size, seq_len, num_experts)
probs = F.softmax(logits, dim=-1)
return probs
# Create gating network
input_dim = 256
num_experts = 8
gating_net = SimpleGatingNetwork(input_dim, num_experts)
# Introduce realistic initialization bias so some experts are initially preferred
# (simulates unequal random init that seeds rich-get-richer dynamics)
with torch.no_grad():
gating_net.gate.bias.data += torch.tensor(
[1.8, -0.6, 1.2, -1.0, 0.9, -1.3, 0.4, -0.8]
)
# Generate sample input
batch_size = 16
seq_len = 128
hidden_states = torch.randn(batch_size, seq_len, input_dim)Now let's compute routing probabilities and expert assignments.
# Get routing probabilities
with torch.no_grad():
routing_probs = gating_net(hidden_states) # (batch, seq, experts)
# Flatten to token-level view
flat_probs = routing_probs.view(-1, num_experts) # (batch*seq, experts)
# Get top-1 expert assignments
expert_assignments = torch.argmax(flat_probs, dim=-1) # (batch*seq,)
num_tokens = expert_assignments.shape[0]Total tokens: 2048 Number of experts: 8 Routing probabilities shape: torch.Size([2048, 8])
The output confirms we have 2,048 tokens in total (16 sequences 128 length). The routing probability tensor shape matches our expectation, containing a probability distribution over the 8 experts for every token in the batch.
Computing Token Fractions
Let's compute the fraction of tokens routed to each expert. The function below implements the token fraction formula directly: for each expert, count the tokens assigned to it and divide by the total.
def compute_token_fractions(expert_assignments, num_experts):
"""Compute the fraction of tokens assigned to each expert."""
counts = torch.zeros(num_experts)
for i in range(num_experts):
counts[i] = (expert_assignments == i).sum()
fractions = counts / len(expert_assignments)
return fractions
token_fractions = compute_token_fractions(expert_assignments, num_experts)Token fractions per expert: Expert 0: 0.697 ███████████████████████████ Expert 1: 0.000 Expert 2: 0.189 ███████ Expert 3: 0.000 Expert 4: 0.099 ███ Expert 5: 0.000 Expert 6: 0.014 Expert 7: 0.000
Even with random initialization and a biased setup to simulate early-training conditions, we can already see substantial imbalance emerging. E0 and E2 capture significantly more than their fair share in this deterministic sample, while the other experts remain underutilized. This is precisely the kind of initial asymmetry that, left unchecked, will amplify into collapse. Let's compute our load balance metrics to quantify this exactly.
def compute_load_metrics(token_fractions):
"""Compute load balance metrics."""
num_experts = len(token_fractions)
# Load imbalance factor
lif = num_experts * token_fractions.max().item()
# Coefficient of variation
mean_frac = 1.0 / num_experts
std_frac = token_fractions.std().item()
cv = std_frac / mean_frac
# Entropy (normalized)
probs = token_fractions + 1e-10 # Avoid log(0)
entropy = -(probs * torch.log(probs)).sum().item()
max_entropy = np.log(num_experts)
normalized_entropy = entropy / max_entropy
return {
"load_imbalance_factor": lif,
"coefficient_of_variation": cv,
"normalized_entropy": normalized_entropy,
"max_token_fraction": token_fractions.max().item(),
"min_token_fraction": token_fractions.min().item(),
}
metrics = compute_load_metrics(token_fractions)Load Balance Metrics: Load Imbalance Factor: 5.578 (ideal: 1.0) Coefficient of Variation: 1.929 (ideal: 0.0) Normalized Entropy: 0.414 (ideal: 1.0) Max/Min Token Fraction: 0.697 / 0.000
The load imbalance factor above 1.0 indicates that the most popular expert receives more than its fair share of tokens. The normalized entropy below 1.0 shows that the distribution is not uniform. These numbers represent a model that has not even started training, and the biases we introduced are small compared to what develops during actual training runs. This gives you a sense of how aggressive the rich-get-richer dynamics can be.
Simulating Training Dynamics
To see how imbalance develops during training, let's simulate the rich-get-richer dynamics by iteratively biasing the gating network toward popular experts. This simulation compresses the feedback loop that plays out over real training runs into a controlled demonstration.
def simulate_training_imbalance(
num_experts, num_steps, reinforcement_strength=0.1
):
"""Simulate how expert preferences evolve during training without balancing."""
# Start with uniform expert scores
expert_scores = torch.zeros(num_experts)
history = {"step": [], "fractions": [], "metrics": []}
for step in range(num_steps):
# Convert scores to probabilities
probs = F.softmax(expert_scores, dim=0)
# Sample token assignments based on probabilities
num_tokens = 1000
assignments = torch.multinomial(probs, num_tokens, replacement=True)
# Count tokens per expert
fractions = torch.zeros(num_experts)
for i in range(num_experts):
fractions[i] = (assignments == i).sum().float() / num_tokens
# Record history
history["step"].append(step)
history["fractions"].append(fractions.clone())
history["metrics"].append(compute_load_metrics(fractions))
# Simulate rich-get-richer: boost scores of popular experts
expert_scores += reinforcement_strength * (fractions - 1 / num_experts)
return history
# Run simulation
history = simulate_training_imbalance(
num_experts=8, num_steps=100, reinforcement_strength=0.4
)
The simulation demonstrates how quickly expert utilization diverges from uniform. After only 100 steps, the load imbalance factor has risen substantially, with some experts receiving several times their fair share of tokens while others receive almost none. In a real training run, this feedback loop operates over hundreds of thousands of gradient steps, giving it ample time to reach pathological extremes.
Visualizing Expert Collapse
Let's run a longer simulation with stronger reinforcement to observe expert collapse. We increase both the number of steps and the reinforcement strength to show what happens when no balancing mechanism intervenes.
# Simulate collapse with stronger dynamics
collapse_history = simulate_training_imbalance(
num_experts=8, num_steps=200, reinforcement_strength=0.3
)
The stacked area chart shows expert collapse in action. By the end of training, nearly all tokens are routed to just one or two experts. The other experts, despite consuming memory and having trainable parameters, contribute nothing to model outputs. Notice how the collapse happens in stages: initially, some experts grow while others shrink, then the dominant experts gradually squeeze out all competitors until the final collapsed state is reached.
Analyzing Router Entropy
Router entropy provides another view into load balance. Let's examine how the softmax output distribution changes as collapse progresses. Entropy is particularly valuable because it begins declining earlier than other metrics. This provides a leading indicator of collapse before it becomes obvious in the token fraction plots.
def compute_router_entropy(expert_scores):
"""Compute entropy of the router's softmax output."""
probs = F.softmax(expert_scores, dim=0)
entropy = -(probs * torch.log(probs + 1e-10)).sum()
return entropy.item()
def simulate_with_entropy(num_experts, num_steps, reinforcement_strength=0.2):
"""Track router entropy during simulated training."""
expert_scores = torch.zeros(num_experts)
entropies = []
max_entropy = np.log(num_experts)
for step in range(num_steps):
# Record entropy
entropy = compute_router_entropy(expert_scores)
entropies.append(entropy / max_entropy) # Normalized
# Sample and reinforce
probs = F.softmax(expert_scores, dim=0)
assignments = torch.multinomial(probs, 1000, replacement=True)
fractions = torch.zeros(num_experts)
for i in range(num_experts):
fractions[i] = (assignments == i).sum().float() / 1000
expert_scores += reinforcement_strength * (fractions - 1 / num_experts)
return entropies
entropy_history = simulate_with_entropy(
num_experts=8, num_steps=150, reinforcement_strength=0.5
)
The entropy plot clearly shows the progression toward collapse. Starting at maximum entropy (uniform routing), the entropy stays near its ceiling during the early noisy phase and then falls rapidly as the router concentrates probability mass on fewer experts. This metric provides an early warning signal: entropy begins dropping before token fractions show obvious skew. By monitoring entropy during training, you can detect the onset of imbalance and intervene before the situation becomes irreversible.
Looking at all three metrics together, we can see how they provide complementary information. The Load Imbalance Factor rises sharply once the dominant expert clearly separates from the pack, which makes it good for detecting moderate to severe imbalance. The Coefficient of Variation is sensitive to diffuse imbalance where many experts deviate from their ideal load, even if no single expert is dramatically dominant. Normalized Entropy provides the earliest warning signal, declining smoothly from the very first steps as the distribution begins to concentrate.