Skip to content

MoE support for Qwen drafts (A1, A2, B) - #993

Closed
yushengsu-thu wants to merge 22 commits into
sgl-project:mainfrom
yushengsu-thu:moe-support-qwen
Closed

yushengsu-thu wants to merge 22 commits into
sgl-project:mainfrom
yushengsu-thu:moe-support-qwen

Conversation

@yushengsu-thu

@yushengsu-thu yushengsu-thu commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Sparse-MoE FFN support for SpecForge draft models, for the Qwen family:

This is #991 without its DeepSeek-specific parts: parts A1, A2 and B of the split plan in #990 as 18 commits, plus four follow-up commits, so the stack can be reviewed one step at a time. The DeepSeek-only options (sqrtsoftplus and sigmoid scores, group-limited routing, noaux_tc balancing, the SwiGLU clamp), the lean warm start and the DeepSeek-V4 family (C1) follow in a later PR through the same registries and config fields.

It keeps #811's contracts and registries for the router, the routed experts and the shared expert, with TopKRouter, GroupedExperts and SwiGLUSharedExpert as the built-in implementations. Docstrings and comments are one line, with the explanations in moe/DESIGN.md.

Split plan

PR Commits Head Depends on Changes
A1 1–12 ef134da main +1,818 / −17
A2 13–16 4f007ba A1 +594 / −42
B 17–18 c05ebc1 A2 +583
follow-up 19 3bda0a9 B +10 / −5
follow-up 20 c44212a 19 +101 / −92
follow-up 21 15459cb 20 +20 / −8
follow-up 22 64f2c61 21 +66 / −73

B only adds family files; it does not change the framework core.

Commit 19 binds the class a registry lookup returns to a name before using it, in the four component builders (A1's files) and in serving_architecture (A2's). When the stack is split, each change goes with the part that owns its file.

