(fsdp2 dev.) feat(moe): expert parallelism for MoE drafters on the FSDP2 backend (stacked on #915 + #931)

[yushengsu-thu]

Draft, stacked on #915 (FSDP2 backend) and #931 (@sherlockwu's MoE stack #811/#812 rebased on main). This branch is moe-stack-rebased (#931) + a merge of codex/fsdp2-backend (#915; one conflict in training/backend.py, resolved by applying the MoE checkpoint-naming converters around the FSDP2 state-dict seam) + one feature commit. Until those land, review the last commit: git diff 7965b568..12c3f9d2.

Motivation

The MoE stack trains routed experts as ordinary FSDP parameters, so every micro-batch re-gathers all of them: for the DeepSeek-V4-Flash drafter that is 32.7B expert parameters (~65 GB per micro-batch, 46 s/step on 8x B200 per the stack's last commit), which is why that commit freezes warm-started experts. #882/#883 add expert parallelism on the FSDP1 side, but only with NO_SHARD, and carry their own grad-norm partitioning, optimizer owner-split and per-rank checkpoint shards.

This PR adds expert parallelism on the FSDP2 backend so FSDP and EP compose, and the optimizer, grad-norm and checkpoint paths stay exactly as they are.

What this adds

• training.expert_parallel_size (FSDP2 only; requires a sharded fsdp_sharding, tp_size=1, no sequence parallelism, a world size divisible by it, and n_routed_experts divisible by it).
• init_distributed builds a 2-D (efsdp, ep) mesh (ep = consecutive ranks, one node for ep <= GPUs per node); ParallelConfig carries it.
• specforge/modeling/draft/moe/expert_parallel.py: the EP layout, shard_expert_parameter (the stacked [E, ...] weights become DTensor Shard(0) over ep), and the differentiable seams gather_tokens (all-gather, backward reduce-scatter) and scatter_outputs (reduce-scatter, backward all-gather).
• MoELayer.forward under EP: all-gather the group's tokens, run the replicated router on them, compute only the locally owned experts for all of them, reduce-scatter the fp32 partial outputs, add the shared expert on the local tokens. The EP axis is carved out of data parallelism, so ranks keep distinct data and the dense part of the draft is not computed twice.
• GroupedExperts.apply_expert_parallel and a local-slice dispatch for both sorted_loop and grouped_mm (sorting by expert makes this rank's experts one contiguous run of slots). Zero-valued terms keep the gathered input, combine weights and expert parameters on the graph when a rank receives no token, so every rank issues the same collectives and FSDP2 sees the same gradient set.
• FSDP2TrainingBackend: applies EP before wrapping; a per-parameter shard_placement_fn (ShardPlacementResult, available in torch 2.13) puts expert slices on the efsdp mesh and everything else on the full mesh; rescales expert gradients by 1/ep at the optimizer boundary; get_model_state_dict(full_state_dict=True) still gathers the full [E, ...] tensors, so the MoE naming converters, resume and export are unchanged.
• Guards: the FSDP1 backend and NO_SHARD reject EP with explicit errors.
• Recipe examples/configs/online/disaggregated/external/deepseek-v4-flash-dspark-moe-ep4-disaggregated.yaml (the MoE arm with backend: fsdp2, expert_parallel_size: 4: the DP4 trainer is one EP group, 16 of 64 experts per rank, experts trained) and docs (customization guide, training guide, recipe runbook, examples README, moe/DESIGN.md).

Why no owner-split, per-rank shards or norm partitioning

With expert slices as DTensors over (efsdp, ep), every parameter element lives on exactly one rank: the optimizer's local-shard sum of squares all-reduced over WORLD is already the exact global norm, Adam state is per shard, and DCP gathers full tensors on rank 0. Gradient scale: FSDP2 averages dense gradients over the world and expert gradients over efsdp = world / ep, but an expert owner's gradient already sums its whole group's tokens, so the backend multiplies expert gradients by 1/ep. The 4-rank test checks every gradient against a plain data-parallel reference.

Tests

CPU gloo, torch 2.13.0 / transformers 5.12.1 (no GPU on this machine):

• tests/test_modeling/test_moe_expert_parallel.py: 2-rank layer parity (outputs, input gradients, expert-slice gradients, router and shared-expert partial gradients summing to the reference, routing counts over the gathered tokens, official-naming checkpoint); a rank with no selected expert still completes every collective; 4-rank FSDP2TrainingBackend with ep=2 (efsdp=2) and ep=4 (efsdp=1): all gradients vs the data-parallel reference, full checkpoint in official naming on rank 0, zero-then-reload round trip through set_model_state_dict; the unsharded layer is unchanged.
• tests/test_config/test_schema.py: the new field's validators and world-size check.

python -m unittest tests.test_modeling.test_moe_expert_parallel tests.test_config.test_schema \
  tests.test_modeling.test_moe_deepseek_v4 tests.test_config.test_recipe_readme \
  tests.test_config.test_example_draft_config_wiring tests.test_runtime.test_package_architecture \
  tests.test_runtime.test_cli_config_build
# Ran 96 tests - OK (skipped=3)

Also green: tests.test_modeling.test_dflash2, test_draft_registry, tests.test_runtime.test_export, test_model_loading. black, isort and autoflake (pre-commit pins) are clean on the changed files.

Not done yet / follow-ups

• No GPU run. The next step is the DSV4-Flash EP=4 arm on one 8x B200 node (two TP2 capture servers + the DP4 trainer, per the runbook): moe/* load metrics, step time against the frozen-experts arm, and acceptance. The grouped_mm path under EP only runs on CUDA and is untested here.
• EP adds one host sync per MoE layer per micro-batch (this rank's slot bounds in the sorted routing); test_dflash_host_syncs will need an EP exception once EP recipes join that gate.
• Partial outputs are reduce-scattered in fp32; a bf16 variant would halve that traffic.
• Tokens per rank must match across the EP group (true for the DFlash family's fixed anchor shapes).
• EP combined with moe_freeze_experts (frozen, sliced experts ignored by FSDP) is wired but untested.
• torchtitan's fixed-capacity / DeepEP dispatchers could be plugged in behind the same gather/scatter seam later; this PR stays on torch 2.13 and adds no dependency.
• This takes a different route from feat: expert parallelism for draft models with sharded experts #882/feat(moe): partition the routed experts across an expert-parallel group #883 (FSDP and EP together, DTensor slices, token all-gather instead of replicated tokens). Happy to converge with @fg11991 and @sherlockwu on one of them.