执行摘要

• 一句话：修复 allgather_cp 模式下 routing replay 专家对齐逻辑
• 推荐动作：建议精读 cp_utils.py 中新增的 prepare_routed_experts_for_routing_replay 函数，其展示了在 CP+TP 并行下如何对齐 rollout 与训练阶段的专家索引。该设计模式值得复用。同时关注后续是否追加测试覆盖。

功能与动机

在 allgather_cp 模式下，原有的专家令牌对齐逻辑未考虑 CP（Context Parallel）分组，导致路由回放（Routing Replay）时专家索引错位，进而影响 MoE 模型训练正确性。PR 标题 'Fix R3 for allgather_cp' 直接点明了修复目标。

实现拆解

1. 提取公用填充函数 _pad_routed_experts：在 cp_utils.py 中新增 _pad_routed_experts，用于将专家索引张量按指定大小 padding，填充值循环取模 num_experts。该函数替代了 actor 中原有的 pad_func 闭包。

2. 新增核心对齐函数 prepare_routed_experts_for_routing_replay：同样位于 cp_utils.py，接收 rollout 产出的专家索引序列和 token 序列，返回对齐后的专家张量。函数内部首先对每个序列 padding 1 个 token，然后根据 allgather_cp 标志选择不同路径：

   • 若 allgather_cp=True：将所有序列拼接后，按 CP_size * pad_size 对齐 padding，然后 chunk 分到当前 CP rank。
   • 若 allgather_cp=False：沿用原有逻辑，对每个序列调用 slice_with_cp 后拼接，再按 pad_size 对齐 padding。
     最后若 sequence_parallel=True，按 TP rank 切分序列。

3. 简化 actor.py 中的 fill_routing_replay：移除原有内联的 pad_func 定义、断言循环、手动的 padding 和 slice_with_cp 调用，直接调用 prepare_routed_experts_for_routing_replay。同时更新导入语句，从 cp_utils 导入新函数而非 slice_with_cp。

4. 移除不再需要的 slice_with_cp 导入：actor.py 中原来导入 slice_with_cp，现在仅需 prepare_routed_experts_for_routing_replay，因此从导入列表中删除 slice_with_cp。

5. 配置参数传递：新函数通过关键字参数接收 num_experts、data_pad_size_multiplier、sequence_parallel 和新增的 allgather_cp，这些参数从 self.args 直接传入，无需在函数内部再次获取。

关键文件：

• slime/backends/megatron_utils/cp_utils.py（模块 后端；类别 source；类型 core-logic；符号 _pad_routed_experts, prepare_routed_experts_for_routing_replay）: 新增了两个核心函数 _pad_routed_experts 和 prepare_routed_experts_for_routing_replay，实现了 allgather_cp 分支的专家对齐逻辑，是整个修复的核心。
• slime/backends/megatron_utils/actor.py（模块 后端；类别 source；类型 core-logic；符号 fill_routing_replay）: 重写了 fill_routing_replay 方法，移除大量内联逻辑并委托给 prepare_routed_experts_for_routing_replay，简化了代码。

关键符号：_pad_routed_experts, prepare_routed_experts_for_routing_replay, fill_routing_replay

关键源码片段

slime/backends/megatron_utils/cp_utils.py

新增了两个核心函数 _pad_routed_experts 和 prepare_routed_experts_for_routing_replay，实现了 allgather_cp 分支的专家对齐逻辑，是整个修复的核心。

# slime/backends/megatron_utils/cp_utils.py
​
def _pad_routed_experts(experts: torch.Tensor, pad: int, num_experts: int) -> torch.Tensor:
    # 如果不需要 padding 则直接返回
    if pad == 0:
        return experts
    _, num_layers, topk = experts.shape
    # 生成填充值：使用 arange 并取模 num_experts，确保填充的专家索引在有效范围内
    pad_experts = (
        torch.arange(
            pad * num_layers * topk,
            device=experts.device,
            dtype=experts.dtype,
        ).reshape((pad, num_layers, topk))
        % num_experts
    )
    # 拼接原张量与填充张量
    return torch.cat([experts, pad_experts], dim=0)
​
​
def prepare_routed_experts_for_routing_replay(
    rollout_routed_experts: Sequence[torch.Tensor],
    tokens: Sequence[torch.Tensor],
    *,
    num_experts: int,
    data_pad_size_multiplier: int,
    sequence_parallel: bool,
    allgather_cp: bool,
) -> torch.Tensor:
    # 对齐 rollout 产出的路由专家元数据与训练 token 布局
    assert len(rollout_routed_experts) == len(tokens)
    for experts, token_ids in zip(rollout_routed_experts, tokens, strict=False):
        # 专家数应比 token 数少 1（最后一个 token 不参与路由预测）
        assert experts.shape[0] == token_ids.shape[0] - 1
​
    # 先对每个序列的最后一个位置 padding 一个 expert
    padded_experts = [_pad_routed_experts(experts, 1, num_experts) for experts in rollout_routed_experts]
    pad_size = mpu.get_tensor_model_parallel_world_size() * data_pad_size_multiplier