Commit 20 moves the aux loss from the aux_loss balance controller into the router. TopKRouter adds moe_aux_loss_coeff times the loss that dflash_config.moe_aux_loss_type selects (only "switch", transformers' loss, for now), and balance only names a policy that shifts the scores used to pick experts. The Qwen preset sets balance="none"; its aux loss, layer output and gradients are bitwise unchanged. The moe/load_frac_* metrics are dropped, since moe/load_max_ratio carries the same information. When the stack is split, commit 20 goes with A2, which added aux_loss balancing, except the preset and serving-config change, which goes with B.

Commit 21 turns the aux loss on by default: dflash_config.moe_aux_loss_coeff defaults to 0.001, the router_aux_loss_coef of transformers' Qwen MoE configs, and 0 turns it off. The Qwen recipe already sets 0.001, so its training is unchanged. A later family whose recipe balances without an aux loss, such as DeepSeek-V4's aux-loss-free bias, sets the coefficient to 0.

Commit 22 moves the shared expert from moe/shared.py into moe/experts.py, next to the routed experts whose swiglu it already used, and removes shared.py. The classes, registries and shared_experts.* parameter names do not change. When the stack is split, it goes with A1, which added both files, and A2's sigmoid shared-expert gate then edits experts.py.

Verification

  • Every commit passes lint.
  • At the end of each part, the tree equals the reviewed reference, the full local suite has no failures beyond main's, and main's three GPU CI steps pass on 8× H200, one job at a time:
Part end Local suite Capture gate Liger unittest discover
A1 (ef134da) 1,247 passed; same failures as main 3 OK 18 OK 1,592 OK
A2 (4f007ba) 1,260 passed; same failures as main 3 OK 18 OK 1,605 OK
B (c05ebc1) 1,268 passed; same failures as main 3 OK 18 OK 1,613 OK
19 (3bda0a9) 1,268 passed; same failures as main 3 OK 18 OK 1,613 OK
20 (c44212a) 1,269 passed; same failures as main 3 OK 18 OK 1,614 OK
21 (15459cb) 1,270 passed; same failures as main 3 OK 18 OK 1,615 OK
22 (64f2c61) 1,270 passed; same failures as main 3 OK 18 OK 1,615 OK
  • FlashAttention is not built on the test box (main's CI compiles flash-attn 2.8.3.post2), so its tests are skipped.

Notes

  • Serving the Qwen MoE drafter needs SGLang's Qwen3MoeDSparkModel, which is not upstream yet.
  • The DeepSeek-V4 family (C1) and DeepSeek-V4's official drafter (C2) stay in MoE draft support #990 and follow in separate PRs.

Part 1 of 5 of the sgl-project#811 skeleton commit (ec6dfda), split for review.

- _registry.py: the named-registry helper the MoE registries share.
- config.py: MoEConfig, the moe_preset registry and resolve_moe_config().
  Architecture keys use the target checkpoints' native HF names at the
  draft-JSON top level; training-only knobs live under dflash_config as
  moe_*.
- __init__.py: the package entry point, exporting the config API.
- tests: config resolution, preset defaults and overrides, validation.
Part 2 of 5 of the sgl-project#811 skeleton commit (ec6dfda), split for review.

- router.py: RoutingResult, the score-function registry and the Router
  contract (hidden states -> top-k choices + combine weights) with its
  registry.
- balance.py: the BalanceController contract (selection-only score shift,
  optional aux loss, metrics, deferred updates) and the "none" policy.
- experts.py: the RoutedExperts contract and the experts-backend registry.
- shared.py: the SharedExpert contract and its registry.
- tests: stub components for every contract, registered by the tests, and
  the registries' error messages.
…egies

Part 3 of 5 of the sgl-project#811 skeleton commit (ec6dfda), split for review.

- layer.py: MoELayer composes gate (router + balance controller), experts
  and shared_experts under the official attribute names; build_ffn() is the
  dense/MoE switch, so dense drafts get the kernel provider's MLP verbatim.
- hooks.py: apply_pending_balance_updates, collect_moe_aux_loss and
  collect_moe_metrics over any module tree.
- dflash.py: decoder layers build their FFN through build_ffn;
  _init_weights reaches the MoE layers' bare Parameters; the model forward
  applies pending balance updates in training. The moe import is placed in
  isort order (the original commit failed the isort hook here).
- strategies/base.py: the DFlash and DSpark strategies add moe/* load
  metrics.
- tests: layer composition, the model-level hooks and the DFlash wiring.
Part 4 of 5 of the sgl-project#811 skeleton commit (ec6dfda), split for review.

- state_dict.py: the to/from_checkpoint_state_dict boundary between the
  modules' native layout and the official file naming, with a registry of
  converters. FSDP's full-state-dict hooks need the module's own FQNs, so
  the rename cannot live in state_dict().
- init.py: WarmStartPlan, select_target_experts and plan_warm_start.
- The boundary is applied in the FSDP backend's save and load,
  materialize_draft, warm start, and both exporters.
- tests: the conversion boundary and warm-start planning.
Part 5 of 5 of the sgl-project#811 skeleton commit (ec6dfda), split for review.
Together the five parts are that commit, except for its two lint fixes:
the moe import in dflash.py is in isort order (part 3), and the tests do
not import the unused MoEConfig.

- moe/DESIGN.md: responsibility, why the draft matches the target's MoE,
  the package layout, the contracts that matter, extension points.
- customization.md: the MoE section for DFlash-family drafts.
RoutingResult carries the differentiable pre-selection scores (optional), and
BalanceController.observe() receives the whole result instead of bare counts,
so a policy can build an auxiliary balance loss for the same forward.
@yushengsu-thu yushengsu-thu mentioned this pull request Oct 19, 2026
3 of 4 tasks
- base.py: one Registry, RoutingResult (moved from router.py), MetricValue
  (moved from balance.py) and the TrainingState contract the components
  share; _registry.py goes.
- config.py: architecture keys at the draft-JSON top level and training
  keys under dflash_config.moe_* as explicit maps; expert counts must be
  positive integers; a draft is MoE only if it names a moe_preset. Fields
  for implementations this PR does not add (group-limited routing, the
  SwiGLU clamp, a bias update rate) are dropped; the shared-expert gate
  returns with its implementation. experts_backend and shared_expert stay
  with the expert registries.
- layer.py: build_ffn refuses n_routed_experts without a moe_preset
  instead of falling back to a dense FFN.
- router, balance, experts, shared: import the shared types from base.
sgl-project#811's router, routed-expert and shared-expert contracts and their
registries stay; this adds the built-in implementations, so a draft that
names a preset builds a working layer and the tests run real code.

- router.py: TopKRouter, registered as "topk", picks the top-k on the
  balance controller's selection scores and combines with the raw scores,
  all in fp32; a softmax score function is registered and
  get_score_function goes (SCORE_FUNCTIONS.get does the same). Routers
  implement TrainingState and receive tokens_per_sample.
- experts.py: GroupedExperts, registered as "grouped", holds the routed
  experts as three stacked [E, out, in] parameters with sorted-segment
  dispatch (sorted_loop, or torch._grouped_mm on CUDA) and returns a
  fresh fp32 accumulator, which the RoutedExperts contract now requires.
- shared.py: SwiGLUSharedExpert, registered as "swiglu". Both expert
  contracts implement TrainingState.
- layer.py: MoELayer adds the shared expert into the fp32 routed sum,
  casts once, and passes tokens_per_sample to the router.
- tests: the router, experts and shared-expert stubs go; the layer tests
  run the built-in components, with reference checks of the routed sum,
  the grouped_mm path and the single cast, and a check that the layer
  builds the implementations the config names.
- layer.py: MoELayer drives the router, balance controller, routed
  experts and shared expert through one TrainingState loop (components(),
  in a fixed order) for observe, apply_pending_update, aux_loss, metrics
  and reset_parameters, where it used to call the balance controller
  alone; MoELayer.apply_pending_balance_update is now
  apply_pending_update. The model-level hooks (iter_moe_layers,
  apply_pending_balance_updates, collect_moe_aux_loss,
  collect_moe_metrics) move here and hooks.py goes.
- balance.py: BalanceController implements TrainingState, whose no-op
  defaults replace its own copies; the timing rules live in DESIGN.md.
- dflash.py: the pending-update comment points at DESIGN.md.
- tests: the components run through one loop, pending updates reach every
  layer, and the DFlash forward applies them in training only.
Checkpoint files name an MoE layer's tensors the way the target family
does, while modules keep a native layout (stacked [E, out, in] experts).
checkpoint.py replaces sgl-project#811's state_dict.py converters with layouts
registered by name: the preset picks one (moe_layout overrides it), only
that layout runs (sgl-project#811's converters all ran on every state dict), and it
maps one MoE layer at a time.

- checkpoint.py: CheckpointLayout(to_file, from_file), the LAYOUTS
  registry with the "native" layout (files keep module names), and
  to/from_checkpoint_state_dict(state, owner), where owner is the draft
  model or its config; dense drafts pass through unchanged.
- init.py goes: WarmStartPlan and select_target_experts have no user in
  this PR.
- config.py: MoEConfig.layout (draft JSON key moe_layout, preset
  default); an unknown layout fails when the config resolves.
- families/__init__.py: where target families register a preset and a
  layout; core modules never import it.
- Every file read and write converts through the draft's layout:
  FSDPTrainingBackend save/load, trainer resume, warm start,
  AutoDraftModel.from_pretrained (MoE files are read memory-mapped and
  regrouped, which HF from_pretrained cannot do), materialize_draft and
  the HF exporter. to_sglang (EAGLE3 only) goes back to main's version.
- __init__.py: the public API is the contracts, the registration
  functions, the layer and its hooks, and the checkpoint boundary; the
  registries and build functions stay in their modules.
- tests: the layout config keys, a per-expert test layout's round trip
  (with layer-prefix disambiguation), owner resolution, the native
  identity, and a from_pretrained round trip through the layout.
sgl-project#811 added the moe/* metrics inside the DFlash and DSpark strategies'
forward_loss and collected the auxiliary loss without adding it to the
objective.

- strategies/base.py: DraftTrainStrategy.finalize_step_output adds the
  draft's summed MoE auxiliary loss to whichever term backpropagates
  (loss, or the loss_terms numerator scaled by this micro-batch's
  denominator, the weighting the objective gets) and the layer-averaged
  moe/* metrics; drafts without MoE layers get their output back
  unchanged. _moe_metrics and its two call sites go.
- controller.py: TrainerCore calls it after forward_loss on every
  training micro-step; evaluation calls forward_loss alone, so eval
  losses exclude the auxiliary term.
- test_dflash_host_syncs.py: the optimizer-boundary test's stand-in
  strategy subclasses DraftTrainStrategy, the type TrainerCore takes.
- tests: the aux loss reaches the gradient with and without loss_terms,
  the moe/* metrics reach a strategy with no MoE code, and a dense
  step's output is untouched.
- moe/DESIGN.md: rewritten for the framework core: the package layout,
  the rules (core never imports families, every choice is explicit, a new
  family only adds files), and the contracts: detection keys on the preset, one
  TrainingState loop, deferred updates, the accumulation order, attribute
  names, naming converted at the file boundary, and model-level terms for
  every strategy. sgl-project#811's notes that still hold stay: the Megatron-Core
  split, why a preset per target family (serving, latency), and the
  extension points, including experts backends and shared experts.
- customization.md: drafts opt in by naming a moe_preset; architecture
  and training-only keys, losses and metrics without strategy code,
  checkpoints and exports through the layout, how new families and
  components register.
- training/DESIGN.md: DraftTrainStrategy owns finalize_step_output, which
  TrainerCore calls after forward_loss.
- balance.py: load_balancing_loss, transformers' per-layer auxiliary
  load-balancing loss (E * sum_e f_e * p_e), and AuxLossController,
  registered as "aux_loss": no selection bias; in training it emits
  aux_loss_coeff times that loss over the micro-batch's tokens and
  reports aux_loss, load_frac_max and load_frac_mean.
- tests: the loss matches transformers' load_balancing_loss_func and the
  controller scales it by aux_loss_coeff.
- config.py: shared_expert_gate, an architecture key: "none" or a
  per-token "sigmoid" (validated).
- shared.py: with "sigmoid", SwiGLUSharedExpert scales its output per
  token by sigmoid(gate(x)), gate a bias-free Linear(hidden, 1).
- tests: the gated output against a reference; the key is validated.
From-scratch DFlash-family drafters feed every router a token-independent
common mode (mask-token embeddings plus near-uniform attention over the
target features), so every token picks the same top-k and an aux loss
alone does not recover.

- config.py: moe_router_bias (architecture key): a fixed fp32 router
  logit bias; moe_router_center ("sample") and moe_router_center_momentum
  (training keys).
- router.py: TopKRouter centers its input on each sample's token mean in
  training and on a running mean (input_mean, moved in
  apply_pending_update outside checkpoint regions and averaged across
  ranks) in eval; bias and input_mean are fp32 buffers that no optimizer
  owns. fold_router_centering turns the running mean into a router logit
  bias (gate.bias = -W @ mu, added to an existing bias) for exports;
  fold_router_centering_config is the matching config change.
- base.py: apply_keeping_fp32 keeps the router's fp32 buffers exact
  through module-wide dtype casts: the folded bias is the common-mode
  logit, large next to the token spread, so bf16 rounding could flip
  top-k choices.
- checkpoint.py: to_serving_state_dict folds then names;
  exported_config_updates gives the config.json changes.
- to_hf.py: exports fold the centering and rewrite config.json so the
  export reloads as the draft it describes.
- tests: per-sample centering, the running mean and its momentum, exact
  buffers under casts and meta materialization, the fold against the
  trained routing, the router bias, and a centered export that reloads.
- checkpoint.py: a layout may name the config.json fields its serving
  engine reads (serving_config) and map a dense draft's serving class to
  the engine's class for the same draft with the family's MoE FFN
  (serving_architectures); serving_config, serving_architecture and
  exported_config_updates expose them.
- normalize_dflash_export.py: MoE exports get the serving class and the
  routing fields; dense exports are untouched.
- tests: the fields and the class through a test layout; the normalizer
  leaves dense exports alone.
- families/qwen.py: the qwen3_5_moe preset (softmax scores, renormalized
  top-k, aux_loss balancing, one sigmoid-gated SwiGLU shared expert), the
  qwen layout (the per-expert gate_proj/up_proj/down_proj naming SGLang's
  Qwen2MoeSparseMoeBlock loader reads) and the serving config for
  Qwen3MoeDSparkModel, which refuses recipes that class cannot run.
- configs/qwen3.8-27b-dspark-moe.json: a DSpark draft for Qwen3.8-27B
  with 512 experts of width 512, top-10, a 2048-wide shared expert,
  router centering and aux loss 0.001.
- tests: the layer against transformers' Qwen3.5-MoE block, the aux loss
  against its load_balancing_loss_func, the layout's names and round
  trip, the export fold and serving config, and the normalizer's class.
- qwen3.8-27b-dspark-moe-h200-disaggregated.yaml: four FP8 capture
  servers plus a 4-rank FULL_SHARD trainer with optimizer offload,
  training the MoE drafter in configs/qwen3.8-27b-dspark-moe.json from
  scratch.
- examples/configs/README.md: the recipe's entry.
@yushengsu-thu
yushengsu-thu force-pushed the moe-support-qwen branch 3 times, most recently from d3131a0 to 1d47604 Compare October 20, 2026 03:27
The four component builders and serving_architecture looked an
implementation up and used it in the same expression; they now bind it
to a name first.

- balance.py, router.py, experts.py, shared.py: build_* binds
  controller_cls, router_cls, experts_cls or shared_cls, then calls it.
- checkpoint.py: serving_architecture binds the layout, then reads its
  serving classes, as serving_config does.
@yushengsu-thu
yushengsu-thu force-pushed the moe-support-qwen branch 2 times, most recently from 1c091fc to 3bda0a9 Compare October 20, 2026 06:58

@yushengsu-thu yushengsu-thu left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

check

Comment thread specforge/modeling/draft/moe/router.py Resolved
Comment thread specforge/modeling/draft/moe/balance.py Resolved
Comment on lines +34 to +40
@abc.abstractmethod
def forward(self, x: torch.Tensor, routing: RoutingResult) -> torch.Tensor:
"""Weighted routed sum of ``x`` (``[T, hidden]``) as a fresh fp32 tensor."""

@abc.abstractmethod
def reset_parameters(self, std: float) -> None:
"""Initialize the expert weights with the draft's ``initializer_range``."""

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I should add. This is very common

@@ -0,0 +1,77 @@
{

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it common to give this configs??
Refter to config and see how to provide the right config

Comment on lines +106 to +141
# -- model-level hooks over any module tree -------------------------------------


def iter_moe_layers(module: nn.Module) -> Iterator[MoELayer]:
for sub in module.modules():
if isinstance(sub, MoELayer):
yield sub


def apply_pending_balance_updates(module: nn.Module) -> None:
"""Apply stashed routing statistics; call first in the training forward."""
for layer in iter_moe_layers(module):
layer.apply_pending_update()


def collect_moe_aux_loss(module: nn.Module) -> Optional[torch.Tensor]:
"""Sum of the layers' already-scaled auxiliary losses, or ``None``."""
total: Optional[torch.Tensor] = None
for layer in iter_moe_layers(module):
loss = layer.aux_loss()
if loss is not None:
total = loss if total is None else total + loss
return total


def collect_moe_metrics(
module: nn.Module, prefix: str = "moe/"
) -> Dict[str, MetricValue]:
"""Layer-averaged scalar metrics under ``prefix``; ``{}`` without MoE layers."""
sums: Dict[str, MetricValue] = {}
n = 0
for layer in iter_moe_layers(module):
n += 1
for key, value in layer.metrics().items():
sums[key] = value if key not in sums else sums[key] + value
return {f"{prefix}{key}": value / n for key, value in sums.items()}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hook functions , is necessary ?

Comment on lines +52 to +55
if self.training:
self.last_counts = routing.counts.detach()
for component in self.components():
component.observe(routing)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • Expert forward output仍能算。
  • 主 loss 仍能訓練 Router。
  • 但 MoE aux loss 完全消失。
  • load-balancing metrics 消失。
  • 自訂 component 的 observe() 不會執行。
  • moe_aux_loss_coeff 形同無效
    這些事情要在哪裡做好?

Comment thread specforge/modeling/draft/moe/shared.py Resolved
The aux loss acts only in backward, yet it came from the "aux_loss"
balance controller, a selection-time policy whose
adjust_selection_scores returns the scores unchanged. The router now
adds the loss itself whenever moe_aux_loss_coeff > 0, and balance only
names a policy that shifts the scores used to pick experts.

- router.py: TopKRouter adds moe_aux_loss_coeff times the selected aux
  loss in each training forward and reports it as the aux_loss metric.
- balance.py: AUX_LOSSES registers load_balancing_loss as "switch". The
  "aux_loss" controller is removed with its load_frac_* metrics:
  load_frac_max is load_max_ratio / E and load_frac_mean is always 1/E.
- config.py: dflash_config.moe_aux_loss_type picks the aux loss
  (default "switch").
- families/qwen.py: the preset sets balance="none"; the serving config
  accepts only "none".
- __init__.py: exports register_aux_loss.
- DESIGN.md, customization.md: the two knobs and the new key.
- tests: the router's aux loss against transformers, an unknown aux loss
  name, and the Qwen preset's balance.

The Qwen recipe's aux loss, layer output and gradients are bitwise
unchanged.
dflash_config.moe_aux_loss_coeff now defaults to 0.001, the
router_aux_loss_coef of transformers' Qwen MoE configs, so an MoE draft
adds the router's load-balancing loss unless its JSON sets the
coefficient to 0.

- config.py: aux_loss_coeff defaults to 0.001.
- DESIGN.md, customization.md: the default and how to turn it off.
- tests: the default adds the loss and 0 turns it off.

The Qwen recipe already sets 0.001, so its training is unchanged.
The routed experts and the shared expert are both expert FFNs, and the
shared expert already used experts.py's swiglu, so one module now holds
both. The classes, registries and shared_experts.* parameter names do
not change.

- experts.py: SharedExpert, SHARED_EXPERTS, register_shared_expert,
  build_shared_expert and SwiGLUSharedExpert, moved unchanged.
- shared.py: removed.
- __init__.py, layer.py, tests: import them from experts.py.
- DESIGN.md: the module list.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants