Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

作者/机构: Wenhan Ma†‡∗, Hailin Zhang‡, Liang Zhao‡, Yifan Song†‡, Yudong Wang†‡, Zhifang Sui†⋄, Fuli Luo⋄
†State Key Laboratory of Multimedia Information Processing, School of Computer Science, Peking University ‡LLM-Core Xiaomi

A1 主要贡献

本文研究了在混合专家（MoE）模型中，强化学习（RL）训练的不稳定性问题。

核心问题与研究目标:
* 问题: 在大型语言模型（LLM）的强化学习训练中，尤其是在MoE模型中，路由机制常引入不稳定性，甚至导致灾难性的训练崩溃。
* 分析: 作者发现MoE模型的训练阶段和推理阶段之间存在显著的路由行为差异。这种不一致性，加上即使在相同条件下，路由框架在重复前向传播中也可能产生不同的专家选择，是导致不稳定的根本原因。
* 目标: 解决MoE模型在强化学习中的训练-推理不一致性问题，从而稳定训练过程。

创新点与主要贡献:
1. 系统性地识别和分析了MoE模型中训练与推理之间的路由分布差异，并强调了这些差异在训练不稳定性中的关键作用。
2. 提出了“部署路由回放”（Rollout Routing Replay, R3）方法。该方法通过在推理（Rollout）阶段捕获路由分布，并在训练阶段直接回放这些分布来对齐训练和推理过程中的路由行为。如图1左侧所示，R3记录推理引擎的路由选择，并将其应用于训练引擎，从而显著缩小训练与推理之间的策略差异。
3. 通过实验验证了R3的有效性。R3显著降低了训练-推理策略的KL散度（图1右上），并减少了极端差异的出现。在多个RL设置（包括多/单mini-step和基础/SFT模型）中，R3在稳定性和整体性能方面均优于现有方法如GSPO和TIS，有效防止了训练崩溃（图1右下）。

A3 背景知识与关键发现

2 预备知识

符号表示。本文考虑一个由$\theta$参数化的自回归语言模型，表示为策略$\pi_{\theta}$，它从查询$x \in D$生成响应$y$。序列的似然由$\pi_{\theta}(y|x) = \prod_{|y|}^{t=1} \pi_{\theta}(y_t|x, y_{<t})$给出，其中$|y|$是序列长度。$\pi_{infer}$和$\pi_{train}$分别表示在推理引擎和训练引擎中运行的策略。</p>

近端策略优化 (PPO)。PPO【Schulman等人，Proximal policy optimization algorithms，2017，arXiv】是策略优化的基石算法。对于给定的查询$x$，PPO通过最大化以下目标来更新策略$\pi_{\theta}$：
$$\mathcal{J}_{\text{PPO}}(\theta) = \mathbb{E}_{s_t, a_t \sim \pi_{\text{old}}, \omega \sim \mathcal{U}(0,1)} \left[ \left| \frac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} \min \left( u_i(\theta)\hat{A}_i, \text{clip}(u_i(\theta), 1-\epsilon, 1+\epsilon)\hat{A}_i \right) \right| \right]$$
序列$y$中token $y_t$的重要性采样比率$w_t(\theta)$定义为：
$$ w_t(\theta) = \frac{\pi_{\text{train}}(\theta)(y_t|x, y_{<t})}{\pi_{\text{train}}(\theta_{\text{old}})(y_t|x, y_{<t})} $$<br /> token $y_t$的优势$\hat{A}_t$通常由一个独立的价值模型估计，而$\epsilon$是重要性比率的裁剪范围。为简洁起见，本文省略了KL正则化项。

训练-推理不一致性问题。一个关键的不一致性源于使用独立的引擎进行部署（rollout）和训练的普遍做法，其中数据通过推理策略（$\pi_{infer}$）采样，但损失函数使用训练策略（$\pi_{train}$）计算，如方程1所示。这种策略不匹配导致了强化学习中的训练不稳定性。在MoE模型中，作者认为这个问题主要源于路由器的不一致性。本文提出的解决方案旨在缓解这个根本问题，使其具有广泛的适用性，并且与近期的策略优化框架如GRPO【Shao等人，Deepseekmath: Pushing the limits of mathematical reasoning in open language models，2024，arXiv】、GSPO【Zheng等人，Group sequence policy optimization，2025，arXiv】和DAPO【Yu等人，Dapo: An open-source llm reinforcement learning system at scale，2025，arXiv】正交且兼容。

