Swift MoE 模型 GRPO 训练中的 Router ReplayR2/R3原理、参数与源码实现解析【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift在 Swift 的 MoE 模型强化学习GRPO训练中训练引擎Megatron与推理引擎vLLM之间的路由不一致会显著放大训练-推理偏差甚至导致训练崩溃。Router Replay 通过在训练前向中回放固定的路由掩码来消除这一偏差按掩码来源分为 R2Vanilla Routing Replay与 R3Rollout Routing Replay两种策略。读完本篇你将理解三个策略训练/旧/采样之间的路由失配如何分解进重要性采样比掌握--router_replay_mode的取值与版本要求并能从 Swift 源码级调用链记录、切片、回放、清理验证 R2/R3 的完整实现。背景MoE RL 训练中的三个策略在 MoE 模型的 GRPO 训练中存在三个共享同一份模型权重、但路由行为可能不同的策略阶段策略记号路由结果说明训练策略$\pi_\theta$$e^{\pi}_t$执行梯度更新时的模型旧策略$\pi_{\theta_{\text{old}}}$$e^{\pi}_{\text{old},t}$本批更新前的模型状态采样策略$\mu_{\theta_{\text{old}}}$$e^{\mu}_{\text{old},t$推理引擎如 vLLM中的采样策略权重与旧策略相同但因 kernel 实现差异、精度等因素产生不同的路由结果其中 $\pi_{\theta_{\text{old}}}$ 与 $\mu_{\theta_{\text{old}}}$ 在采样时刻权重完全一致但由于推理引擎与训练引擎的实现差异如算子实现同一输入也可能得到不同的路由结果。训练-推理失配的分解参照 GSPO 论文中的推导token 级重要性采样比可以分解为两个因子$$ \frac{\pi_\theta(y_t|x, y_{t})}{\mu_{\theta_{\text{old}}}(y_t|x, y_{t})} \underbrace{\frac{\pi_{\theta_{\text{old}}}(y_t|x, y_{t})}{\mu_{\theta_{\text{old}}}(y_t|x, y_{t})}}{\text{训练-推理差异 (training-inference discrepancy)}} \times \underbrace{\frac{\pi\theta(y_t|x, y_{t})}{\pi_{\theta_{\text{old}}}(y_t|x, y_{t})}}_{\text{策略陈旧 (policy staleness)}} $$对 MoE 模型而言专家路由与这两个因子深度耦合训练-推理差异训练引擎与推理引擎之间的路由不一致$e^{\pi}{\text{old},t} \neq e^{\mu}{\text{old},t}$会放大输出分布的偏移策略陈旧随着 mini-batch 更新推进路由本身也会漂移$e^{\pi}t \neq e^{\pi}{\text{old},t}$进一步偏离采样策略。Router Replay 的作用正在于此通过固定前向过程中的专家选择从模型结构层面消除第一类失配同时将第二类失配控制在可控范围内。R2Vanilla Routing ReplayR2 的核心思想是在梯度更新时回放旧策略在训练引擎中确定的路由$e^{\pi}_{\text{old},t}$。原理在训练前向过程中先用旧策略权重做一次前向记录每个 MoE Router 层选中的专家索引然后强制训练模型 $\pi_\theta$ 在前向时使用这些索引$$ g_{\text{replay},i} \frac{I^{\pi}{\text{old},i} \cdot \exp(s{\text{train},i})}{\sum_j I^{\pi}{\text{old},j} \cdot \exp(s{\text{train},j})} $$其中 $I^{\pi}{\text{old}}$ 是旧策略的路由掩码$s{\text{train}}$ 是训练模型计算出的 router logits。注意 softmax 仍然作用在训练的 logits 上因此梯度可以正常回传到 router 权重——回放的是选哪些专家而不是打分结果。性质场景行为第一个 mini-batchon-policy$\theta \theta_{\text{old}}$故 $e^{\pi}t e^{\pi}{\text{old},t}$目标策略不变无偏后续 mini-batchesoff-policy$\theta \neq \theta_{\text{old}}$故 $e^{\pi}t \neq e^{\pi}{\text{old},t}$目标策略改变有偏但策略陈旧度被控制Swift 中的 R2 实现R2 完全发生在训练引擎内部无需推理引擎配合。从 MegatronGRPOTrainer 的实现可以看到其调用时序在计算旧策略 per-token logps 之前将全局路由回放动作设为RECORDif self.enable_routing_replay: if self.args.router_replay_mode R2: RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD) if self.args.router_replay_mode R3: RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD) old_per_token_logps_packed, routing_topk_idx self.compute_per_token_logps(...) if self.enable_routing_replay: batch[routed_experts] routing_topk_idx RouterReplay.clear_global_indices() RouterReplay.clear_global_router_replay_action()即RECORD动作下做一次无梯度的旧策略前向get_router_replay_data 随后从本 rank 各 Router 实例的recorded_topk_idx中取出路由索引并对索引做合法性校验必须落在[0, num_experts)内与压缩存储专家数 ≤ 256 时用uint8否则int32得到形状为[1, seq_len, num_moe_layers, topk]的张量挂到 batch 的routed_experts字段上。进入训练 step 后MegatronBaseTrainer 在forward_backward_func之前将全局动作置为REPLAY_FORWARD使训练前向与反向都使用刚才记录的路由优化器 step 完成后清理全局索引与动作为下一个 batch 的RECORD做准备。R3Rollout Routing ReplayR3 的核心思想是在训练前向中回放采样策略在推理引擎中确定的路由$e^{\mu}_{\text{old},t}$。原理在推理引擎如 vLLM采样时额外记录每个 token 在每一层 MoE Router 上的专家索引路由掩码然后把这些掩码传入训练引擎强制 $\pi_\theta$ 在前向时使用它们$$ g_{\text{replay},i} \frac{I^{\mu}{\text{old},i} \cdot \exp(s{\text{train},i})}{\sum_j I^{\mu}{\text{old},j} \cdot \exp(s{\text{train},j})} $$与 R2 相比R3 回放的正是采样时实际走过的路由因此它直接消灭了训练-推理差异这一项而不仅仅是消除训练引擎内部的记录-使用差异。与其他方法的兼容性R3 与GSPO正交可以组合使用以进一步提升稳定性R3 与TIS截断重要性采样修正组合可能没有额外收益——R3 已经从源头消除不一致TIS 的额外修正可能是冗余的在 off-policy 训练中Router Replay 与Clipping都是必要的。Swift 中的 R3 实现R3 的实现横跨推理与训练两侧推理侧vLLM。MegatronRolloutMixin 在构建 vLLM 引擎时若router_replay_mode R3会强制断言 vLLM ≥ 0.14.0 并开启enable_return_routed_expertsif self.args.router_replay_mode R3: assert check_vllm_version_ge(0.14.0), \ The enable_return_routed_experts attribute is not supported. Please upgrade vllm to 0.14.0 or higher vllm_engine_kwargs[enable_return_routed_experts] True返回的routed_experts作为 RolloutOutput 采样结果 的一个NumpyArray字段随输出返回。若 vLLM tensor parallel size 1由于 vLLM 的路由捕获主机侧缓存只存在于 TP rank 0其余 rank 的输出中routed_experts为空rollout_mixin.py#L856-L868 会通过broadcast_object_list从 TP 主 rank 广播到所有 TP rank保证每个 rank 的样本都携带完整路由信息。训练侧Megatron。routed_experts进入 micro-batch 后forward_step 在每次前向前弹出该字段并调用 set_router_replay_data 注入到当前 PP 段各 MoE Router 的target_indices反向结束后动作切换为REPLAY_BACKWARD保证反向传播沿用同一路由。由于 vLLM 按全部 Transformer 层含 dense 层上报routed_experts而 Megatron 中只有 MoE 层存在 router且各 rank 只持有并行切分后的本地层get_local_topk_idx_for_current_rank 负责三级切片PP 切片按moe_layer_freq找出当前 pipeline rank 的全局层区间只保留 MoE 层的索引CP 切片context parallel 1 时按zigzag/contiguous分区模式配合cu_seqlens_q切分序列维SP 切片scatter_to_sequence_parallel_region按序列并行切分。计算旧策略 logps 的公共路径 compute_per_token_logps 则根据动作分派REPLAY_FORWARDR3时断言 batch 中必须存在routed_experts并注入RECORDR2时前向后取回记录的索引。此外由于回放的路由掩码可能产生重复的专家索引使routing_map.sum() num_tokens * topkapply_router_replay_patch 会对MoEAlltoAllTokenDispatcher.preprocess打补丁从 routing map 实际求和推导num_out_tokens避免 alltoall 切分尺寸不匹配。Router Mask 缓存R3 相关论文还提出路由掩码可以与 KV Cache 一起缓存。对相同的前缀 tokenMoE Router 的输出是确定的因此路由掩码可以随前缀 KVCache 一起存储并复用。这在多轮 Agent 场景工具调用中尤为重要——无需重新 prefill 前缀即可获取路由掩码整体 rollout 的时延开销小于 3%。Swift 使用方式参数与环境要求通过--router_replay_mode参数选择路由回放策略定义见 MegatronArguments也在 Megatron-SWIFT 命令行参数文档 中有记录取值说明disabled默认不进行路由回放R2Vanilla Routing Replay在训练引擎中记录旧策略路由并回放R3Rollout Routing Replay从推理引擎导出路由掩码并在训练中回放环境要求R3 要求 vLLM ≥ 0.14.0以支持返回routed_experts信息对应源码中的enable_return_routed_experts开关Router Replay 目前仅在 Megatron 后端可用要求 megatron-core ≥ 0.16.0依赖megatron.core.transformer.moe.router_replay模块apply_router_replay_patch 在版本不满足时直接断言报错提示升级。从源码结构看该能力挂在 Megatron GRPO 训练链路上参数在 MegatronArguments 中声明MegatronBaseTrainer 在router_replay_mode ! disabled时置位enable_routing_replay并应用 dispatcher 补丁R2/R3 的差异最终体现为全局RouterReplayAction的取值RECORD/REPLAY_FORWARD以及routed_experts数据来源训练引擎自记录 vs vLLM 导出。与训练-推理偏差修正的关系Router Replay 与 训练-推理失配修正Training-Inference Mismatch 中描述的重要性采样IS修正是互补的两层手段IS 修正在 loss 层面通过重要性权重修正概率偏差属于事后补偿Router Replay通过固定路由在模型结构层面消除偏差的来源属于源头治理。两者的关系可以对照 R3 的兼容性结论理解R3 已经消除了训练-推理差异因子此时再叠加 TIS 类修正往往是冗余的而 R2 只消除了训练引擎内部的记录差异推理引擎与训练引擎之间的残余失配仍可结合 IS 修正处理。参考论文《Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers》Router Replay 的原始方法论文论文《Group Sequence Policy Optimization》GSPO训练-推理失配分解的出处论文《Stabilizing Reinforcement Learning with LLMs: Formulation and Practices》Megatron Core 官方文档中的 Router Replay API 指南文中涉及的实现文件均可在当前仓库中继续深入阅读router_replay_utils.py、Megatron GRPO Trainer、Rollout Mixin、Trainer 公共工具。【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600 LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300 MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考