Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

TL;DR

This work addresses the instability of reinforcement learning with Mixture-of-Experts by identifying training–inference routing discrepancies as a key root cause. It introduces Rollout Routing Replay (R3), a simple yet effective mechanism that reuses inference-time routing masks during training to align routing behavior, while preserving gradient flow and keeping rollout overhead under $3\%$. Empirical results show substantial reductions in training–inference $D_{KL}$ divergence and extreme token discrepancies, leading to improved stability and performance over GSPO and TIS across multiple RL tasks, including multi-turn dialogues. The approach is compatible with both on-policy and mini-batch off-policy training and can be integrated with routing mask caching for efficiency, offering a practical solution for stabilizing RL in MoE models with broad applicability to real-world LLM workflows.

Abstract

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.

Figures (10)

• Figure 1: Left: Illustration of the Rollout Routing Replay (R3). Top right: Training and inference discrepancies before and after applying R3. Bottom right: Reinforcement learning training performance before and after applying R3.
• Figure 2: (a): Illustration of the training-inference discrepancy in the MoE model. (b): Illustration of the training-inference discrepancy in the MoE+R3 model. (c): Illustration of the training-inference discrepancy in the Dense model. (d): Extreme Token Distribution Function, calculated based on Equation \ref{['eq:etdf']}.
• Figure 3: Router discrepancy analysis
• Figure 4: Probabilities obtained by performing forward propagation twice using the Megatron
• Figure 5: Analysis of training–inference collapse. The plot shows the estimated training–inference KL divergence and the extreme token distribution function $F(\tau=2)$ (Eq. \ref{['eq:etdf']}) at each training step.