exp - feat(moe): qwen3 preset and the Qwen3.8-27B DFlash2 MoE ablation arm (stacked on #811/#812)

[yushengsu-thu]

Draft, stacked on #811 / #812. This branch is main + a merge of kan/moe-2-dsv4 (the MoE FFN skeleton and the DeepSeek-V4 preset; eight commits, one docs conflict resolved by keeping both sections) + one commit of its own (1e1c335d). The MoE package itself is unchanged from #812. Once #811/#812 land, this PR shrinks to the last commit; until then review that commit (git diff 6e3f0842..1e1c335d).

What this adds

A MoE arm for the Qwen3.8-27B DFlash2 drafter, trained on the existing FSDP backend (no expert parallelism), as a dense-vs-MoE ablation pair in the style of #812's DSV4-Flash arm.

• qwen3 MoE preset (modeling/draft/moe/presets.py): the Qwen3-MoE routing recipe — softmax scores over all experts, top-k combine weights renormalized (norm_topk_prob), no scaling, no shared expert, plain SwiGLU. Balancing is the aux-loss-free noaux_tc selection bias already in the package; dflash_config.moe_aux_loss_coeff adds the complementary balance loss when a from-scratch run collapses onto a few experts.
• num_experts as the Qwen spelling of the expert count (moe/config.py): the package's architecture keys are "the target checkpoint's native HF names", and Qwen checkpoints say num_experts where DeepSeek says n_routed_experts. Both resolve to MoEConfig.n_routed_experts; stating both with different values is an error.
• configs/qwen3.8-27b-dflash2-moe.json: the checked-in dense DFlash2 draft plus the MoE keys only — 16 routed experts of width 4352, top-4, so the activated FFN width (4 × 4352) equals the dense 17408 at 4× the FFN parameters (about 5.9B total vs 1.9B dense). Training knobs under dflash_config: moe_bias_update_rate: 0.001, moe_dispatch: grouped_mm.
• examples/.../managed-local/qwen3.8-27b-dflash2-moe-4server-dp4-disaggregated.yaml: identical to the dense 4-server/DP4 recipe except the draft JSON and run/store names, so the A/B diff is the FFN. Capture servers, flow control and Mooncake settings are shared.
• Docs: customization.md (preset + spelling), the Qwen3.8-27B runbook (new "MoE-FFN arm" section with the memory and balancing notes), examples/configs/README.md, moe/DESIGN.md.

Why FSDP is enough here

Expert weights are stacked [E, out, in] parameters that FSDP1 (use_orig_params) shards like any other; with SHARD_GRAD_OP the unsharded bf16 draft is about 12 GB per rank during a step, well within the 4+4 layout's trainer budget. grouped_mm compute scales with routed tokens, not expert count. EP (#882/#883) only becomes necessary for DeepSeek-class expert sets.

Tests

tests/test_modeling/test_moe_qwen3.py (13 tests, CPU): preset recipe; num_experts/n_routed_experts agreement, conflict and dense cases; checked-in config = dense config + MoE keys and the iso-activated-width invariant; the recipe differs from the dense arm only in draft_model_config, cache_dir, run_id, output_dir, control_dir, consumer_state_dir; softmax top-k weights renormalize to one and match a per-token dense reference; bf16 forward keeps the fp32 balance bias and the bias moves selection only; deferred balance update; official checkpoint naming round-trips with no shared-expert keys; serving_fields(); DFlash2 integration (two MoE layers train, DFlash2's zero-initialized conv projections survive MoE init, bias update applied on the next forward, strict reload through the naming boundary, dense DFlash2 config still builds Qwen3MLP).

Run locally on CPU with torch 2.13.0 / transformers 5.12.1:

python -m unittest tests.test_modeling.test_moe_qwen3 tests.test_modeling.test_moe_deepseek_v4 \
  tests.test_modeling.test_dflash2 tests.test_modeling.test_draft_registry \
  tests.test_config.test_example_draft_config_wiring tests.test_config.test_recipe_readme \
  tests.test_config.test_schema tests.test_runtime.test_package_architecture \
  tests.test_runtime.test_export tests.test_runtime.test_model_loading
# Ran 171 tests — OK (skipped=5: CUDA-only)

autoflake, isort and black (pinned versions from .pre-commit-config.yaml) are clean on the changed Python files.

Not yet run: a GPU smoke of the recipe itself (the trainer's moe/* load metrics, step time and acceptance against the dense arm). That is the next step before this leaves draft.

Out of scope / follow-ups

• Serving: SGLang's DFlash2 draft path (dflash.py) builds a dense DFlashMLP, so the MoE arm is measured training-side (moe/* metrics, spec_generate) until a MoE-aware DFlash draft class exists on the serving side (SGLANG_EXTERNAL_MODEL_PACKAGE can host it without patching SGLang).
• Expert parallelism (feat: expert parallelism for draft models with sharded experts #882/feat(moe): partition the routed experts across an expert-parallel group #883).
• A Qwen3.5-MoE preset (sigmoid-gated shared expert, aux loss) if a Qwen3.5 MoE target is trained.

🤖 Generated with Claude Code

[yushengsu-thu]

Closing for now: the MoE-drafter training we need runs on the existing #811/#812 stack (DeepSeek-V4-Flash DSpark MoE arm, one 8-GPU node) without this change. The qwen3 preset and the Qwen3.8-27B MoE ablation config remain on the fork branch yushengsu-thu/SpecForge:feat/qwen3.8-27b-dflash2-moe in case a Qwen-family MoE drafter is picked up later.

[yushengsu-thu]

Reopening: we are going to run the Qwen3.8-27B MoE arm on H200 with settings identical to the dense arm (same recipe, only the draft JSON differs) and attach the results here.