[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training by jayhenry · Pull Request #2056 · InternLM/xtuner

jayhenry · 2026-08-28T12:54:04Z

Summary

This PR adds an optional dispatcher="moonep" path for node-local BF16 MoE training. Native FSDP2 remains the sole owner of expert parameters, optimizer state, and checkpoint identity; MoonEP owns only the communication/VMM execution workspace.

Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.

flowchart LR
    A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
    B --> C["Dispatch + grouped GEMM + combine"]
    C --> D["BF16 duplicated-gradient return"]
    D --> E["FSDP reduce-scatter"]
    E --> F["FP32 shard gradient + optimizer state"]

Design

MoonEP is imported lazily only when selected. The integration validates a versioned backend capability before allocating resources, so DeepEP, All2All, AGRS, and non-EP configurations do not acquire a MoonEP dependency.

The hot path is ordered with CUDA events and GPU-side EP barriers. The profiler regression gate verifies:

no complete home-expert weight copy in the direct path;
no complete local duplicated-gradient temporary;
no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

FSDP post-all-gather hooks install the VMM landing tensors while preserving the original DTensor parameters. The backward path returns BF16 home gradients to those parameter edges; existing FSDP communication then reduce-scatters them and produces FP32 sharded gradients for the FP32 optimizer update.

Supported scope

PyTorch 2.12.1+cu132.
Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
TP1.
MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

Not claimed in this first version: Expert TP, FP8 experts, cross-node MoonEP, FSDP no_sync, pipeline parallelism, or decoding.

External MoonEP dependency

This XTuner branch expects MoonEP XTuner integration API v2. The validated companion implementation is currently commit c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

The companion MoonEP patch still needs to be published separately before this PR can be merged or reproduced outside the development environment. The XTuner package does not add MoonEP as a mandatory dependency.

Validation

The formal workload is Qwen3.5-35B-A3B on one node with 8 H200 GPUs, FSDP2 + EP4, BF16 parameters/reduction, Direct VMM landing, Triton grouped GEMM, and 20 training steps. Each comparison changes only the dispatcher.

Gate	DeepEP median tokens/s	MoonEP median tokens/s	MoonEP / DeepEP
MTP disabled, steps 6-20	5980.999	6047.073	101.10%
MTP1, steps 6-20	5357.164	5408.172	100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

MoonEP forward/MTP/Domino/SP/compile: 16 passed.
DCP/HF/offload/optimizer/lifecycle: 10 passed.
Dispatcher and grouped-GEMM regressions: passed.
FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

See the acceptance report for the full configuration, per-step metrics, profiler evidence, and reproduction commands.

Result

The first-version objective is met for single-node BF16 FSDP2 + EP training: MoonEP matches DeepEP numerically and reaches equivalent steady-state throughput while preserving FSDP parameter/checkpoint ownership and a host-sync-free communication hot path.

jayhenry · 2026-09-01T13:54:27Z

@claude review

jayhenry · 2026-09-02T08:55:04Z

            assert isinstance(loss_ctx, list) and len(loss_ctx) == len(seq_ctx), (
                "seq_ctx_list and loss_ctx_list must be lists of the same length"
            )
+            if self._moonep_runtime is not None and len(seq_ctx) > self.config.intra_layer_micro_batch:


改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

jayhenry · 2026-09-02T09:02:24Z

+        staging_reference: bool,
+        num_sms: int = 64,
+    ) -> None:
+        self._backend = require_moonep_backend()


直接检查是否安装 moonep，不做过度检查。
另外，在文件顶层 try import moonep 即可，后续引用也直接用 moonep，去掉 self._backend

jayhenry · 2026-09-02T10:00:45Z

        # moe forward
        # (e, )
-        tokens_per_expert = torch.histc(topk_ids, bins=self.n_routed_experts, min=0, max=self.n_routed_experts)
+        tokens_per_expert = torch.bincount(topk_ids.flatten(), minlength=self.n_routed_experts)


torch.bincount 存在强制 Host-Device 同步，改回原来的 torch.histc

jayhenry · 2026-09-02T11:41:11Z

            n_routed_experts=n_routed_experts,
            ep_group=process_group,
            tp_group=tp_group,
            ep_tp_group=ep_tp_group,


先hardcode bf16，去掉这个参数，assert 不支持 fp8

jayhenry · 2026-09-02T11:46:05Z

-            min=0,
-            max=self.n_routed_experts,
-        )  # .view(self.ep_mesh.size(), -1)
+        tokens_per_expert = torch.bincount(topk_ids.flatten(), minlength=self.n_routed_experts)