3 训练-推理差异

强化学习框架中的训练-推理差异经常导致训练不稳定和模型崩溃。本节将证明，在MoE模型中，这种策略不匹配被显著放大，主要源于不一致的路由分布。此外，作者观察到，即使是相同训练框架的多次运行也可能产生不同的token概率，进一步加剧了RL训练的不稳定性。

3.1 MoE模型中训练与推理的策略差异

实验设置。作者使用MoE模型Qwen3-30B-A3B【Yang等人，Qwen3 technical report，2025，arXiv】来分析训练和推理引擎之间的策略差异。首先，使用SGLang推理引擎为2048个数学问题生成响应，并保存每个生成token的概率，共产生约2000万个响应token。然后，将这些响应输入Megatron以获得训练引擎分配的相应概率。作者使用多种度量来量化这两个概率分布之间的差异。

KL散度估计。设$T$为所有响应token的集合，作者使用Schulman【Schulman，Approximating kl divergence，2020，http://joschu.net/blog/kl-approx.html】提出的$k_3$方法来估计训练和推理概率分布之间的KL散度 ：
$$D_{\mathrm{KL}}(\pi_{\mathrm{train}}(\theta)||\pi_{\mathrm{infer}}(\theta)) = \frac{1}{|T|} \sum_{t \in T} \left[ \frac{\pi_{\mathrm{train}}(\theta)(t)}{\pi_{\mathrm{infer}}(\theta)(t)} - 1 - \log \frac{\pi_{\mathrm{train}}(\theta)(t)}{\pi_{\mathrm{infer}}(\theta)(t)} \right]$$
计算结果显示，Qwen3-30B-A3B (MoE)的估计KL散度为$1.535 \times 10^{-3}$，而Qwen3-8B (密集模型基线)的KL散度为$6.4 \times 10^{-4}$。

可视化分析。为了可视化MoE模型的训练-推理差异，作者随机采样了1000万个响应token，并绘制了散点图，x轴为SGLang概率，y轴为Megatron概率。点围绕$y = x$线的集中程度表示一致性程度。如图2a和2c所示，与Qwen3-8B相比，Qwen3-30B-A3B的散点带要宽得多，揭示了更大的训练-推理差异。

极端Token分布分析。为了量化模型在训练和推理期间行为的差异，作者引入了极端Token分布函数$F(\tau)$，定义为：
$$F(\tau)=\frac{1}{|T|} \sum_{t \in T} \mathbb{I}\left[\max \left(\frac{\pi_{\text {train }}(\theta)(t)}{\pi_{\text {infer }}(\theta)(t)}, \frac{\pi_{\text {infer }}(\theta)(t)}{\pi_{\text {train }}(\theta)(t)}\right)>\tau\right].$$
该函数测量了极端token的比例——即训练分布$\pi_{train}(t)$和推理分布$\pi_{infer}(t)$之间的概率比率超过阈值$\tau$的token。图2d绘制了该函数$F(\tau)$随阈值$\tau$的变化。该图显示，对于$\tau > 2$，Qwen3-30B-A3B模型中的极端token比例比Qwen3-8B模型大一个数量级。这一显著差距表明MoE模型中存在高得多的训练-推理可变性。

3.2 MoE模型中训练与推理的路由差异

路由的非连续性是关键。从函数连续性的角度看，MoE模型和密集模型之间的关键区别在于路由引入的非连续性。在MoE模型中，路由器输入的微小扰动可能导致选择完全不同的专家，从而引起层输出的巨大变化。而密集模型没有显式的专家选择过程，不会表现出这种现象。

