Source-linked AI summary
Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers
Wenhan Ma, Hailin Zhang, Liang Zhao, Yifan Song, Yudong Wang, Zhifang Sui, Fuli Luo
TL;DR
MoE reinforcement learning can become unstable because separate inference and training engines produce inconsistent routing and policies, with existing approaches not fully resolving the resulting off-policy issue. The paper proposes Rollout Routing Replay (R3), which reuses inference-time routing distributions during training; across multiple settings, it reduces divergence, stabilizes training, and outperforms existing methods.
Problem
Separate inference and training engines create policy and routing discrepancies in MoE reinforcement learning, while existing approaches do not fully resolve the intensified off-policy issue.
Method
R3 captures inference-time routing distributions during sequence generation and replays them during training to align expert selection while preserving gradient flow.
Results
Across multiple RL settings, R3 substantially reduces training–inference divergence, stabilizes training, and consistently outperforms existing methods.
Takeaways & Limitations
Aligning training and inference in MoE models provides a practical solution for improving RL stability.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper identifies training–inference routing discrepancies as a source of MoE reinforcement-learning instability and proposes Rollout Routing Replay (R3) to align them. Across multiple RL settings, R3 reduces discrepancies and improves stability and performance relative to existing approaches.
- Motivation: Separate inference and training engines can produce divergent token probabilities, contributing to catastrophic RL collapse.Existing remedies include importance sampling and specialized kernels, but the paper states they do not fully resolve the intensified off-policy issue in MoE models.
- Motivation: MoE routing distributions are identified as a pivotal instability factor because dynamic expert selection increases training–inference policy discrepancies.The paper frames routing distributions as the root cause to address rather than discarding data with large discrepancies.
- Method: R3 captures inference-engine routing distributions during sequence generation and replays them in the training engine to align routing behavior.This approach targets the routing mechanism directly while preserving the distinction between inference and training phases.
- Effects: R3 reduces the number of tokens with significant training–inference probability discrepancies by approximately an order of magnitude.The reduction is associated with a substantial decrease in KL divergence between logits produced by the two engines.
- Results: Across multi-/single-mini-step and Base/SFT settings, R3 outperforms GSPO and TIS in stability and overall performance.The experiments cover multiple RL settings and report consistent improvements in the method comparison.
2 Preliminaries
This section formalizes autoregressive policy likelihoods, PPO optimization, and the policy mismatch created when rollout data comes from inference while training losses use a separate policy. In MoE models, the mismatch mainly stems from router inconsistency and can destabilize RL.
- Notation: An autoregressive policy π_θ generates response y from query x, with sequence likelihood factorized over token-level conditional probabilities.The sequence length is |y|, and each token probability conditions on x and preceding response tokens.
- PPO: PPO updates π_θ by maximizing a clipped objective based on token-level importance-sampling ratios and estimated advantages.The importance-sampling ratio compares current and old training-policy probabilities; the advantage is estimated by a separate value model.
- Policy mismatch: Using separate rollout and training engines means data is sampled with π_infer while the loss is computed with π_train, creating a policy mismatch.The paper states that this mismatch causes training instability in reinforcement learning and mainly stems from router inconsistency in MoE models.
- Policy mismatch: Variation in the old training policy adds noise to the importance-sampling ratio, which can destabilize or break the RL process.This connects router-related nondeterminism to the reliability of policy updates.
3 Training-Inference Discrepancies
MoE models exhibit larger and structurally distinct training–inference discrepancies than dense models, driven substantially by inconsistent routing decisions. These discrepancies also arise across repeated forward passes and can destabilize RL through unreliable policy updates.
- 1.535×10^-3 versus 6.4 × 10^-4: Qwen3-30B-A3B has higher estimated training–inference KL divergence than the Qwen3-8B dense baseline.
- Qwen3-30B-A3B shows a much wider probability scatter band than Qwen3-8B, indicating greater training–inference discrepancy.
- For τ>2, Qwen3-30B-A3B has an extreme-token fraction an order of magnitude larger than Qwen3-8B.The Extreme Token Distribution Function measures the proportion of tokens whose training-to-inference probability ratio exceeds τ.
- Approximately 10% of routers select different experts between training and inference, while 94% of tokens differ in at least one layer.
- The mean routing difference is approximately 6 routers per token, demonstrating substantial routing-distribution mismatch.
- Repeated Megatron forward passes can produce different probabilities for identical sequences, adding noise to the old policy and destabilizing importance sampling.The reported KL divergence between repeated forward-pass distributions is 8.4 × 10^-4.
4 Rollout Routing Replay
Rollout Routing Replay reuses inference-time expert-selection masks during training while retaining training-logit gradients. This aligns routing across engines, reduces probability discrepancies, and adds less than 3% rollout latency.
- The conventional MoE forward pass selects the top-K experts from router logits and combines their outputs using gating weights.
- R3 replays the inference routing mask during training while applying softmax to training logits, preserving gradient flow.
- Using the inference mask aligns the experts selected during training replay with those selected during inference.
- Below 3%: storing and retrieving router masks adds less than 3% latency during rollout.
- Caching routing masks with prefix KVCache preserves R3 efficiency for repeated prefixes and multi-turn agent interactions.
- 1.5×10^-3 to 7.5 × 10^-4: R3 reduces Qwen3-30B-A3B training–inference KL divergence toward the dense-model value of 6.4 × 10^-4.The accompanying analysis reports an order-of-magnitude reduction in tokens with large discrepancies for the MoE model.
5 Experiments
The experiments evaluate R3 across MoE reinforcement-learning settings, comparing its performance and stability with existing methods. R3 improves scores, reduces training–inference inconsistency, and prevents observed collapses without slowing rollout.
- Experimental setup: R3 caches rollout routing and replays it during old-policy recomputation and policy updating, alongside comparisons with GRPO, TIS, and GSPO.The implementation uses VeRL, Megatron for training, and SGLang for inference.
- Overall performance: 1.29 points: GRPO+R3 outperforms GSPO in the multi-mini-step setting, while GSPO+R3 adds 0.95 points over GSPO.In the single-mini-step setting, R3 exceeds TIS by 5.58 points on the SFT model and 1.51 points on the base model.
- Training stability: Three single-mini-step runs without R3 collapse, whereas R3 keeps F(τ=2) below 10^-4 for most of training.Without R3, the SFT model trained with GRPO exceeded F(τ=2)=0.1 after 60 global steps, indicating more than twofold probability differences for 10% of tokens.
- Optimization and generation behavior: R3 produces smaller gradient norms, smoother sequence growth, and more stable entropy during single-mini-step training with the base model.Without R3, entropy increases later and fluctuates heavily.
- Multi-turn task reinforcement learning: On the multi-turn software-engineering task, GRPO without R3 collapses after around 90 steps, while GRPO+R3 remains stable and reaches Pass@1 38.6.The score is 6.8 points higher than GRPO, and Router Mask Caching avoids slowing rollout speed.
6 Related Works
The related-work discussion situates MoE RL instability in routing inconsistencies between rollout and training. R3 addresses this by recording inference-time routing decisions and replaying them during training.
- Background: MoE models scale capacity efficiently by sparsely activating specialized experts, while RLVR refines reasoning, mathematics, and coding abilities.The discussion connects these model and training trends to the instability problem addressed by the paper.
- Routing inconsistency: Router sensitivity and framework differences can assign different experts to the same response during rollout and training.The paper identifies these inconsistent assignments as a source of heightened MoE RL instability.
- Compared methods: Rollout Routing Replay records inference-time routing decisions and replays them during training to enforce consistency without computational overhead.This is contrasted with Recompute Routing Replay, which caches routing during recomputation and replays it only during updating.
7 Conclusion
The conclusion identifies training–inference routing discrepancies as the primary instability source in MoE reinforcement learning. It presents R3 as a practical alignment method supported by results across multiple RL settings.
- Conclusion: R3 reuses inference-time routing distributions during training to align expert selection while preserving gradient flow.The method is designed to address the identified routing discrepancy directly.
- Conclusion: Across multiple RL settings, R3 reduces training–inference divergence, stabilizes training, and consistently outperforms existing methods.The conclusion emphasizes alignment between training and inference as important for MoE reinforcement learning.
A Detailed Evaluation Results and Training Metrics
The appendix contains detailed evaluation results and training metrics for the experiments. The referenced figures cover evaluation outcomes and training behavior.
- Detailed evaluation results: Figure 7 provides detailed evaluation results for the experiment in Section 5.
- Training metrics: Figure 8 provides detailed training metrics for the experiment in Section 5.
B Additional Experiments of Reasoning SFT Model
The experiments evaluate R3 during reinforcement learning on a reasoning-SFT version of Qwen3-30B-A3B. Without R3, training collapses, whereas R3 keeps training robust and stabilizes several training metrics.
- Experimental Setup: Qwen3-30B-A3B-Base was fine-tuned on the Mixture-of-Thoughts 5 long-reasoning dataset to produce a reasoning-SFT model for reinforcement learning experiments.The study compares GSPO and GRPO with GRPO+R3 under specified learning-rate and mini-step settings.
- Training Stability: GSPO without R3 collapses at training step 90 when mini_step = 4.
- Training Stability: GRPO without R3 collapses at training step 40 when mini_step = 1.
- Training Stability: Training with R3 remains robust and shows more stable sequence-length growth, entropy, and gradient norms.
C Detailed Training Metrics of Multi-Turn Reinforcement Learning
Figure 10 presents detailed training metrics and evaluation results for the RL training process of Qwen3-30B-A3B on the SWE task.
- Training Metrics: Figure 10 reports detailed training metrics for the RL training process on the SWE task.
- Evaluation Results: Figure 10 reports evaluation results for the RL training process on the SWE task.
- Model: The reported model is Qwen3-30B-A3B.