执行摘要

• 一句话：修复 ROCm 上弹性 EP 缩放准确率并消除冗余重排通信开销
• 推荐动作：值得精读，尤其是 rank_load_imbalance 的 replica‑aware 不均衡度计算和 ROCm 平台专有条件的取舍。设计者对分布式数值一致性的强调和 review 中的集体决策讨论具有借鉴意义。

功能与动机

分布式测试 test_elastic_ep.py 在 ROCm 上 GSM8K 准确率从正常值崩溃至 0.19/0.016，冗余重排导致约 16s 的横向 all-to-all 颠簸使引擎饥饿，详见 PR body。

实现拆解

1. 防止 expert 拓扑缓冲区被清零：在 vllm/model_executor/layers/fused_moe/routed_experts.py 的 update_expert_map_info() 中，将五个 expert 拓扑缓冲区（_expert_map、expert_mask、expert_global_to_physical、expert_physical_to_global、expert_local_to_global）的 register_buffer 调用添加 persistent=False 参数。这些缓冲区是每 rank 派生的元数据，不应出现在 state_dict() 中，从而避免 ROCm dummy 权重初始化的 zero_() 和 batch_transfer_weights() 的覆盖。
2. 计算负载不均衡并跳过低增益重排：在 vllm/distributed/eplb/eplb_state.py 的 rearrange() 中增加跳过逻辑。定义 rank_load_imbalance() 函数，根据全局 expert 负载和物理到逻辑映射计算每个 rank 的负载不均衡度（考虑 replica 分摊）。比较当前映射和新映射的不均衡度相对改进，若改进 <5% 则设置 skip_rearrange = True。
3. 条件性跳过重排：skip_rearrange 为 True 时跳过 rearrange_expert_weights_inplace() 和 _commit_eplb_maps()，避免不必要的 all-to-all 通信。该逻辑仅在 ROCm 平台、非 profile 模式、非缩放场景（rank_mapping is None）且当前映射有效时生效。

关键文件：

• vllm/distributed/eplb/eplb_state.py（模块 负载均衡；类别 source；类型 core-logic；符号 rank_load_imbalance）: 核心逻辑修改：定义 rank_load_imbalance 函数，在 rearrange() 中插入 ROCm 专用跳过低增益重排的逻辑，避免 16s 冗余通信。
• vllm/model_executor/layers/fused_moe/routed_experts.py（模块 MoE 层；类别 source；类型 data-contract）: 数据契约修改：将 expert 拓扑缓冲区注册为非持久化，避免 ROCm dummy 权重初始化清零和 batch_transfer_weights 覆盖。

关键符号：rank_load_imbalance, update_expert_map_info

关键源码片段

vllm/distributed/eplb/eplb_state.py

核心逻辑修改：定义 rank_load_imbalance 函数，在 rearrange() 中插入 ROCm 专用跳过低增益重排的逻辑，避免 16s 冗余通信。

# 仅在 ROCm、非 profile、非缩放、且当前映射有效时启用跳过检查
skip_rearrange = False
if (
    current_platform.is_rocm()
    and not is_profile
    and rank_mapping is None
    and bool((eplb_model_state.physical_to_logical_map >= 0).all())
):
    logical_loads = global_expert_load_window.float()
    ep_size = ep_group.size()
​
    def rank_load_imbalance(
        mapping: torch.Tensor,
        logical_loads: torch.Tensor = logical_loads,
        ep_size: int = ep_size,
    ) -> float:
        # 将映射移动到负载张量所在设备并转为 long
        mapping = mapping.to(
            device=logical_loads.device,
            dtype=torch.long,
        )
        # 统计每个 physical expert 被映射到的 replica 数量
        replica_counts = torch.zeros_like(logical_loads)
        replica_counts.scatter_add_(
            dim=1,
            index=mapping,
            src=torch.ones_like(mapping, dtype=logical_loads.dtype),
        )
        # 将每个 logical expert 的负载均匀分摊到其所有 replica
        loads_per_replica = torch.gather(
            logical_loads / replica_counts.clamp_min(1),
            dim=1,
            index=mapping,
        )
        # 按 ep_size 重组并求和得到每个 rank 的总负载
        loads_per_rank = loads_per_replica.reshape(
            logical_loads.shape[0], ep_size, -1
        ).sum(dim=(0, 2))
        mean_load = loads_per_rank.mean()
        if mean_load == 0:
            return 1.0
        return (loads_per_rank.max() / mean_load).item()
