Load Balancing › The Cost of Imbalance in Distributed Training
Language AI HandbookAll chapters Load Balancing › The Cost of Imbalance in Distributed TrainingOne section · Sections in this chapter7 / 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
Load imbalance has particularly severe consequences when experts are distributed across devices. In expert parallelism, each expert resides on a separate accelerator, and the system's throughput is constrained by the slowest device. This is a basic result from parallel computing theory known as Amdahl's Law: the speedup from parallelization is limited by the fraction of work that cannot be parallelized, or in this case, by the time required for the most overloaded device. Let's quantify the efficiency loss that results from uneven workload distribution.
Think of the distributed training scenario as a team of runners in a relay race where all runners must cross the finish line before the next leg begins. If some runners are much faster than others, the fast runners finish early and wait while the slow runners complete their legs. The total time is determined by the slowest runner, not the average. Load imbalance in expert parallelism creates exactly this situation: the device hosting the overloaded expert is the slow runner, and all other devices wait for it to finish before training can proceed.
def compute_parallel_efficiency(token_fractions):
"""
Compute parallel efficiency when experts are on separate devices.
In expert parallelism, all devices must wait for the slowest one.
Efficiency is the ratio of useful compute to total compute time.
"""
num_experts = len(token_fractions)
# Each device's workload is proportional to its token fraction
# Time is determined by the most loaded device
max_load = token_fractions.max().item()
# Average useful work per device
avg_load = 1.0 / num_experts
# Efficiency: useful work / time spent
# Time = max_load, total useful work = avg_load * num_experts = 1.0
efficiency = avg_load / max_load if max_load > 0 else 0
return efficiency
# Compute efficiency for different imbalance levels
imbalance_levels = []
efficiencies = []
for step in range(0, len(collapse_history["fractions"]), 10):
fractions = collapse_history["fractions"][step]
lif = compute_load_metrics(fractions)["load_imbalance_factor"]
eff = compute_parallel_efficiency(fractions)
imbalance_levels.append(lif)
efficiencies.append(eff)
The relationship is stark: efficiency is inversely proportional to the load imbalance factor. When the most-loaded expert receives twice its fair share of tokens (LIF = 2), parallel efficiency drops to 50%. When one expert receives all tokens (LIF = ), efficiency approaches since only one device does useful work while the others wait.
Consider training a large MoE model with 64 experts distributed across 64 GPUs. At perfect balance, all 64 GPUs compute in parallel, giving approximately 64x speedup over a single GPU. If load collapses such that one expert handles 50% of tokens (LIF = 32), efficiency drops to about 3%, meaning that 62 out of 64 GPUs spend most of their time waiting. The effective throughput is barely better than what a single GPU could achieve. Load balancing is therefore a practical necessity for serious MoE training at scale.