Skip to content

[Feature] Integrate UltraEP Runtime for MoE EP Load Balancing - #2050

Open
ShilohYu wants to merge 12 commits into
InternLM:mainfrom
ShilohYu:feat/ultraep
Open

[Feature] Integrate UltraEP Runtime for MoE EP Load Balancing#2050
ShilohYu wants to merge 12 commits into
InternLM:mainfrom
ShilohYu:feat/ultraep

Conversation

@ShilohYu

@ShilohYu ShilohYu commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

背景

在 MoE Expert Parallel（EP）训练中，router post-gating 后的专家负载不均衡会使端到端吞吐受最慢 rank 限制。UltraEP 根据当前 layer / microbatch 的实际负载创建热点专家副本，并将部分 token 重路由到 replica slot。

本 PR 将 UltraEP 作为 Xtuner 的可选运行时接入，同时保持职责边界：

  • router 和训练语义继续使用 logical expert；
  • dispatcher 继续负责 token 通信；
  • Xtuner/FSDP/optimizer 继续管理 master expert 参数；
  • UltraEP 管理 replica weight、replica grad、placement、reroute 和 weight sync；
  • replica buffer 不进入模型参数、optimizer state 或 checkpoint。

主要改动

UltraEP runtime

  • 新增 MoEConfig.ultraep_cfg 和 UltraEPManager；
  • 支持 logical expert 到 physical expert 的 placement 和 reroute；
  • 支持 master/replica 权重同步；
  • 支持 replica gradient reduce；
  • 增加 force-balanced benchmark proxy；
  • 增加运行时诊断和 step NVTX 标记。

TE grouped GEMM

  • 支持 master weight 与 replica weight 使用不同 storage base；
  • master WGrad 返回 Xtuner autograd；
  • replica WGrad 写入 UltraEP 管理的 replica buffer；
  • 通过 XTUNER_GROUP_GEMM=te 启用 TEGroupedGEMM adapter；
  • 可通过 XTUNER_TE_GEMM_BACKEND=auto|cublas|cutlass|torch 进一步选择 TE 内部 backend。

支持范围

当前已验证单机八卡下EP8×DP1 单microbatch/多microbatch、FSDP EP4×DP2设置，底层dispatcher为deepep。暂不支持 activation recompute、MTP expert layers、all2all/agrs dispatcher、FP8 grouped experts 和 expert bias。

实验结果

实验使用单机 8×H200、DeepEP、BF16、torch.compile 和缩小 Qwen3-235B-A22B 结构。正式 timing 使用 25 steps，step 1–5 warmup，step 6–25 统计。

EP8×DP1，single-microbatch

模式 Step time (s) 全局吞吐 (tokens/s) TFLOPS/GPU (TFLOPS) MFU (%) 相对 force (%) 峰值显存 (GB)
baseline 1.2638 104,031 264.72 26.75 83.98 107.37
UltraEP off 1.2678 103,721 263.89 26.67 83.73 107.33
UltraEP on (R=1) 1.1730 111,758 285.19 28.82 90.22 103.38
force-balanced 1.0582 123,877 316.14 31.95 100.00 101.89

UltraEP on 相比 UltraEP off：

  • Step time 降低约 7.48%；
  • 全局吞吐提升约 7.75%；
  • TFLOPS/GPU 提升约 8.08%；
  • MFU 提升约 2.15 个百分点；
  • 峰值显存减少约 3.95 GB；
  • 达到 force-balanced 吞吐的 90.22%。

baseline、UltraEP off 和 UltraEP on 的 loss 曲线稳定下降，未观察到 NaN、Inf 或异常发散。force-balanced 用于均衡性能上限参考，不作为正常路由正确性基线。

FSDP EP4×DP2

该实验验证 FSDP local expert shard、reshard_after_forward=True、TE grouped GEMM 以及 UltraEP/FSDP gradient reduce 顺序。

模式 Step time (s) 全局吞吐 (tokens/s) TFLOPS/GPU (TFLOPS) MFU (%) 相对 force (%) 峰值显存 (GB)
baseline 1.1863 ± 0.0005 110,357.5 ± 46.4 236.73 23.92 92.41 85.98 ± 0.04
UltraEP off 1.1746 ± 0.0011 111,472.8 ± 101.4 241.51 24.41 93.35 82.52 ± 0.04
UltraEP on (R=1) 1.1590 ± 0.0033 112,927.9 ± 311.5 247.97 25.06 94.57 83.13 ± 0.12
force-balanced 1.0980 ± 0.0228 119,417.2 ± 2,072.3 276.78 27.97 100.00 82.04 ± 0.00

UltraEP on 相比 UltraEP off：

  • Step time 降低约 1.33%；
  • 全局吞吐提升约 1.31%；
  • TFLOPS/GPU 提升约 2.68%；
  • MFU 提升约 0.65 个百分点；
  • 峰值显存增加约 0.61 GB；
  • 达到 force-balanced 吞吐的 94.57%。

EP8×DP1，multi-microbatch（MB=2）

模式 Step time (s) 全局吞吐 (tokens/s) TFLOPS/GPU (TFLOPS) MFU (%) 相对 force (%) 峰值显存 (GiB)
baseline 1.6846 ± 0.0131 78,517.2 ± 629.2 200.75 ± 1.61 20.29 ± 0.16 61.45 115.040 ± 0.152
UltraEP off 1.6326 ± 0.0457 80,869.2 ± 2,290.5 206.77 ± 5.85 20.90 ± 0.59 63.29 110.299 ± 0.119
UltraEP on (R=1) 1.4370 ± 0.0704 92,412.7 ± 4,188.5 236.28 ± 10.71 23.88 ± 1.08 72.32 110.132 ± 0.409
force-balanced 1.0328 ± 0.0431 127,784.0 ± 4,607.0 326.73 ± 11.78 33.02 ± 1.19 100.00 109.201 ± 0.001