路由分布差异分析。基于此，作者进一步分析了MoE模型中训练和推理之间的路由分布差异。实验使用SGLang【Zheng等人，Sglang: Efficient execution of structured language model programs，2024，arXiv】和Qwen3-30B-A3B模型【Yang等人，Qwen3 technical report，2025，arXiv】为2048个数学问题生成响应。对于每个响应，从推理引擎收集所有token（包括输入token）的路由分布。然后，将这些序列输入Megatron引擎进行前向传播，得到训练引擎观察到的路由分布。作者从不同层面比较了这两组路由信息：
* 路由器层面比较：对每个token和每个MoE层，统计MoE路由器做出不同专家选择的数量，并计算这些差异的频率。图3a显示，大约10%的路由器在训练时选择了与推理时不同的专家。
* Token层面比较：对每个token，统计其在所有层中MoE路由器做出不同专家选择的数量，并计算这些情况的频率。图3b显示，94%的token在至少一个层的前向传播中选择了不同的专家。
* 序列层面比较：对于一个序列，计算每个token的路由分布差异，然后在所有token上取平均，得到每个序列的平均差异，并绘制直方图（图3c）。结果显示，每个token的平均差异约为6个路由器。
这些发现表明，在训练和推理过程中，MoE模型表现出路由分布的差异。

3.3 同一框架内同一序列重复前向传播的差异

实验与发现。作者在Megatron框架下对同一组序列进行了两次前向传播，得到了两个概率分布。遵循3.1节的流程，计算了这些分布之间的KL散度并绘制了结果（散点图见图4，KL散度为$8.4 \times 10^{-4}$）。

对RL的影响。结果表明，对于MoE模型，在Megatron引擎中，即使输入序列完全相同，两次前向传播的最终输出概率也可能不同。在强化学习设置中，这种变化为旧策略$\pi_{train}(\theta_{old})$的计算增加了噪声。这种噪声使得重要性采样比率不可靠，可能导致强化学习过程不稳定甚至中断。

A2 方法细节

本节详细描述了部署路由回放（Rollout Routing Replay, R3）的实现、其在多轮对话中的缓存支持，以及其对训练-推理差异影响的分析。

4.1 实现

常规MoE前向传播。首先描述训练框架内MoE层的常规前向传播过程。考虑序列在第$t$个token，位于第$l$个Transformer块的MoE层。设该层在训练期间的输入为$x_{train}$。路由器logits计算如下：
$S_{\text{train}} = \mathbf{X}_{\text{train}}\mathbf{W}_{r},$
其中$W_r$表示路由器的线性权重矩阵。设专家数量为$M$，要选择的专家数量为$K$。训练期间，路由器根据logits选择top-K个专家，这由一个二进制掩码表示：
$I_{\text{train}} = \text{TopKMask}(s_{\text{train}}, K),$
其中$I_{train} \in \{0, 1\}^M$且$\sum_i I_{train,i} = K$。然后，通过对所选专家的logits应用softmax来产生门控权重：
$$ g_{\text {train }, i}=\frac{I_{\text {train }, i} \exp \left(s_{\text {train }, i}\right)}{\sum_{j=1}^{M} I_{\text {train }, j} \exp \left(s_{\text {train }, j}\right)} \quad \text { for } i=1, \ldots, M . $$
最后，MoE层的输出计算为专家输出的加权和：
$$y_{\text{train}} = \sum_{i=1}^{M} g_{\text{train},i} \mathcal{E}_i(\mathbf{x}_{\text{train}}),$$
其中$\mathcal{E}_i(\cdot)$表示第$i$个专家网络。