​
    if allgather_cp:
        # allgather_cp 模式：所有序列拼接后，按 CP_size * pad_size 整体对齐，再分片到各 CP rank
        routed_experts = torch.cat(padded_experts, dim=0)
        cp_size = mpu.get_context_parallel_world_size()
        cp_rank = mpu.get_context_parallel_rank()
        global_pad_size = cp_size * pad_size
        pad = (global_pad_size - routed_experts.size(0) % global_pad_size) % global_pad_size
        routed_experts = _pad_routed_experts(routed_experts, pad, num_experts)
        routed_experts = routed_experts.chunk(cp_size, dim=0)[cp_rank]
    else:
        # 非 allgather_cp 模式：每个序列先 slice_with_cp 再拼接，最后按 pad_size 对齐
        routed_experts = [
            slice_with_cp(experts, lambda x, pad: _pad_routed_experts(x, pad, num_experts))
            for experts in padded_experts
        ]
        routed_experts = torch.cat(routed_experts, dim=0)
        pad = (pad_size - routed_experts.size(0) % pad_size) % pad_size
        routed_experts = _pad_routed_experts(routed_experts, pad, num_experts)
​
    if sequence_parallel:
        # 若开启序列并行，按 TP rank 切分序列维度
        tp_rank = mpu.get_tensor_model_parallel_rank()
        tp_size = mpu.get_tensor_model_parallel_world_size()
        seqlen = routed_experts.size(0)
        assert seqlen % tp_size == 0
        start = seqlen // tp_size * tp_rank
        end = seqlen // tp_size * (tp_rank + 1)
        routed_experts = routed_experts[start:end]
​
    return routed_experts

slime/backends/megatron_utils/actor.py

重写了 fill_routing_replay 方法，移除大量内联逻辑并委托给 prepare_routed_experts_for_routing_replay，简化了代码。

# slime/backends/megatron_utils/actor.py ( 改动后 )
​
def fill_routing_replay(self, data_iterator, num_microbatches, rollout_data):
    if "rollout_routed_experts" not in rollout_data:
        raise ValueError(
            "rollout_routed_experts is required in rollout_data when use_rollout_routing_replay is set."
        )
​
    from megatron.core.transformer.transformer_block import get_num_layers_to_build
    from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
    from slime.utils.routing_replay import RoutingReplay
​
    for iterator in data_iterator:
        iterator.reset()
​
    # 循环处理每个 micro-batch
    for _ in range(sum(num_microbatches)):
        batch = data_iterator[0].get_next(["rollout_routed_experts", "tokens"])
        # 委托给 cp_utils 中的对齐函数，传入所需配置
        rollout_routed_experts = prepare_routed_experts_for_routing_replay(
            batch["rollout_routed_experts"],
            batch["tokens"],
            num_experts=self.args.num_experts,
            data_pad_size_multiplier=self.args.data_pad_size_multiplier,
            sequence_parallel=self.args.sequence_parallel,
            allgather_cp=self.args.allgather_cp,
        )
​
        routing_replay_offset = 0
        for vp_stage, model in enumerate(self.model):
            config = model.module.config
            num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage)
            offset = get_transformer_layer_offset(config, vp_stage=vp_stage)
            for layer_id in range(offset, offset + num_layers_to_build):
                # 跳过稠密层（非 MoE 层）
                if isinstance(config.moe_layer_freq, int):
                    if layer_id % config.moe_layer_freq != 0:
                        continue
                elif isinstance(config.moe_layer_freq, list):
                    assert len(config.moe_layer_freq) > 0
                    # ... 后续逻辑不变

评论区精华

本 PR 无公开 review 讨论，合并人即作者，直接合并。

• 暂无高价值评论线程

风险与影响

• 风险：1）修改涉及 fill_routing_replay 核心训练循环，若 prepare_routed_experts_for_routing_replay 存在逻辑错误，可能导致所有 MoE 模型训练失败。2）移除了 actor 中原有的断言与检查，若新函数未完全等价替代，可能遗漏错误。3）未新增测试，依赖现有测试覆盖，需关注 CI 结果。4）allgather_cp 分支新增了 chunk 操作，若 CP size 不能整除会导致断言失败。
• 影响：直接影响使用 allgather_cp=True 且启用 use_rollout_routing_replay 的 MoE 模型训练，此前可能训练异常，修复后恢复正确。对非 allgather_cp 模式无影响。代码可维护性提升，后续 routing replay 逻辑统一在 cp_utils.py 中修改。影响范围为 Megatron 后端训练流程。
• 风险标记：核心路径变更, 缺少测试覆盖, 并行一致性

关联脉络

• PR #2173 [docker] Update SGLang patch for PD R3 routed experts: 同样涉及 routed experts 的验证与配置修改，与本文的 routing replay 对齐逻辑属于同一功能线。
• PR #2169 Merging profiling info into router: 重构了 sglang 补丁与 routing 相关逻辑，本文进一步将 routing replay 逻辑从 actor 提取到 cp_utils，延续了重构趋势。