MoE FFN support for DFlash-family drafts by yushengsu-thu · Pull Request #1003 · sgl-project/SpecForge

yushengsu-thu · 2026-10-21T06:44:27Z

Summary

An MoE FFN for DFlash-family drafts (DFlash, Domino, DSpark), starting with the Qwen3.5-MoE family (Qwen3-Next / Qwen3.5 / Qwen3.8). A draft config turns it on by naming a moe_preset; every other draft keeps its dense MLP. Load balancing follows torchtitan: a selection bias moved once per optimizer step, plus a small load-balance loss.

I re-write +factor the PRs (#811, #812, #800) to make it more maintainable/readable and easier add new MoE layer to draft models in the future

Layout

lines of code:
production code: +557 / −1
doc and config: +155 / −0
ci test: +861 / −5

specforge/modeling/draft/moe/
  layer.py     MoELayer, build_ffn
  experts.py   grouped_mm_experts, RoutedExperts, SharedExpert
  router.py    TopKRouter (top-k selection, selection_bias, per-step expert counts)
  balance.py   load_balance_loss, dispatch_balance, count_tokens_per_expert;
               set_moe_token_mask, take_moe_aux_loss, step_moe_balancing (the hooks the model calls)
  families/qwen3_5_moe.py   the MoELayer from a Qwen3.5-MoE config (registered as moe_preset "qwen3_5_moe")

configs/qwen3-8b-moe-dspark.json                           # a 2-layer 8x1536 top-1 DSpark draft for Qwen3-8B
examples/configs/online/disaggregated/managed-local/qwen3-8b-moe-dspark.yaml   # its online recipe

Wiring:

# specforge/modeling/draft/dflash.py: the decoder layer's FFN, dense unless the draft names a moe_preset
-        self.mlp = kernels.make_mlp(config)
+        self.mlp = build_ffn(config, dense=kernels.make_mlp)

Note: the moe code 90% is quick simple and intuitive. However, to train moe and and balance experts, we must introduce selection_bias (dsk's) and axu_loss (Qwen's) mechanism in balance.py to change the router score. Those parts are not that intuitive. If you feel confused about that part. Please kindly let me know and I can divide the PRs into:

pr1. moe e2e training
pr2. moe router balance - axu_loss
pr3. moe router balance - election_bias
for you

Design notes

layer.py: the MoE layer. MoELayer puts the router, the routed experts and the optional shared expert together; build_ffn gives each draft layer a dense MLP or an MoE layer.
experts.py: the experts. RoutedExperts computes the outputs of the experts each token was routed to; SharedExpert is an always-on expert added on top.
router.py: the router. TopKRouter picks each token's top-k experts and keeps the selection bias and the step's expert counts.
balance.py: load balancing. Provides the balance loss, the balance metrics, and the hooks the model calls to add the loss and update the selection bias once per optimizer step.
families/: model families. Each family builds the MoE layer from its model's config, selected by moe_preset; qwen3_5_moe.py is the Qwen3.5-MoE family and reads the balancing settings moe_router_aux_loss_coef, moe_router_aux_loss_type (batch_wise by default), moe_router_selection_bias and moe_router_bias_update_rate.

Verification

Both runs share this setup:

Model: Qwen3-8B target; DSpark draft, 2 layers, 8 experts × dim: 1536, top-1, no shared expert (configs/qwen3-8b-moe-dspark.json).
Dataset: PerfectBlend regenerated by Qwen3-8B (temperature 0, no thinking), 1.42M conversations, max length 8192.
Training: online, 2 SGLang capture servers + 6 DDP trainers on 8 H200; 132 samples/step; 2 epochs (21,528 steps); lr 6e-4, cosine, 4% warmup; bf16; seed 42.
Runs: 2 epoch

Result: the two runs differ very little. Walk accept length is 4.27 vs 4.26 after epoch 1 and 4.64 vs 4.60 after epoch 2, slightly in favor of #1003, with cross-entropy and accuracy just as close. Both runs stay balanced, within 2% of 1 at the end of each epoch.

Out of scope

Serving: the SGLang draft classes, the normalizer's MoE branch and the served selection_bias are the stacked PR serving: SGLang draft classes for the MoE DSpark/DFlash drafts (stacked on #1003) #1016.
A DeepSeek MoE family: the router and experts keep its contract and seams (score_fn, selection_bias, route_scale, the expert activation), but no such family is built here.
Activation checkpointing for MoE drafts, and torchtitan's token-weighted loss averaging across micro-batches.

Revised version

MoE FFN for DFlash-family drafts, staged: module, selection bias, aux loss (final) #1031

configs/qwen-8b-moe.json is qwen3-8b-dspark.json plus the Qwen3-MoE config fields (num_experts, num_experts_per_tok, moe_intermediate_size, norm_topk_prob, router_aux_loss_coef, decoder_sparse_step, mlp_only_layers). Nothing reads them yet, so the draft still builds the dense FFN. The managed-local recipe trains it online on one 8-GPU node: one capture server on GPU 0 and seven trainer ranks.

Qwen3DFlashDecoderLayer now builds its FFN with build_ffn(config, dense=kernels.make_mlp) from the new specforge.modeling.draft.moe package. A draft config without a moe_preset gets dense(config), the same native or Liger MLP as before, so every existing draft builds unchanged. A config that names a moe_preset raises NotImplementedError until the MoE layer lands.

…-8b-moe

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

… recipe Mirror exp/qwen3-8b-dspark-moe-8expert-1activate-dense-same: 2 layers, GQA 8/2, block 16, no confidence head, is_causal, 8 experts x 1536 top-1 without renorm, no shared expert, aux 0.001, and its 12-epoch recipe with Liger. Main has no ddp_bucket_cap_mb, flatten/cache_vocab_projection or consumer_dispatch, so those stay out; router centering is not part of this branch.

yushengsu-thu · 2026-10-23T05:23:17Z

training and save ckpt commits config until here

planning.py rejects model.use_liger_kernel=true for dspark: on main only the dflash algorithm wires Liger kernels into its draft. The exp branch could set it because it merged sgl-project#963. The first online smoke of this recipe failed on it.

TopKRouter keeps a Switch-style load-balance loss (the MoE aux loss) on gate.aux_loss after every training forward: E * sum_e load_e * mean_prob_e over the kept tokens, K when balanced. Scores are L1-normalized per token so sigmoid and sqrtsoftplus routers share the formula. The qwen3_5_moe family passes router_aux_loss_coef through; nothing consumes the loss yet.

MoELayer.token_mask marks this batch's kept slots; the training forward passes it to the router so the load-balance loss skips padding blocks. _forward_draft_blocks sets it for DFlash, Domino and DSpark from block_keep_mask. MoE drafts refuse activation checkpointing: a recompute would not see this per-batch state.

take_moe_aux_loss pops the routers' per-forward loss as (coefficient-weighted mean, raw mean) and clears the per-batch state; the weighting then follows each objective's own normalization. OnlineDFlashModel folds the weighted mean where each objective is built: the DFlash numerator the trainer normalizes globally, Domino's loss and loss_terms, DSpark's plain loss; the raw mean is reported as moe/aux_loss. The router clears aux_loss at every forward so no training value survives into eval. Dense drafts are untouched.

…step TopKRouter keeps a persistent fp32 selection_bias that shifts the top-k picks and a per-step tokens_per_expert count. The counts accumulate in forward; at the step boundary the trainer calls the strategy's on_optimizer_step, the DFlash-family model forwards it to step_moe_balancing, which pools every router's counts in one all-reduce and applies the zero-mean sign rule (rate x sign(mean - count)). Enabled per draft config with dflash_config.moe_router_selection_bias; the rate is moe_router_bias_update_rate (0.001 per optimizer step, as in DeepSeek-V3).

moe_router_selection_bias true with the DeepSeek-V3 rate of 0.001 per optimizer step, next to the 0.001 load-balance loss.

yushengsu-thu · 2026-10-23T07:32:18Z

balancing related features done and use step-level (instead of micro-step) loss and mechanism

yushengsu-thu · 2026-10-26T03:00:45Z

above verification: moe draft model training result (2epoch)

…Strategy

…token Same draft as qwen3-8b-moe16x768-16A1-dspark.json with the 12288-wide FFN split eight ways instead of sixteen (same total expert parameters, twice the active width per token), router balancing at the sgl-project#1003 defaults (aux loss 0.005, selection-bias step 0.001).

maocheng23

I'd recommend splitting this PR into two: one for the MoE module itself and one for load balancing (the split you offered in the description works for me).

The module part should be self-contained and easy to review. The balancing part reaches into several stages of training: the router forward, the loss assembly of each objective, the strategy, and the trainer step. That makes it feel a bit invasive and hard to follow. With a separate PR we could land the module quickly and give the balancing design a focused review and cleanup.

For the balancing PR, it may be worth following torchtitan's structure:

inject the aux-loss gradient from inside the router, so objectives don't need to fold it into the loss;
pass the token mask as an explicit forward argument instead of setting it on the router;
keep the step-boundary work (count all-reduce, bias update, metrics) in a single hook that the trainer registers once.

maocheng23 · 2026-10-26T07:36:33Z

+            # Qwen3.8's config omits it; transformers always renormalizes.
+            norm_topk_prob=getattr(config, "norm_topk_prob", True),
+            initializer_range=std,
+            # torchtitan's units: 1 when balanced, one coefficient per layer.


We are supporting torch titan?

No torchtitan dependency; the comment only cites the torchtitan loss units (1 when balanced) becasue we refer it

maocheng23 · 2026-10-26T07:43:33Z

+from specforge.modeling.draft.moe.balance import (
+    set_moe_token_mask,
+    step_moe_balancing,
+    take_moe_aux_loss,


why moe only in dflash?
This should be a common part for other draft models, e.g. eagle?

Here I think the coupling part is the baclancing?

The coupling was the balancing; #1031 moves it to a trainer-registered hook. EAGLE3 wiring is one line, follow-up.

jiapingW · 2026-10-27T08:08:49Z

@yushengsu-thu There are three concerns that I think need to be handled.

Introduce a small router base class with explicit balancing capabilities.

Replace discovery through hasattr(module, "tokens_per_expert") with an explicit contract rather than assuming auxiliary losses, masks, selection bias, and update methods are all available. Keep the current algorithm in TopKRouter; complex model families can provide separate implementations.
Separate auxiliary-loss computation, selection-bias updates, and diagnostics.

When aux_loss_coef == 0, skip auxiliary-loss computation and its per-microbatch collectives. If selection-bias updates are enabled, keep local expert counts and aggregate them at the optimizer-step boundary; moe/step_balance can reuse those counts. Using a selection bias for routing does not by itself require updating it. If both mechanisms are disabled, skip balancing statistics unless diagnostics explicitly require them. Diagnostic-only calculations should run without gradients. For whole-step metrics, accumulate counts from every relevant invocation in the reported step, while performing diagnostic aggregation and reporting at the requested logging frequency.
Separate router computation from state updates for activation checkpointing.

Router forward currently accumulates expert counts and stores auxiliary losses, while the wrapper clears the token mask and auxiliary-loss state. Backward recomputation could count tokens twice, lose the original mask, or use a different cumulative load.

The loss is 1 when balanced and the coefficient applies per layer; the torchtitan reference read as a dependency.

The MoE layer, experts, top-k router, the qwen3_5_moe family and the build_ffn switch from sgl-project#1003, without the load balancing: no selection bias, no load-balance loss, no per-step hook. Those follow as stage 2 and stage 3 of this PR.

…ient 0 MoELayer takes aux_loss_coef and aux_loss_type (microbatch_wise, or the default batch_wise over the step's counts pooled across ranks) and, in training with a non-zero coefficient, leaves this forward's loss in aux_loss for the balancing hook to fold before backward. At coefficient 0 nothing is computed: no graph and no per-micro-batch collective. The loss formula and units are sgl-project#1003's, moved from the router to the layer.

…the step hook MoEBalancing.before_backward collects each MoE layer's loss times its coefficient, adds the sum to the loss and, when the strategy reports loss_terms, to the numerator scaled by the denominator, and logs the raw mean as moe/aux_loss; TrainerCore runs the registered hooks before every backward. The arithmetic is sgl-project#1003's if_add_moe_aux_loss, done once here instead of at the end of each objective.

…_wise 0.005 in the DSpark recipe dflash_config.moe_router_aux_loss_coef and moe_router_aux_loss_type (batch_wise by default, microbatch_wise the alternative) reach MoELayer from the family builder; the default coefficient 0 keeps drafts without the keys loss-free. The Qwen3-8B MoE DSpark recipe trains with batch_wise 0.005 next to the expert bias, as sgl-project#1003 did.

yushengsu-thu · 2026-10-28T02:12:02Z

@maocheng23 Done in #1031, one concern per commit:

Split: the module lands alone as stage 1. afd6ad4
Token mask as a forward argument; nothing is set on routers. 0e31878
One hook the trainer registers once: count all-reduce, bias update, metrics. e8dda61
Aux loss: that hook folds it before backward instead of injecting, since SpecForge's trainer normalizes by the global denominator. 117e99c 8f8d5e2

yushengsu-thu · 2026-10-28T02:14:14Z

@jiapingW Done in #1031, one concern per commit:

Router base class: Router declares the contract (forward, expert bias, its update); found by type, not hasattr. e5c2991
Counts stay local on MoELayer and are pooled once per optimizer step by the hook, which also updates the bias and reports moe/step_balance. afe63b2 e8dda61
Aux loss at coefficient 0: not computed, no collectives; bias_update_rate 0 keeps a bias without updating it; the metric is detached. 117e99c 8f8d5e2
Router forward is stateless; the mask is a forward argument. The layer still counts in forward, so MoE drafts keep refusing activation checkpointing. 0e31878

…draft starts a fresh window tokens_per_expert is not saved, and from_pretrained builds the model on the meta device, so without this the counts of a warm-started draft (model.draft_checkpoint_path) held uninitialized memory until the first optimizer step reset them. sgl-project#1003 zeroed them in the router's reset_parameters; the layer does the same now. The from_pretrained test also checks that a trained expert_bias round-trips.

yushengsu-thu · 2026-10-28T06:17:55Z

@jiapingW @maocheng23
All Revised code is here: #1031
stacked PRs:

MoE FFN for DFlash-family drafts: the module #1033 (module),
MoE drafts: the expert bias (stacked on #1033) #1034 (expert bias),
MoE drafts: the load-balance loss (stacked on #1034) #1035 (load-balance loss)

Rebased from sgl-project#1003 onto sgl-project#1031 (dev-moe-ffn): the served router reads dflash_config.moe_router_enable_expert_bias and loads gate.expert_bias, the names sgl-project#1031 gives the DeepSeek-style bias in its checkpoints. Co-Authored-By: Claude Mythos 5.1 <noreply@anthropic.com>

select now receives the shifted scores (choice) and only picks the top-k, so a family that overrides select, such as DeepSeek-V3's group-limited routing, needs no knowledge of the bias. The weights still come from the unbiased scores. No behavior change: the bitwise probe against sgl-project#1003 stays identical.

…t#1003 code, 5-epoch cosine) Co-Authored-By: Claude Mythos 5.1 <noreply@anthropic.com>

…p script and py-spy dump Training finished 2026-11-02 12:35 UTC (53,820 steps, 2.45 s/step); epochs 1-5 on HF yushengsu/dense-1-8-only-5-epoch. Train acc at the last step of each epoch 0.6032 / 0.6441 / 0.6582 / 0.6714 / 0.6624 (1,000-step means 0.6117 / 0.6420 / 0.6564 / 0.6610 / 0.6696). The managed producer hung in streaming_ref_channel.poll at shutdown again (PR sgl-project#1003 code). Co-Authored-By: Claude Mythos 5.1 <noreply@anthropic.com>

yushengsu-thu added 5 commits October 20, 2026 14:54

moe: add routed and shared SwiGLU experts with a grouped_mm path

0d8083d

moe: add a top-k router with DeepSeek-V4 routing seams

3b8493c

moe: build MoELayer from the qwen3_5_moe family and enable it in qwen…

23b5dc0

…-8b-moe

yushengsu-thu mentioned this pull request Oct 21, 2026

MoE FFN support for DFlash-family drafts (WIP) yushengsu-thu/dev-moe-draft#2

Closed

yushengsu-thu changed the title ~~MoE FFN support for DFlash-family drafts~~ [WIP] MoE FFN support for DFlash-family drafts Oct 21, 2026

yushengsu-thu marked this pull request as ready for review October 21, 2026 20:42

yushengsu-thu requested review from FlamingoPg, FrankLeeeee, shuaills and sleepcoo as code owners October 21, 2026 20:42

Copilot AI balanced review requested due to automatic review settings October 21, 2026 20:42

Copilot AI reviewed Oct 21, 2026

yushengsu-thu mentioned this pull request Oct 21, 2026

[Feature] MOE Draft model Training #979

Open

3 of 4 tasks

yushengsu-thu force-pushed the dev-moe-draft branch from e5ec8db to 23b5dc0 Compare October 23, 2026 03:06

yushengsu-thu added 6 commits October 22, 2026 23:06

configs: turn on the selection bias in qwen-8b-moe

0f59e18

moe_router_selection_bias true with the DeepSeek-V3 rate of 0.001 per optimizer step, next to the 0.001 load-balance loss.

yushengsu-thu added 5 commits October 23, 2026 15:39

moe: add torchtitan's batch-wise load-balance loss as a router option

28c4e0b

moe: log the step's pooled load balance next to the per-micro-batch loss

bf453d9

moe: switch the load-balance loss and step balance to torchtitan's units

17a1d68

configs: balance qwen-8b-moe with the batch-wise loss next to the bias

b78c2c3

moe: move the balance math and training hooks into balance.py

68d1e22

yushengsu-thu added 2 commits October 25, 2026 19:39

moe: pick the aux-loss load next to the loss formula

7c0f46e

moe: keep the token mask on the router

b7da853

yushengsu-thu added 2 commits October 26, 2026 00:02

Merge branch 'main' into dev-moe-draft

6a51ee0

tests: make the DDP accumulation test's strategy implement DraftTrain…

cde2f05

…Strategy

yushengsu-thu mentioned this pull request Oct 26, 2026

[Exp.] DSpark v2 MoE draft (stacked on #1003) #1022

Draft

4 tasks

maocheng23 reviewed Oct 27, 2026

View reviewed changes

moe: say what the balancing comments mean without naming torchtitan

ca7062a

The loss is 1 when balanced and the coefficient applies per layer; the torchtitan reference read as a dependency.

maocheng23 mentioned this pull request Oct 27, 2026

[6/6] feat(attention): DSA and QSA draft attention training #1030

Open

yushengsu-thu mentioned this pull request Oct 27, 2026

MoE FFN for DFlash-family drafts, staged: module, selection bias, aux loss (final) #1031

Open

This was referenced Oct 28, 2026

MoE FFN for DFlash-family drafts: the module #1033

Open

MoE drafts: the expert bias (stacked on #1033) #1034

Open

MoE drafts: the load-balance loss (stacked on #1034) #1035

Open

maocheng23 mentioned this pull request Oct 29, 2026

[9/10] Version draft settings and separate serving serialization #1045

Draft

yushengsu-thu added a commit to yushengsu-thu/SpecForge that referenced this pull request Oct 31, 2026

experiments: dense-1-8-only-5-epoch recipe and scripts (PR sgl-projec…

77dc157

…t#1003 code, 5-epoch cosine) Co-Authored-By: Claude Mythos 5.1 <noreply@anthropic.com>

Conversation

yushengsu-thu commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Layout

Design notes

Verification

Out of scope

Revised version

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu commented Oct 23, 2026

Uh oh!

yushengsu-thu commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yushengsu-thu commented Oct 26, 2026

Uh oh!

maocheng23 left a comment

Choose a reason for hiding this comment

Uh oh!

maocheng23 Oct 26, 2026

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

maocheng23 Oct 26, 2026

Choose a reason for hiding this comment

Uh oh!

yushengsu-thu Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

jiapingW commented Oct 27, 2026

Uh oh!

yushengsu-thu commented Oct 28, 2026

Uh oh!

yushengsu-thu commented Oct 28, 2026

Uh oh!

yushengsu-thu commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yushengsu-thu commented Oct 21, 2026 •

edited

Loading

yushengsu-thu commented Oct 23, 2026 •

edited

Loading

yushengsu-thu Oct 27, 2026 •

edited

Loading

yushengsu-thu commented Oct 28, 2026 •

edited

Loading