部署路由回放（R3）。现在介绍R3。假设在推理阶段，MoE层的输入为$x_{infer}$。路由器计算推理logits $s_{infer} = x_{infer}W_r$，并从中获得路由掩码$I_{infer} = \text{TopKMask}(s_{infer}, K)$。R3的核心思想是在训练的前向传播中重用推理时的路由掩码$I_{infer}$，但仍然对训练时的logits应用softmax以保留梯度流。具体来说，在这条回放路径上，“回放”门控权重$g_{replay}$的计算方式如下：
$$\delta_{\text{replay},i} = \frac{l_{\text{infer},i} \exp(s_{\text{train},i})}{\sum_{j=1}^M l_{\text{infer},j} \exp(s_{\text{train},j})} \quad \text{for } i=1,\dots,M.$$
然后，这些回放权重被用来组合训练专家的输出，产生回放输出$y_{replay}$：
$$ \mathbf{y}_{\text{replay}} = \sum_{i=1}^{M} g_{\text{replay},i} \mathcal{E}_i(\mathbf{x}_{\text{train}}). $$
这种设计有两个主要目的：(a) 对齐训练和推理：使用$I_{infer}$确保了训练期间使用的专家与推理期间选择的专家相匹配，消除了专家选择上的不匹配。(b) 保留梯度数据流：通过仅回放掩码，梯度仍然可以流回logits而不干扰计算图，这有助于有效地优化路由器。

4.2 路由器掩码缓存与多轮对话支持

利用KVCache机制。许多推理引擎使用KVCache前缀缓存策略【Kwon等人，Efficient memory management for large language model serving with pagedattention，2023，SOSP】来避免对先前看过的上下文进行冗余的预填充计算，从而显著减少多轮交互中的总计算量。作者观察到，缓存的路由器掩码具有相似的属性：对于相同的前缀token，MoE路由器应产生相同的结果。因此，来自推理引擎的路由掩码$I_{infer}$可以与前缀KVcache一起被缓存。

实现与优势。具体来说，对于每一层和每个token前缀，相应的路由掩码与KVCache一起存储。当相同的前缀出现并命中缓存时，掩码可以被重用，无需重新计算。这使得R3能够与前缀缓存机制无缝集成。在智能体（agent）场景中，缓存路由掩码尤其有益。许多智能体任务，如软件工程【Jimenez等人，Swe-bench: Can language models resolve real-world github issues?，2024】和网页浏览【Wei等人，Browsecomp: A simple yet challenging benchmark for browsing agents，2025，arXiv】，涉及自回归生成和工具调用之间的多轮交互。为了提高效率，这些过程直接重用前几轮的KVCache。路由掩码缓存使R3在RL智能体任务中保持高效，无需重新预填充以生成路由掩码，这对于训练大规模、先进的MoE模型至关重要。

4.3 R3对训练-推理差异的经验性分析

实验评估。为了评估R3在减少训练-推理差异方面的有效性，作者使用Qwen3-30B-A3B模型重复了3.1节中描述的过程。在此过程中，作者缓存在SGLang上推理时获得的路由分布，并在Megatron框架内进行回放。在获得推理引擎概率和训练引擎概率后，使用公式2估计了训练和推理之间的KL散度。

结果分析。结果显示，应用R3后，训练和推理之间的KL散度从$1.5 \times 10^{-3}$降低到$7.5 \times 10^{-4}$，这接近于密集模型观察到的$6.4 \times 10^{-4}$。这直观地表明训练-推理差异的减少。作者还用R3绘制了训练-推理差异比率的累积分布图（图2d）。该图表明，对于MoE模型，应用R3将具有较大训练-推理差异的token频率降低了一个数量级。

A4 实验环境

任务与数据集:
* 任务: 数学推理。
* 训练数据集: 收集并筛选了来自多个开源数据集（包括BigMath【Albalak等人，2025】、ORZ【Hu等人，2025】等）的问题，共约10万个可验证的数学问题。
* 评估数据集: AIME24, AIME25, AMC23, 和 MATH500 (level 5)。
* 评估指标: AIME24和AIME25报告Avg@32，AMC23报告Avg@16，MATH500 (level 5)报告Avg@4。为了公正评估，记录每个训练运行中每5个全局步骤的模型性能，并报告观察到的最高性能及其对应的训练步骤。

模型架构:
* Qwen3-30B-A3B-Base【Yang等人，2025】。
* Qwen3-30B-A3B-SFT：在通用指令遵循数据集上由Base模型微调而来。

硬件配置:
* 未明确指定GPU型号/数量/平台、网卡型号/数量、CPU架构/型号、连接关系。

