exp.[rebase] MoE FFN for DFlash-family drafters (#811/#812) on current main, to run the DeepSeek-V4-Flash MoE arm

[yushengsu-thu]

Draft. This is @sherlockwu's MoE FFN stack (#811 + #812, branch kan/moe-2-dsv4) rebased onto current main, opened so the DeepSeek-V4-Flash DSpark MoE drafter arm can be trained on today's main while those PRs wait for review. All eight commits keep their original author; nothing in the MoE package was changed. It is not meant to supersede #811/#812: close this once they land or Kan rebases them himself.

What the rebase touched

• kan/moe-2-dsv4 was 151 commits behind main. git rebase origin/main replayed the eight commits cleanly except one hunk in docs/recipes/deepseek-v4-flash-dspark-disaggregated.md, where main had added the "On AMD MI355X" section at the same place the stack adds "MoE-FFN arm (ablation)". Both sections are kept, AMD first.
• main had already moved docs/advanced_features/customization.md to docs/sections/advanced_features/; git followed the rename, so the stack's MoE section landed in the new location.

Checks

CPU, torch 2.13.0 / transformers 5.12.1, on the rebased tree:

python -m unittest 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 158 tests — OK (skipped=5: CUDA-only)

The lint job will inherit the pre-existing autoflake findings on the stack's own files that @fg11991 noted on #883; they are not introduced by the rebase.

Plan for this draft

Run examples/configs/online/disaggregated/external/deepseek-v4-flash-dspark-moe-disaggregated.yaml (two TP2 capture servers + DP4 trainer on one 8x B200 node, per the runbook) on this branch and attach the run notes here: whether the recipe starts and trains on current main, moe/* load metrics, step time, and any fixes the rebase needs.

🤖 Generated with Claude Code

[yushengsu-thu]

Smoke run of the DeepSeek-V4-Flash DSpark MoE arm on this branch (GB300, 50 optimizer steps)

Setup. One 4x GB300 node (284 GB each, aarch64, lmsysorg/sglang:v0.5.18-cu130 arm64 image: torch 2.13.0+cu130, sglang 0.5.18, transformers 5.12.1, flashinfer 0.6.17, mooncake-transfer-engine-cuda13 0.3.12.post1, spec-capture patch applied). The 8-GPU recipe layout was folded onto 4 GPUs: two TP1 capture servers (GPU 0/1, the 155 GiB FP8/FP4 checkpoint fits one GB300; --moe-runner-backend flashinfer_mxfp4 works on sm_103 after the usual DeepGEMM warmup + trtllm FP4 MoE autotune) and a DP2 trainer (GPU 2/3) with training.accumulation_steps=64, so the global batch stays 128 and the recipe's math is unchanged. Data: prepare_data.py --dataset sharegpt --sample-size 8000 (7997 prompts after filtering), 2 prompt epochs, max_steps=50, no wandb. Everything else is the checked-in yaml.

Three 50-step runs with identical overrides, capture servers kept up and only store_id changed:

run	step time (steps ≥5)	ce@1	ce@3	ce@10	ce@25	ce@50	loss@50	acc@50	walk_accepted_length@50	l1@50	experts unused @10 → @50	grad_norm max
deepseek-v4-flash-dspark-moe (as checked in)	8.9 s	58.5	494	37.3	23.2	14.9	3.24	0.028	1.035	1.96	0.86 → 0.74	108
deepseek-v4-flash-dspark (dense control)	6.6 s	60.3	232	38.8	23.6	15.8	3.38	0.020	1.021	1.98	–	149
MoE + dflash_config.moe_aux_loss_coeff: 0.01	8.9 s	57.8	499	29.2	26.4	18.3	3.65	0.008	1.012	1.98	0.71 → 0.57	107

What works.

• The rebased stack trains end to end on current main: MoE draft builds, FSDP (SHARD_GRAD_OP) over the stacked [E, out, in] experts, grouped GEMM dispatch on CUDA, deferred noaux_tc bias updates, moe/* metrics on the scalar channel, and the step-50 checkpoint is written in the official per-expert naming (960 layers.{l}.mlp.experts.{i}.w{1,2,3}.weight keys + 5 gate.bias) through the FSDP backend's state-dict boundary.
• GPU unit tests on the box: test_moe_deepseek_v4 (incl. test_grouped_mm_matches_sorted_loop on CUDA), test_dflash2, test_dflash_offline_launch — 55 tests, all pass on aarch64/GB300.
• Cost at this layout: MoE arm 8.9 s/step vs dense 6.6 s/step (≈1.35x) and about 140–154 GB per trainer rank; the MoE training_state.pt is 112 GB vs 15 GB dense.

What to watch.

• Early-step instability (ce spike to 200–500 at step 2–3, grad norm 100–150 before clipping at 1.0) shows up in both arms, so it is the recipe/data at LR 6e-4 from scratch, not the MoE path. Both arms are at the same place after 50 steps (ce ≈ 15, acc ≈ 2–3%, walk_accepted_length ≈ 1.0, i.e. no draft contribution yet) — 50 steps is far too short to say anything about accept length; the +5% result in DeepSeek-V4 MoE FFN for DFlash-family drafters + DSV4-Flash ablation arm #812 was at ~1.9k steps.
• Routing collapse from scratch, exactly as the noaux_tc.py docstring warns: by step 2, 74% of the 64 experts receive no tokens and moe/load_max_ratio sits at the 64/6 = 10.67 ceiling; the selection bias grows linearly (0.064/step at this accumulation) but had only recovered to 74% unused by step 50. moe_aux_loss_coeff=0.01 helps (57% unused at step 50, trend still improving) but the aux term is ~0.013 against a loss of ~4, so 0.01 is probably too small for a from-scratch drafter; the checked-in ablation JSON leaves it at 0. Worth stating in the runbook which balancing setting the validated run used, or defaulting the MoE arm to a non-zero aux coefficient.

Logs: train_run1.log (MoE), train_dense.log, train_moe-aux.log on the devbox; summary table produced by a small log parser. Happy to run the full 2-epoch recipe on this layout if useful (≈1.9k steps × 9 s ≈ 4.7 h).

🤖 Generated with Claude Code