Repository navigation
MoE drafts: the load-balance loss (stacked on #1034) - #1035
Open
yushengsu-thu wants to merge 23 commits into
Open
yushengsu-thu wants to merge 23 commits into
yushengsu-thu wants to merge 23 commits into
Conversation
The MoE layer, experts, top-k router, the qwen3_5_moe family and the build_ffn switch, with the Qwen3-8B MoE DSpark config and recipe. No load balancing: the router is stateless and the trainer is untouched.
Router fixes what every family's router provides, forward(x) -> scores, indices, weights, and is the point where other families (DeepSeek) add their own routing; TopKRouter keeps the algorithm. Load balancing builds on this contract in a later PR.
This was referenced Oct 28, 2026
yushengsu-thu
marked this pull request as ready for review
October 28, 2026 06:18
yushengsu-thu
requested review from
FlamingoPg,
FrankLeeeee,
shuaills and
sleepcoo
as code owners
October 28, 2026 06:18
yushengsu-thu
force-pushed
the
dev-moe-aux
branch
4 times, most recently
from
October 29, 2026 22:28
a09dca4 to
9178715
Compare
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.
SharedExpert.reset_parameters, and MoELayer.reset_parameters covers the shared expert and its gate too, so resetting the layer means the whole layer. HF's init already reached these Linears; the values are the same.
Within an MoE draft (moe_preset), ffn_layer_types marks the layers that keep the dense MLP (intermediate_size); the others take the family's MoE layer (moe_intermediate_size per expert). No list means every layer is MoE, as before. The DSpark MoE config trains ["dense", "moe"].
…he score Router implements the shared steps, top-k selection and the combine weights (renormalized, scaled), and forward always returns scores, indices, weights. A family overrides the steps it needs: TopKRouter adds the linear gate and its softmax score; a DeepSeek-style router would override score (fp32 sigmoid) and select (group-limited). MoELayer takes any Router. TopKRouter's numerics are unchanged.
The step that gathers, renormalizes and scales the chosen experts' scores returns the weights that combine their outputs; routing_weights says so (combine did not).
yushengsu-thu
force-pushed
the
dev-moe-aux
branch
from
October 29, 2026 22:36
9178715 to
d5d204e
Compare
Follows the draft-layer layout of sgl-project#1025: the MoE FFN is a common FFN component next to the dense one, in specforge/modeling/draft/layers/ffn/moe. A pure move plus import paths; layers/__init__.py and layers/ffn/__init__.py stay empty so sgl-project#1025's versions merge in cleanly.
yushengsu-thu
force-pushed
the
dev-moe-aux
branch
from
October 29, 2026 23:15
d5d204e to
0133a8e
Compare
yushengsu-thu
force-pushed
the
dev-moe-aux
branch
from
October 30, 2026 02:42
0133a8e to
85149ed
Compare
3 of 4 tasks
Sigmoid top-k sums can underflow to 0 and turn the weights into NaN; HF DeepSeek-V3/V4 and torchtitan's route_norm_epsilon guard it with 1e-20. For softmax the sum is at least K / E, so the addition is a bitwise no-op in fp32: Qwen parity and the e2e losses are unchanged.
yushengsu-thu
force-pushed
the
dev-moe-aux
branch
from
October 30, 2026 21:33
85149ed to
888cc60
Compare
The shipped draft config now matches the 8 x 1536 experiment: 5 layers, 32 query / 8 KV heads, every FFN an MoE layer (8 experts, top-1), and the recipe trains for 6 epochs with Liger on. ffn_layer_types stays explicit so the per-layer choice is visible; the dense + MoE mix is still covered by the tests.
The DFlash-family model passes moe_token_mask, the kept-block mask expanded to draft tokens, on every draft call; the draft and its decoder layers pass it on, and only an MoELayer takes it. Nothing reads it yet: the load balancing that follows counts with it. The stand-in drafts in the loss tests accept the argument.
TrainerCore.register_step_hook takes a StepHook: before_backward adjusts each micro-batch's StepOutput before the backend sees it, on_step runs after every optimizer step and its metrics join that step's. Nothing registers one yet; MoE load balancing will, once, at assembly.
A fp32 expert_bias buffer (None unless a family asks for one) shifts the top-k picks, never the weights; bias_update_rate 0 keeps a loaded bias fixed; update_expert_bias is one DeepSeek-V3 sign step with a zero-mean delta, called with the step's pooled kept-slot counts.
A non-persistent int64 tokens_per_expert buffer; in training each forward adds its kept slots (moe_token_mask) with scatter_add_. Integer buffers survive dtype casts and stay exact; reset_parameters zeroes the counts since from_pretrained builds on the meta device.
MoEBalancing finds a model's MoELayers by type and, at every optimizer step, all-reduces their 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). trainer.py registers it when the model has MoE layers. MoE drafts refuse activation checkpointing: the layers count in forward, a recompute would count twice.
…ecipe turns them on dflash_config.moe_router_enable_expert_bias and moe_router_bias_update_rate (0 keeps a loaded bias fixed), both off by default.
Router contract and bias step, the layer's masked counts, the step hook with a mocked all-reduce, the mask through DFlash and DFlash2 layers, TrainerCore running the registered hook, a 2-rank pooled update, the launch gate's bias assertions, the host-sync MoE case with the bias, and a from_pretrained round trip of the bias with zeroed counts.
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.
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.
… 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. Folded rather than injected: SpecForge's trainer normalizes gradients by the global denominator, so an injected gradient would need a per-micro-batch scale handed to the layer before backward.
…pe trains with batch_wise 0.005 dflash_config.moe_router_aux_loss_coef and moe_router_aux_loss_type; the default coefficient 0 keeps drafts without the keys loss-free.
The loss against transformers' load_balancing_loss_func with and without the mask, the layer's two load types and its skip at coefficient 0 (no collective), the hook's fold with and without loss_terms, the gradient parity of the hook's fold against an objective-side fold through TrainerCore, a 2-rank batch-wise load, the launch gate's aux metric, the host-sync MoE case with the loss on, and the family's aux knobs.
yushengsu-thu
force-pushed
the
dev-moe-aux
branch
from
October 30, 2026 22:20
888cc60 to
ec8b69c
Compare
This was referenced Oct 31, 2026
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.
Summary
The load-balance loss for the MoE drafts, stacked on #1034; increment: yushengsu-thu/SpecForge@dev-moe-bias...dev-moe-aux
MoELayercomputes the loss (E * sum_e share_e * P_e, 1 when balanced; torchtitan's units) over themicrobatch_wiseorbatch_wiseload (the step's pooled counts; default). At coefficient 0 nothing runs.MoEBalancing.before_backwardfoldssum coef_l * aux_linto the loss and theloss_termsnumerator, logsmoe/aux_loss. Folded, not injected: the trainer normalizes by the global denominator.moe_router_aux_loss_coef,moe_router_aux_loss_type; the recipe trainsbatch_wise0.005.router.pyis unchanged. The loss is a statistic over the routed tokens, like the counts, so it lives onMoELayer(token mask, per-layer coefficient and type) with the math inbalance.py. It reads only the router'sscoresandindices, so everyRoutersubclass gets it; the per-token L1 normalization ofscoresis a no-op for softmax and lets sigmoid scores (DeepSeek) work too.Formula and units are #1003's; bitwise equal on the MoE path, within noise over 500 steps. Supersedes #1017.