Stabilisierung von MoE-Reinforcement Learning durch Abgleich von Trainings- und Inferenz-Routern

Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

October 13, 2025

papers.authors: Wenhan Ma, Hailin Zhang, Liang Zhao, Yifan Song, Yudong Wang, Zhifang Sui, Fuli Luo

cs.AI

papers.abstract

Reinforcement Learning (RL) hat sich als entscheidender Ansatz zur Verbesserung der Fähigkeiten großer Sprachmodelle erwiesen. Bei Mixture-of-Experts (MoE)-Modellen führt der Routing-Mechanismus jedoch häufig zu Instabilität, die sogar zu einem katastrophalen Zusammenbruch des RL-Trainings führen kann. Wir analysieren die Trainings-Inferenz-Konsistenz von MoE-Modellen und identifizieren eine signifikante Diskrepanz im Routing-Verhalten zwischen den beiden Phasen. Darüber hinaus kann das Routing-Framework selbst unter identischen Bedingungen bei wiederholten Vorwärtsschritten divergierende Expertenauswahlen liefern. Um diese grundlegende Inkonsistenz zu adressieren, schlagen wir Rollout Routing Replay (R3) vor, eine Methode, die Routing-Verteilungen aus der Inferenz-Engine aufzeichnet und sie während des Trainings wiedergibt. R3 reduziert die KL-Divergenz der Trainings-Inferenz-Policy erheblich und mildert extreme Diskrepanzen, ohne die Trainingsgeschwindigkeit zu beeinträchtigen. Umfangreiche Experimente in verschiedenen Settings bestätigen, dass R3 das RL-Training erfolgreich stabilisiert, Zusammenbrüche verhindert und Methoden wie GSPO und TIS übertrifft. Wir sind überzeugt, dass diese Arbeit eine neue Lösung zur Stabilisierung von RL in MoE-Modellen bieten kann.

English

Reinforcement learning (RL) has emerged as a crucial approach for enhancing the capabilities of large language models. However, in Mixture-of-Experts (MoE) models, the routing mechanism often introduces instability, even leading to catastrophic RL training collapse. We analyze the training-inference consistency of MoE models and identify a notable discrepancy in routing behaviors between the two phases. Moreover, even under identical conditions, the routing framework can yield divergent expert selections across repeated forward passes. To address this foundational inconsistency, we propose Rollout Routing Replay (R3), a method that records routing distributions from the inference engine and replays them during training. R3 significantly reduces training-inference policy KL divergence and mitigates extreme discrepancies without compromising training speed. Extensive experiments on various settings confirm that R3 succeeds in stabilizing RL training, preventing collapse and outperforming methods such as GSPO and TIS. We believe this work can offer a new solution for stabilizing RL in MoE models.