[Feature] Decouple expert parallelism from FSDP sharding (dp2ep layout)

[silencelamb]

Summary

This PR decouples MoE expert parallelism from FSDP sharding (the "dp2ep" layout of torchtitan): EP becomes a sub-dimension of the FSDP shard dimension instead of an orthogonal mesh axis. Dense parameters are sharded over the full data-parallel shard group and are no longer replicated across EP ranks; routed experts keep their EP split and are FSDP-sharded only over the remaining efsdp = dp_shard / ep ranks. The switch is FSDPConfig.decouple_ep_fsdp (default False; the legacy layout is untouched). The design note is docs/design/decouple_ep_fsdp.md.

flowchart LR
    R["root mesh (replicate, efsdp, ep)"] --> EP["ep_mesh = root[ep]<br/>token dispatch / expert ownership"]
    R --> DS["fsdp_mesh = flatten(efsdp, ep) = dp_shard<br/>FSDP group of dense params"]
    R --> EF["expert_fsdp_mesh = root[efsdp]<br/>FSDP group of routed experts on top of EP"]
    R -. "replicate > 1 (HSDP)" .-> HS["hsdp_mesh = root[replicate, dp_shard]<br/>expert_fsdp_mesh = root[replicate, efsdp]"]

Loading

With the legacy (fsdp, ep) mesh, fsdp = world / ep, so EP=8 on 8 GPUs keeps a full copy of every dense parameter and its fp32 optimizer state on every rank. For GLM-5.2-30B that pins the allocator at its ceiling (2–3 allocation retries per step) and makes EP8 about 7× slower than EP4. With the decoupled layout EP4 goes from 99.5 to 83.5 GiB at equal step time, and EP8 from 120.0 GiB / 11.8 s to 76.3 GiB / 1.64 s per step (8×H200, production recipe).

Changes

• FSDPConfig.decouple_ep_fsdp, validated by a pydantic model_validator (ep | dp_shard, positive sizes); the runtime mesh checks raise ValueError.
• MoE._init_decoupled_device_mesh: one root mesh (replicate, efsdp, ep); fsdp_mesh = root[efsdp, ep]._flatten("dp_shard") keeps its meaning for existing consumers, expert_fsdp_mesh is new, and HSDP + EP coexist on this path.
• Two-level fully_shard: every MoEBlock is wrapped on expert_fsdp_mesh first, then each decoder / MTP layer on the dense mesh; expert all-gathers are prefetched together with the next layer. The expert groups sit inside the layer's reentrant activation checkpoint, so FSDP2's per-call automatic reshard is off for them: hooks on the outer layer unit reshard them once after the layer's forward and once after its backward, and their backward prefetch points at the layer below.
• _scale_and_reduce_grad_decoupled: routed experts get grad.div_(ep) on top of FSDP's reduce-scatter, FSDP-ignored fully replicated fp32 parameters get one coalesced all-reduce, and the legacy manual cross-EP all-reduce of dense gradients is gone.
• Float8: per-class padding chunk counts and tile-wise reduce meshes (dense: dp_shard chunks, rank stride 1; experts: efsdp chunks, rank stride ep).
• HF / DCP checkpoints need no adapter changes: LoadSpec derives the shard history from the new DTensor placements.
• RL weight sync: BaseModel._fsdp_gather_group picks the gather group per LoadSpec, and WeightIterator._param_owner gathers compose-model parameters with the submodule that owns them, so efsdp / HSDP shards of a language tower are gathered instead of being mistaken for EP-local shards.
• examples/v1/config/sft_glm5p2.py: DECOUPLE_EP_FSDP and HSDP_SHARDING_SIZE environment switches.
• Tests: fake-process-group L0 placement tests for the legacy and decoupled layouts (8 / 16 / 64 ranks), 8-GPU L1–L3 gates, an 8-GPU gate that counts the all-gathers of each FSDP unit under intra-layer micro-batches and full recompute, RL weight-iterator tests; the manual experiment scripts under tests/model/ share their helpers with the gates.

Verification

• pre-commit hooks on the changed xtuner/v1 files (codespell, docformatter, pyupgrade, ruff, ruff-format, pydantic-extra-check) passed; mypy 1.16.1 reports no errors in the lines this PR touches.
• tests/model/test_decoupled_ep_fsdp_mesh.py: 31 passed (fake process groups, CUDA context only).
• tests/rl/test_weight_iterator.py: 13 passed.
• tests/engine/test_decoupled_ep_fsdp_train_engine.py on 8×H200 (torch 2.9.1, tiny random Qwen3-MoE): 7 passed. L1 (ep=1, ep=8 legacy, ep=8 decoupled) and L2 (ep=4 legacy vs. efsdp=2, HSDP + EP with efsdp=1 and efsdp=2): loss curves within 1e-3 relative (observed ~2e-5), total and step-0 per-parameter gradient norms within 1e-2, per-rank parameter memory matching the layout. L3: bit-exact HF export right after from_hf, DCP resume, HF export after resume within bf16 ulps, and cross-layout DCP resharding. The all-gather gate allows at most 2 used all-gathers per expert group and step.
• GLM-5.2-30B on 8×H200, production recipe: EP4 1.688 s / 99.5 GiB → 1.658 s / 83.5 GiB; EP8 11.8 s / 120.0 GiB → 1.64 s / 76.3 GiB. Tile-wise FP8 runs on a 3.4B Qwen3-MoE and on GLM-5.2 stay within run-to-run noise of the legacy layout. Details: reports/GLM52_PR_rerun.md.
• Intra-layer micro-batches with full recompute (Qwen3-30B-A3B, DeepEP, 4 × 16K tokens/GPU/step, same-batch runs): FSDP all-gathers per step 715 → 196 on EP4 with the per-layer reshard. Against the legacy layout, tokens/s goes from −5.3% to +0.6% on 8×H200 (EP4; peak 71.1 GiB vs. 79.4 legacy), and from −6.9% / −9.8% to +0.5% / +0.2% on 2 × 8 H200 (EP8 / EP4), with equal losses. Details: reports/decoupled_ep_fsdp_expert_reshard_zh.md.

Result

decouple_ep_fsdp defaults to False and the execution logic of the legacy path is unchanged; the legacy placements are pinned by the L0 tests. Known limits are listed in §6 of the design note: ExpertTP and decoupling are mutually exclusive, fp32_keys_pattern matching a routed expert is unsupported on both paths, HSDP + EP is validated on single-node scaled-down topologies only, and the two-mesh fully_shard + DeviceMesh._flatten combination is validated on torch 2.8 / 2.9.