[Feature] Integrate UltraEP Runtime for MoE EP Load Balancing

[ShilohYu]

背景

在 MoE Expert Parallel（EP）训练中，端到端吞吐受 token 负载最高的 rank 限制。router post-gating 后产生的专家负载偏斜，会使实际吞吐与理想均衡负载下的性能产生明显差距。UltraEP 是面向 MoE 的在线专家负载均衡运行时。它根据当前 layer / microbatch 的精确 post-gating 负载规划副本，复制热点专家，并将部分 token 重路由到 runtime replica slot。

本 PR 将 UltraEP 接入 Xtuner，同时保持已有职责边界：

• UltraEP 的职责是 placement、reroute 与 replica 生命周期管理；
• dispatcher 仍然是 token 通信后端；
• Xtuner/FSDP/optimizer 继续管理 master expert 参数。

术语解释：router 和上层训练语义看到的是 logical expert；每个 logical expert 对应一个由 Xtuner 管理的 master expert，并可由 UltraEP 创建多个 replica expert。master expert 与 replica expert 共同构成实际执行的 physical expert；只有 dispatcher 和 expert compute 使用 physical expert ID。

主要改动

新增可选的 UltraEP 运行时

• 新增 MoEConfig.ultraep_cfg: UltraEPConfig | None；None 是唯一关闭状态，默认不启用。
• UltraEPConfig 仅承载冗余专家槽位等用户显式配置；UltraEPManagerProvider.from_xtuner_config() 从 MoEConfig 统一派生层数和logical expert 数等运行时 shape 信息。
• 每个 EP group 延迟创建并复用一个 UltraEPManager；关闭时不导入或依赖 UltraEP。
• replica weight 与 FP32 replica grad 由 UltraEP runtime 管理，不注册为模型参数，也不进入 optimizer state 或 checkpoint。

保留 logical routing 语义并引入 physical routing

• 保留 router 输出的 logical topk_ids，供上层训练语义、auxiliary loss 和诊断使用。
• 在 router 与 dispatcher 之间引入 UltraEPLayerRuntime，负责更新 placement、异步同步 master expert weight，以及生成供 dispatch 与 expert compute 使用的 physical expert ID。
• 启用 UltraEP 后，dispatcher 的专家空间由 logical experts 扩展为由 master experts 和 replica experts 组成的 physical experts。

前反向时序

MoE 主链路仍为 attention -> router -> dispatch -> experts -> combine。

• 前向中，UltraEP 在 router 后依次执行 update_placement、异步 weight_sync 和 reroute；权重同步 weight_sync 与 reroute、dispatch 重叠，仅在 expert compute 前等待完成。
• expert compute 使用 dual-weight grouped GEMM，在一次 launch 中处理 master 与 replica expert。
• 反向中，三个 identity-style autograd hook 分别在 expert DGrad 前恢复对应 layer 的 replica weight、在 expert/dispatch backward 后异步启动 replica grad reduce、在 attention backward 后等待归约完成。
• 这样可保证跨 layer 复用 replica buffer 的正确性，并将 replica grad reduce 与 attention backward 重叠。

新增双基址 Grouped GEMM

• 新增 UltraEPGroupedGemm 与 dual-weight Triton grouped-GEMM kernel。
• master 与 replica weight 位于不同存储区域时，该kernel可一次 grouped GEMM launch 内选择正确权重并完成计算。
• master WGrad 返回 Xtuner autograd；replica WGrad 写入 UltraEP 管理的 FP32 buffer，后续归并回对应 master expert。

新增 force-balanced benchmark proxy

• 新增 force_load_balance：前向使用随机 logits 构造近似均衡路由，反向保留原 router logits 的梯度。
• 该模式仅用于评估均衡负载下的性能上限；它改变正常路由语义，不能作为正确性基线。
• 覆盖 torch.compile 兼容性测试。

当前支持范围

当前代码仅接通并验证 dispatcher="deepep" 路径，且显式限制：

• ep_size > 1、expert_tp_size == 1，且 logical expert 数可被 EP size 整除；
• BF16 grouped experts，不支持 expert bias；
• intra_layer_micro_batch == 1；
• DP=1，即 EP group 覆盖整个 world；
• 不支持 activation recompute 和 MTP expert layers；
• 暂不支持 all2all、agrs 等其他 dispatcher。

端到端性能验证

实验使用单机 8×H200、EP=8（TP/DP/SP/PP=1）、DeepEP、BF16、torch.compile，模型为 20 层 Qwen3-235B-A22B 缩小结构（32 experts、top-k=4、expert FFN 2048）。

模式	Step time (s)	全局吞吐 (packed tokens/s)	相对 force-balanced吞吐(%)	峰值显存 (GB)
baseline	1.2638	104,031	83.98%	107.37
UltraEP off	1.2678	103,721	83.73%	107.33
UltraEP on（R=1）	1.1730	111,758	90.22%	103.38
force-balanced proxy	1.0582	123,877	100.00%	101.89

UltraEP off 与 baseline 吞吐接近。开启 UltraEP 后，相比 off：

• step time 降低 7.48%，等价于 1.081× step-time speedup；
• 全局吞吐提升 7.75%；
• 峰值显存降低 3.95 GB；
• 性能达到 force-balanced proxy 的 90.22%。

正确性 sanity check

UltraEP on、UltraEP off 与 baseline 的 loss 曲线均稳定下降，未观察到 NaN、发散或异常拐点。