[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training

[jayhenry]

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; the private VMM workspace owns the physical expert layout.
• Dispatch activations and weights, reuse the existing grouped GEMM path (extended for dynamic expert gradients), and return duplicated expert gradients in BF16 before FSDP reduce-scatter, joined once per layer.
• Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
• Unify the EP dispatcher weight/routing contracts; keep routing counts on device via torch.histc.
• Integrate DCP/HF persistence, activation/router offload, optimizer swap, Muon, and explicit runtime teardown.
• Add a reproducible Qwen3.5 DeepEP/MoonEP acceptance gate (tests/acceptance/).
• 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"]

Loading

Branch status

Rebased onto main and organized as three commits:

1. [Fix] route the grad-norm placement check through RuntimeLayout.is_sharded_placement so FSDP2 + EP _StridedShard bookkeeping placements (no longer a Shard subclass on PyTorch 2.12) contribute their all-reduce.
2. [Feature] the MoonEP integration.
3. [Test] the acceptance gate.

Design notes and the full per-run validation report are kept out of the PR and maintained separately.

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 v3. The validated companion implementation is commit d4494473fb0932dcc35f3a44a0e3b31827f5e282 on the 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

Acceptance harness

tests/acceptance/ (config-locked by tests/engine/test_moonep_acceptance.py, compared by xtuner/_testing/moonep_acceptance.py) runs 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, changing only the dispatcher between the DeepEP and MoonEP legs. Default attention is FlexAttention + deterministic mode; XTUNER_USE_FA3 and XTUNER_DETERMINISTIC are overridable.

Validated behaviour (calibrated FlexAttention + deterministic gate)

On the functionally-equivalent pre-rebase implementation, MoonEP matches DeepEP on every loss and grad-norm curve (cosine similarity >= 0.99, mean relative difference < 1%) and reaches 100-112% of DeepEP steady-state throughput at MTP disabled and MTP1. GPU regression suites (MoonEP forward / MTP / Domino / SP / compile; DCP / HF / offload / optimizer / lifecycle; dispatcher and grouped-GEMM; FSDP2 + EP4 nested-shard grad norm on PyTorch 2.12 and 2.9) passed.

Post-rebase verification on this branch

• Non-GPU regression: 40 passed — acceptance config-lock across all dispatcher/MTP/micro combinations, MoonEP dispatcher contract and six-stage seam, and device-only routing counts.
• GPU, FA3 (non-deterministic; FA3's Hopper backward has no deterministic path at head_dim 256), Qwen3.5-35B-A3B, FSDP2 + EP4, 20 steps, DeepEP vs MoonEP, MTP0 and MTP1, at 32k (micro1) and 16k (micro2) sequence length:
  • Numerical parity preserved: every curve vs DeepEP has cosine >= 0.99997, mean relative difference < 1.2%.
  • Throughput MoonEP/DeepEP: ~0.95 (micro1) / ~0.84 (micro2) — bit-for-bit consistent with both the pre-rebase branch and the 08-28 validated baseline under the same settings, i.e. no regression from the rebase. FA3 + non-deterministic is outside the calibrated gate's regime (which is FlexAttention + deterministic, where MoonEP is throughput-neutral) and this delta is present at the 08-28 baseline as well.

The formal FlexAttention + deterministic GPU gate has not yet been re-executed on the rebased branch; the acceptance harness is unchanged and runs it directly.

Result

The first-version objective is met for single-node BF16 FSDP2 + EP training: MoonEP matches DeepEP numerically, preserves FSDP parameter/checkpoint ownership and a host-sync-free communication hot path, and is throughput-neutral in the calibrated acceptance configuration.

[jayhenry]

@claude review