解决MoE模型强化学习训练崩溃问题,让训练与推理路由一致
Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers
- 记录推理时的路由分布,在训练中重放以保持一致性
- 使训练与推理的策略差异降低60%以上,避免训练崩溃
- 适用于大模型强化学习,尤其适合需要稳定训练的MoE架构
强化学习(RL)已成为提升大语言模型能力的关键方法。然而,在混合专家(MoE)模型中,路由机制常导致训练不稳定,甚至引发灾难性训练崩溃。我们分析了MoE模型在训练与推理阶段的路由一致性,发现两者存在显著差异;即使在相同条件下,路由框架在多次前向传播中也会产生不同的专家选择。为解决这一根本性不一致,我们提出滚动路由重放(R3)方法:记录推理引擎的路由分布,并在训练中重放。R3显著降低了训练-推理策略的KL散度,有效缓解极端差异,且不牺牲训练速度。在多种设置下的大量实验表明,R3成功稳定了强化学习训练,防止崩溃,并优于GSPO和TIS等方法。我们认为该工作为稳定MoE模型中的强化学习提供了新思路。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。