同样要使用不host sync的方式，比如 histc

jayhenry · 2026-09-02T12:33:47Z

 from torch import Tensor


-try:


保留 try except 的import 方式

jayhenry · 2026-09-02T13:40:55Z

+    fsdp_root: nn.Module,
+    targets: Sequence[tuple[str, tuple[nn.Module, nn.Module], tuple[torch.Tensor, torch.Tensor]]],
+) -> tuple[FSDPParam, ...]:
+    """Bind routed expert FSDPParams to their two-generation VMM landings.


将各种检查封装到一个私有函数，保证顶层的线性步骤

jayhenry · 2026-09-02T13:42:10Z

+        self._fixed_tokens_per_rank: int | None = None
+        self._closed = False
+
+    def bind_dispatcher(


改名为 build_dispatcher 是否更合适？

jayhenry · 2026-09-02T13:46:36Z

+    """Own one fresh MoonEP plan and all of its device-side completion
+    edges."""
+
+    def __init__(self, runtime: MoonEPRuntime, *, layer_id: int, grad_slot: int) -> None:


将 runtime 整个传参是否参数过宽？

jayhenry · 2026-09-02T13:47:03Z

+            topk_ids,
+            tokens_per_expert,
+            topk_weights,
+            self,


将 invocation 整个传参是否参数过宽？

jayhenry · 2026-09-02T14:12:11Z

+        # completed home gradients are handed back before FSDP post-backward.
+        with torch.profiler.record_function("MoonEP::prepare_experts"):
+            home_parameters = self._current_home_parameters()
+            local_weights, gradient_targets, weights_ready = workspace.materialize(


这个函数会 prefetch weight 产生GPU通信，如果想要 domino ep 正常工作，是不是应该将它放到 dispatcher.dispatch 函数中？

jayhenry · 2026-09-02T14:38:59Z

+            raise NotImplementedError("MoonEP fixed-S training dispatch does not implement decoding")
+        grad_slot = self._next_gradient_slot
+        self._next_gradient_slot = (grad_slot + 1) % self._runtime._intra_layer_micro_batch
+        return _MoonEPInvocation(self._runtime, layer_id=self._layer_id, grad_slot=grad_slot).dispatch(


现在 dispatcher太浅，只做简单的转发。将 Invocation 退化成一个对象类，并没有 dispatch, combine 等方法，让 dispatcher 本身有具体的dispatch等逻辑，这样改是否合理？

jayhenry force-pushed the moonep branch 2 times, most recently from 06d5929 to 84b40b4 Compare August 31, 2026 09:49

jayhenry commented Sep 2, 2026

View reviewed changes

Comment thread xtuner/v1/module/dispatcher/moonep.py Outdated

jayhenry commented Sep 2, 2026

View reviewed changes

jayhenry added 2 commits September 2, 2026 16:19

feat(moe): extend grouped GEMM for dynamic expert gradients

a87861b

refactor(ep): unify dispatcher weight and routing contracts

4acd90f

jayhenry force-pushed the moonep branch from 84b40b4 to d89d033 Compare September 2, 2026 18:40

jayhenry added 9 commits September 3, 2026 02:29

feat(moonep): integrate BF16 runtime and FSDP expert landing

6f9573a

test(moonep): add Qwen3.5 acceptance gate

8d09bb5

refactor(moonep): own FSDP VMM workspace in XTuner

e8974ff

fix(moe): restore device-only routing counts and Domino width contract

9806669

refactor(moonep): simplify optional backend and BF16 capability checks

e33d6a1

refactor(moonep): make dispatcher own per-layer execution

3212649

refactor(moonep): linearize FSDP VMM landing installation

0a1a01e

fix: recognize FSDP strided shards in grad norm

3f5414f

docs(moonep): record final acceptance results

1a116c4

jayhenry force-pushed the moonep branch from d89d033 to 1a116c4 Compare September 3, 2026 02:30

refactor(moonep): clarify runtime and invocation ownership

ee10d26

jayhenry added 3 commits September 4, 2026 03:10

refactor(moonep): enhance workspace allocation and layout management

0cf3da0

refactor(moonep): defer pair gradient handoff to invocation

1ff9e83

perf(moonep): overlap duplicate-gradient return with upstream backward

0d6c671

Uh oh!

Conversation

jayhenry commented Aug 28, 2026

Summary

Design

Supported scope

External MoonEP dependency

Validation

Result

Uh oh!

jayhenry commented Sep 1, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayhenry Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

jayhenry Sep 2, 2026 •

edited

Loading