软件配置:
* 框架: 使用VeRL【Sheng等人，2024】框架实现R3，训练使用Megatron【Shoeybi等人，2019】，推理使用SGLang【Zheng等人，2024】。
* 训练参数:
* 基线方法:
* GRPO【Shao等人，2024】+ DAPO【Yu等人，2025】的Clip Higher技术（$\epsilon_{low} = 0.2$, $\epsilon_{high} = 0.27$）。
* TIS【Yao等人，2025】，上裁剪阈值$\eta = 2$。
* GSPO【Zheng等人，2025】，序列级重要性采样（$\tau_{low} = 3 \times 10^{-4}$, $\tau_{high} = 4 \times 10^{-4}$）。
* 训练设置:
* 多Mini-Step: mini step数为8，每个mini step训练256个样本，学习率为$1 \times 10^{-6}$。R3在此设置下重计算旧策略和更新策略时都回放路由。
* 单Mini-Step: mini step数为1，所有2048个样本一次性更新，学习率提高到$3 \times 10^{-6}$。此场景不重计算旧策略。
* 其他: Batch size为256，全局batch size为2048 ($K=8$)。最大提示长度2048，最大生成长度30720。采用Dynamic Sampling策略【Yu等人，2025】。不引入专家平衡的辅助损失。

A4 实验结果

主要评估结果如表1所示。

整体性能
* R3表现优异: 在不同场景下，R3均取得了更好的结果。
* 在多mini-step设置中，GRPO+R3的平均分比GSPO高1.29。将R3与GSPO结合，性能进一步提升0.95分。
* 在单mini-step设置中，R3在SFT模型上比TIS高5.58分，在base模型上高1.51分。
* R3与TIS的组合: 将R3与TIS结合并没有带来明显收益，甚至可能降低性能。例如，在SFT模型的单mini-step设置中，TIS+R3比单独使用R3低1.69分。这可能是因为R3已经显著减少了训练和推理之间的策略差异，TIS的额外校正带来的好处微乎其微。

训练稳定性
* R3防止崩溃: 在单mini-step设置中，三种未使用R3的强化学习过程在训练中都发生了崩溃。
* 崩溃原因分析 (图5): 作者绘制了每个训练步骤的训练-推理KL散度和极端token分布函数$F(\tau=2)$。观察到，在每个训练过程中，KL散度和$F(\tau=2)$的值都随训练进程增加。崩溃的训练运行几乎总是伴随着异常高的KL和$F(\tau=2)$值。例如，使用GRPO在单mini-step设置下训练SFT模型时，60个全局步骤后，$F(\tau=2)$的值超过0.1，意味着10%的token在训练和推理框架下的概率差异超过2倍。
* R3的稳定效果: 相比之下，使用R3的训练过程中，$F(\tau=2)$的值大部分时间保持在$10^{-4}$以下。通过对齐训练和推理分布，R3有效地稳定了MoE模型的强化学习。

优化和生成行为 (图6)
* R3提升优化稳定性: R3在训练过程中展现出更稳定的优化、探索行为和生成动态。在单mini-step + base模型的实验组中，R3具有更小的梯度范数、更平滑的序列长度增长模式和更稳定的熵。
* 序列长度: 使用R3时，生成的序列长度在训练初期迅速增长，表明R3能快速捕捉到正确的优化方向。
* 梯度范数: R3始终保持较低的梯度范数，表明优化过程更稳定。
* 生成熵: 使用R3时，熵在大约25个步骤后开始稳定增加，表明模型更早开始探索更好的策略。

A5 结论

本文确定了训练-推理过程中的路由差异是导致MoE模型强化学习不稳定的主要原因。为解决此问题，作者提出了部署路由回放（R3）方法，该方法在训练期间重用推理时的路由分布，以对齐专家选择，同时保留梯度流。在多个RL设置下的实验表明，R3显著减少了训练-推理的差异，稳定了训练过程，并且在性能上持续优于现有方法。研究结果证明了在MoE模型中对齐训练和推理的重要性，并表明R3为提高其稳定性提供了一个实用的解决方案。

A6 附录

A 详细评估结果

B 详细训练指标