Load Balancing › Key Parameters

Understanding the key parameters that govern load balance helps you configure MoE systems effectively and diagnose problems when they arise. These parameters interact in ways that require careful consideration:

• num_experts: The total number of experts in the MoE layer. More experts increase model capacity but make balancing harder because the probability mass must be spread across more options. With 8 experts, each expert's fair share is 12.5%; with 64 experts, it is only 1.56%. Small absolute deviations become more significant in proportional terms as the number of experts grows.

• capacity_factor: A scalar multiplier (typically 1.0 to 1.25) determining the maximum tokens an expert can process relative to uniform load. Lower values enforce stricter balance at the cost of more dropped tokens. Higher values accept more imbalance to reduce token dropping. The right value depends on the severity of your load imbalance and your tolerance for dropped tokens.

• reinforcement_strength: In our simulation, this parameter models the magnitude of the feedback loop. In real training, this corresponds to how strongly the gating network's gradients reinforce existing preferences, which depends on learning rate, batch size, and the auxiliary loss coefficient. Higher reinforcement leads to faster collapse; higher auxiliary loss coefficients counteract this by penalizing concentration.

Choosing these parameters well requires monitoring the load balance metrics we've discussed. Start by establishing baselines during initial training runs, set alerting thresholds based on those baselines, and tune the capacity factor and auxiliary loss coefficient until the metrics stay within acceptable ranges throughout training.