Repository navigation
serving: SGLang draft classes for the MoE DSpark/DFlash drafts (stacked on #1003) - #1016
Draft
yushengsu-thu wants to merge 17 commits into
Draft
yushengsu-thu wants to merge 17 commits into
yushengsu-thu wants to merge 17 commits into
Conversation
This was referenced Oct 25, 2026
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.
…Router on top Router declares what the balancing relies on: the forward signature, the fp32 expert_bias buffer (None when a family does not ask for one), the update rate (0 keeps a loaded bias fixed) and update_expert_bias, one DeepSeek-V3 sign step with a zero-mean delta. TopKRouter keeps the top-k algorithm and adds the bias to the choice only. Nothing passes counts to the router yet; the family does not read the knobs yet.
MoELayer keeps a non-persistent int64 tokens_per_expert buffer and, in training, adds this forward's kept slots to it with scatter_add_; the new token_mask argument marks the kept tokens. Integer buffers survive dtype casts and stay exact through accumulation. The router stays stateless. Nothing pools or resets the counts yet.
…, score the step MoEBalancing finds a model's MoELayers by type and, once per optimizer step, all-reduces their kept-slot counts in one collective, hands each router its pooled load through update_expert_bias, resets the counts and reports moe/step_balance (1 when uniform, E / K when collapsed). Nothing registers or calls it yet.
…the step hook once The DFlash-family draft blocks hand block_keep_mask, expanded to tokens, to DFlashDraftModel.forward as moe_token_mask; the decoder layers pass it on only to an MoELayer. TrainerCore gains register_step_hook, run after each optimizer step with its metrics joining that step's; trainer.py registers MoEBalancing once when the model has MoE layers. MoE drafts refuse activation checkpointing again, since the layers count in forward.
…hem on in the DSpark recipe dflash_config.moe_router_enable_expert_bias (Megatron's name) builds the router with the bias; moe_router_bias_update_rate sets its step, 0 keeps a loaded bias fixed. Both default off, so drafts without the keys train as before. The Qwen3-8B MoE DSpark config enables them at 0.001.
…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.
Resolves controller.py against sgl-project#965's window replay: TrainerCore keeps both _prepared_steps and the step hooks.
…e hook only for modules Dense drafts, and the stand-in drafts in tests, keep their forward signature: OnlineDFlashModel passes moe_token_mask only when the draft has an MoELayer. trainer.py builds MoEBalancing only for nn.Module models, so runtime tests with namespace stand-ins are unaffected.
The fp32-cast check joins the router contract test, the stand-in draft goes (the real DFlash and DFlash2 drafts carry the mask), the two TrainerCore tests become one with two subtests, and one test strategy serves with or without loss_terms.
DSpark accepts use_liger_kernel since sgl-project#963. Liger replaces the draft's RMSNorms; the MoE FFN keeps its grouped_mm experts, which Liger does not cover yet. A 10-step offline run of a Liger-built MoE DSpark draft trained with finite loss, bias and counts.
…its MoE layers The DFlash-family model passes block_keep_mask, which it already has, to the draft like the attention mask; DFlashDraftModel.forward expands it by its own block_size into the token_mask its MoE layers count with. The model no longer inspects the draft for MoE layers, and the moe_token_mask name is gone. The stand-in drafts in the loss tests accept the argument.
…sume it The DFlash-family model passes moe_token_mask, the kept-block mask expanded to draft tokens, on every draft call; the model neither inspects the draft for MoE layers nor branches. The draft hands it to its layers under the same name, and only an MoELayer takes it, for its per-step counts and, with a coefficient, its load-balance loss. The stand-in drafts in the loss tests accept the argument.
…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.
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>
yushengsu-thu
force-pushed
the
dev-moe-serving
branch
from
October 29, 2026 00:43
7ff4a7d to
cdd2354
Compare
Draft
3 of 5 tasks
4 of 5 tasks
This branch has not been deployed
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.
Serving for the MoE DSpark/DFlash drafts of #1031 (stacked on it; only the last commit is this PR).
Added
specforge/serving/sglang_models/dflash_moe.py:Qwen3MoeDSparkModel,DFlashMoEDraftModel(SGLang's draft classes, dense MLP replaced byFusedMoEplus the trained router).specforge/serving/__init__.py,specforge/serving/sglang_models/__init__.py: loaded viaSGLANG_EXTERNAL_MODEL_PACKAGE=specforge.serving.sglang_models.scripts/gates/check_moe_sglang_equivalence.py: served FFN vs trainedMoELayer, per layer.Modified
scripts/gates/normalize_dflash_export.py: exports withmoe_presetget the serving class.tests/test_scripts/test_gate_orchestration.py: the normalizer cases.scripts/gates/README.md: serving an MoE export.Versus the previous #1016: rebased from #1003 onto #1031;
gate.selection_bias/moe_router_selection_biasrenamed togate.expert_bias/moe_router_enable_expert_bias.🤖 Generated with Claude Code