​
    current_imbalance = rank_load_imbalance(
        eplb_model_state.physical_to_logical_map
    )
    proposed_imbalance = rank_load_imbalance(
        new_physical_to_logical_map
    )
    relative_improvement = (
        current_imbalance - proposed_imbalance
    ) / current_imbalance
    skip_rearrange = relative_improvement < 0.05
    if skip_rearrange and is_main_rank:
        logger.info(
            "[EPLB] Skip rearrange: imbalance %.4f -> %.4f "
            "(no material gain)",
            current_imbalance,
            proposed_imbalance,
        )
​
if not skip_rearrange:
    rearrange_expert_weights_inplace(
        eplb_model_state.physical_to_logical_map,
        new_physical_to_logical_map,
        eplb_model_state.model.expert_weights,
        eplb_model_state.expert_buffer,
        ep_group,
        eplb_model_state.communicator,
        is_profile,
        rank_mapping,
    )
    if not is_profile:
        _commit_eplb_maps(
            eplb_model_state,
            new_physical_to_logical_map=new_physical_to_logical_map,
        )

评论区精华

线程 1：tlrmchlsmth 担心不同 rank 的跳过决策因数值差异导致死锁。okorzh-amd 解释所有输入（global_expert_load_window 已 all-reduce、physical_to_logical_map 为全局不变式、new_physical_to_logical_map 来自确定性策略）均 rank 一致，不会分歧。后续更新添加了 all_reduce 确保集体决策。

线程 2：tlrmchlsmth 建议简化代码并考虑在非 AMD 硬件也启用跳过重排（因增加了一次 all‑reduce）。最终决策保留了 ROCm 专用，但代码块被大幅简化。

线程 3：itayalroy 指出 batch_transfer_weights() 排除了 expert_map 但未排除 expert_mask，导致新 rank 的 expert_mask 被发送者覆盖。okorzh-amd 确认并最初添加了排除，但最终通过 persistent=False 从根本上避免此问题。

• 跳过重排的决策分歧与死锁风险 (correctness): 确认输入一致无分歧，且后续代码中通过注释说明；review 后保持当前实现。
• 是否在非 AMD 硬件也启用跳过 (design): 最终保留 ROCm 专用，但代码块被简化，避免引入全局行为变化。
• batch_transfer_weights 中 expert_mask 排除不完整 (correctness): 通过 persistent=False 使这些缓冲区完全不出现在 state_dict 中，从根本上避免覆盖。

风险与影响

• 风险：数值一致性与集体决策：skip_rearrange 的当前实现依赖本地浮点计算，虽作者声称输入一致不会分歧，但未使用 all_reduce 同步决策，在极端浮点误差下存在不同 rank 决策不一致的风险，可能导致集体通信调用（rearrange_expert_weights_inplace）在部分 rank 上被跳过而其他 rank 不跳过的死锁。
  平台隔离：跳过逻辑限定于 ROCm，非 AMD 平台无变更，但引入平台分支增加了维护成本。
  缓冲区持久性更改：将 expert 拓扑缓冲区设为 persistent=False 使其不参与 state dict，不影响模型正确的加载/保存（它们是派生的），但需确保所有使用这些缓冲区的路径均通过 update_expert_map_info() 重建。

• 影响：对用户：ROCm 用户进行弹性 EP 缩放时，GSM8K 准确率从 0.19/0.016 恢复至 ≥0.58，且 Sync-EPLB 的冗余重排被消除，吞吐明显提升。非 ROCm 用户无直接影响。
  对系统：减少每次 step_interval 可能触发的 16s all-to-all 通信，降低网络压力，提升整体调度效率。
  对团队：明确了 ROCm dummy 权重初始化与 CUDA 的行为差异（整数 tensor 清零），以及 EPLB 重排频率对实际负载平衡的敏感性。

• 风险标记：数值决策未 all_reduce 同步, ROCm 专用分支增加维护成本, 缓冲区持久性变更需确保重建路径

关联脉络

• 暂无明显关联 PR