Rethinking MoE Load Balancing as Feedback Control

원문 초록

Mixture-of-Experts (MoE) training is often not stationary: the data mixture changes at curriculum boundaries, when a model is continually pretrained into a new domain, or when the corpus is refreshed. Routing is learned against the current mixture, so each such change disrupts the per-expert load the router had converged to, and existing balancing methods have only been evaluated on stationary data. We rethink load balancing as a feedback control problem, in which the deviation of each expert's load from its fair share is the control error and the per-expert routing bias is the control input, and define a family of controllers, the PID-Balancer, from the proportional, integral, and derivative terms of classical PID control. Two widely used methods belong to this family: the auxiliary balancing loss behaves like a proportional controller, and the loss-free bias update is an integral controller acting on the sign of the error. Reading the router as a static map with a one-step delay, we show that the integral term is the only one that both removes persistent imbalance and averages out the step-to-step fluctuation of the measured load, whereas the proportional and derivative terms pass or amplify that fluctuation into the routing bias. We therefore adopt PID (I), the pure-integral member of the family with a single gain. Two settings are evaluated: (i) training an OLMoE backbone from scratch with 16 to 64 experts per layer, through a mid-run mixture switch, and (ii) continual pretraining of the released OLMoE-1B-7B on a new corpus. PID (I) holds balance before the shift and restores it more tightly afterwards than either the auxiliary loss or loss-free balancing, at comparable perplexity.