Repository navigation
MoE support for Qwen drafts (A1, A2, B) - #993
Closed
yushengsu-thu wants to merge 22 commits into
Closed
yushengsu-thu wants to merge 22 commits into
yushengsu-thu wants to merge 22 commits into
Conversation
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.
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
force-pushed
the
moe-support-qwen
branch
3 times, most recently
from
October 20, 2026 03:27
d3131a0 to
1d47604
Compare
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
force-pushed
the
moe-support-qwen
branch
2 times, most recently
from
October 20, 2026 06:58
1c091fc to
3bda0a9
Compare
yushengsu-thu
commented
Oct 20, 2026
yushengsu-thu
commented
Oct 20, 2026
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``.""" |
Contributor
Author
There was a problem hiding this comment.
I should add. This is very common
yushengsu-thu
commented
Oct 20, 2026
| @@ -0,0 +1,77 @@ | |||
| { | |||
Contributor
Author
There was a problem hiding this comment.
Is it common to give this configs??
Refter to config and see how to provide the right config
yushengsu-thu
commented
Oct 20, 2026
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()} |
Contributor
Author
There was a problem hiding this comment.
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) |
Contributor
Author
There was a problem hiding this comment.
- Expert forward output仍能算。
- 主 loss 仍能訓練 Router。
- 但 MoE aux loss 完全消失。
- load-balancing metrics 消失。
- 自訂 component 的 observe() 不會執行。
- moe_aux_loss_coeff 形同無效
這些事情要在哪裡做好?
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Sparse-MoE FFN support for SpecForge draft models, for the Qwen family:
moe_presetin its draft JSON: MoE FFN skeleton for DFlash-family drafters (contracts, config, seams) #811's contracts and seams, made to run;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_tcbalancing, 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,GroupedExpertsandSwiGLUSharedExpertas the built-in implementations. Docstrings and comments are one line, with the explanations inmoe/DESIGN.md.Split plan
[E, out, in]weights,sorted_loopandgrouped_mmdispatch) and a SwiGLU shared expert as the built-in implementations;MoELayerwith oneTrainingStateloop over its components,build_ffn, and the model-level hooks;finalize_step_output, which adds the aux loss andmoe/*metrics for every strategy;aux_lossbalancing;qwen3_5_moepreset;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_lossbalance controller into the router.TopKRouteraddsmoe_aux_loss_coefftimes the loss thatdflash_config.moe_aux_loss_typeselects (only"switch", transformers' loss, for now), andbalanceonly names a policy that shifts the scores used to pick experts. The Qwen preset setsbalance="none"; its aux loss, layer output and gradients are bitwise unchanged. Themoe/load_frac_*metrics are dropped, sincemoe/load_max_ratiocarries the same information. When the stack is split, commit 20 goes with A2, which addedaux_lossbalancing, 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_coeffdefaults to 0.001, therouter_aux_loss_coefof transformers' Qwen MoE configs, and0turns 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.pyintomoe/experts.py, next to the routed experts whoseswigluit already used, and removesshared.py. The classes, registries andshared_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 editsexperts.py.Verification
unittest discoverNotes
Qwen3MoeDSparkModel, which is not upstream yet.