UltraEP on 相比 UltraEP off：

  • Step time 降低约 11.98%；
  • 全局吞吐提升约 13.61%；
  • TFLOPS/GPU 提升约 14.27%；
  • MFU 提升约 2.98 个百分点；
  • 峰值显存降低约 0.167 GiB；
  • 达到 force-balanced 吞吐的 72.32%。

该配置下 UltraEP on/off 的 loss trajectory acceptance threshold 已通过

@ShilohYu
ShilohYu requested a review from jayhenry August 26, 2026 09:22
@jayhenry
jayhenry requested a review from YanhuiDua August 31, 2026 04:11
@ShilohYu
ShilohYu force-pushed the feat/ultraep branch 5 times, most recently from ae7d6f5 to f73980c Compare September 17, 2026 05:40
@ShilohYu
feat: add force-balanced router benchmark mode
da2d9ab 
Add compile-compatible random-logits routing for balanced benchmark runs.
@ShilohYu
feat(moe): integrate UltraEP runtime with DeepEP dispatcher
ab9d0db 
Integrate UltraEP as an opt-in MoE execution path for DeepEP-based expert
parallel training, and move autograd hooks to the decoder so replica
weights are not snapshotted.

Supported in this initial integration:
- dispatcher="deepep" with ep_size > 1
- BF16 grouped experts without expert bias
- intra_layer_micro_batch == 1
- DP=1, where the EP group covers the whole world
- runtime-owned redundant expert weight/grad buffers outside model parameters,
  optimizer state, and checkpoints
- UltraEP placement/reroute, replica weight sync, fused master+replica grouped
  GEMM, and replica grad reduce for the single-microbatch training path

Not supported yet:
- intra_layer_micro_batch > 1
- activation recompute with UltraEP
- MTP expert layers with UltraEP
- dispatcher="all2all" or dispatcher="agrs"
- FP8 grouped experts
- expert bias
- DP>1/FSDP gradient-reduction ordering
@ShilohYu
feat: use TEGroupedGEMM adapter for MoE grouped GEMM
b2135d7 
Use the externally packaged TEGroupedGEMM backend through a dedicated XTuner adapter. Convert host-side token counts to m_splits, preserve master and UltraEP replica weight and gradient handling, and select the TE backend through get_group_gemm.

Co-authored-by: Cursor <cursoragent@cursor.com>
@ShilohYu
feat(moe): adapt UltraEP to intra-layer multi-microbatch
d8cf120 
Give each in-flight micro-batch its own virtual-layer id and replica
slot so weight-sync, restore, and grad-reduce stay correct when
intra_layer_micro_batch > 1.

Co-authored-by: Cursor <cursoragent@cursor.com>
@ShilohYu
feat(ultraep): adapt FSDP EP4 DP2
7af0d91 
Bind local FSDP expert shards into UltraEP and skip unchanged weight
pointer refresh after reshard.

Co-authored-by: Cursor <cursoragent@cursor.com>
@ShilohYu
perf(ultraep): add optional step nvtx ranges
553d8fc 
Add lightweight phase markers so disabled NVTX ranges stay cheap.
@ShilohYu
refactor(ultraep): hang UltraEP on EPExecutionRuntime and a six-stage…
9e9fef0 
… dispatcher

Create UltraEPModelRuntime in the dispatcher factory and move placement,
sync, restore, and grad-reduce onto a composite dispatcher so the decoder
no longer names UltraEP. prepare_layer_inputs validates the constructor-time
slot width instead of growing it.

Co-authored-by: Cursor <cursoragent@cursor.com>
@ShilohYu
feat(ultraep): keep UltraEP replica slots off ExpertWeightLayout
ae4b860 
Publish the MoonEP-shaped layout envelope for TE grouped GEMM, but leave
UltraEP replica slots out of compiled MoEBlock so later-layer slot reuse
does not trip torch.compile.

Co-authored-by: Cursor <cursoragent@cursor.com>
@ShilohYu
fix(ultraep): create gradient joins inside the microbatch loop
001f208
@ShilohYu
fix(ultraep): address review issues and add FSDP regression tests
62d72e8 
Preserve ordinary MoE list-forward behavior while keeping UltraEP capacity
validation in its runtime. Fix grouped-GEMM test inputs, optional TE
coverage, backend environment isolation, and lint/type issues.

Add opt-in eight-GPU EP8 and EP4/DP2 regression tests covering multiple
microbatches, gradient accumulation, optimizer updates, and coordinated
native teardown. Document the setup and numerical comparison limits.

Validation: 78 targeted tests and 2 eight-GPU integration tests passed;
61 tests passed and 17 skipped without TE. Ruff, formatting, and Pydantic
checks passed. Mypy adds no errors versus the same-environment baseline.
ShilohYu
ShilohYu commented Sep 18, 2026
View reviewed changes
Comment thread xtuner/v1/module/dispatcher/__init__.py
else:
process_group = ep_group
# TODO: remove type ignore here
# TODO: remove ignore

@ShilohYu ShilohYu Sep 18, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

没必要的lint修改，回滚

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jayhenry jayhenry Awaiting requested review from jayhenry
@YanhuiDua YanhuiDua Awaiting requested review from YanhuiDua

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@ShilohYu