Repository navigation
MoE drafts: DeepSeek-V4.1's own DSpark draft (stacked on #1035) - #1073
Open
yushengsu-thu wants to merge 47 commits into
Open
yushengsu-thu wants to merge 47 commits into
yushengsu-thu wants to merge 47 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.
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).
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.
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.
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.
Pure relocation, no behavior change. Split from sgl-project#1023. - specforge/layers/{embedding,linear,lm_head,ring} -> specforge/modeling/draft/layers/ - modeling/draft/{flex_attention,flex_attention_backend}.py -> layers/attention/backends/ - generate_eagle3_mask -> modeling/draft/eagle3_mask.py - modeling/draft/llama3_eagle.py -> modeling/draft/eagle3.py - LlamaMLP -> layers/ffn/dense.py:GatedMLP (re-exported from eagle3 as LlamaMLP) - tests/test_layers -> tests/test_modeling/test_layers (test_decoder.py -> test_ring_attention.py) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
specforge.layers, modeling/draft/{llama3_eagle,flex_attention,flex_attention_backend}
now resolve to the relocated modules (same module objects, so registries and
class identity are preserved) and emit a DeprecationWarning. This commit can
be dropped if a clean break is preferred.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Split from sgl-project#1023. DFlash-family attention moves into specforge/modeling/draft/layers: - attention/base.py: Attention/AttentionContext/AttentionOutput contract and DenseAttention (projections + cache update + backend dispatch + o_proj) - attention/{gqa,mla,rotary}.py: GQA/MHA and expanded-MLA projections, RoPE - attention/backends/dense.py: eager/SDPA/Flash/Flex dispatch, unchanged from Qwen3DFlashAttentionBase.forward - config.py (LayerKernels, AttentionConfig), factory.py (build_attention, build_ffn), decoder.py (pre-norm DecoderLayer) Qwen3DFlashAttention/Qwen3DFlashMLAAttention remain as compatibility constructors; DFlashKernels aliases LayerKernels. Mask handling is unchanged. The parity test embeds the pre-refactor classes verbatim and checks outputs, gradients and two AdamW steps bit for bit across DFlash/DFlash2/DSpark/Domino, GQA/MLA, eager/SDPA, sliding and causal layouts, plus cached block decoding. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Split from sgl-project#1023. DecoderLayer (and DFlash2's convolution decoder) read and write the residual stream through Residual modules: - residual/standard.py: the existing pre-norm add; parameter-free, so default checkpoints, parameter order and numerics are unchanged (parity test still bit-exact) - residual/gr.py: Qwen3.8-Flash-Next gated residual (per-branch Gemma RMSNorm, low-rank gated read, gated write, learned final readout) - residual/mhc.py: DeepSeek-V4.1 single-pass mHC (lagged pre-mix, FP32 coefficients and Sinkhorn, stream-zero initial read) Opt-in via draft_residual in the DFlash-family config. Non-standard residuals force non-reentrant gradient checkpointing and record their non-parameter semantics in the DFlash/Domino/DSpark resume contracts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Split from sgl-project#1023. Dense attention now enforces AttentionConfig's causality and sliding window itself, intersected with the algorithm's mask, using logical query/key positions (AttentionContext.query_positions/key_positions) for packed layouts. Queries with no visible key produce zero output on every backend, including after the o_proj bias (previously eager only); individual empty heads are zeroed before projection. Two changes relative to sgl-project#1023, to keep the default training path as fast as before: - AttentionContext.mask_encodes_visibility: DFlash's per-layer-type masks already encode the same rules, so dense backends no longer rebuild an identical mask (an extra create_block_mask under Flex) in every sliding or causal layer. - Flex requests the LSE only when o_proj has a bias; it already returns zeros for empty rows and heads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The five commits before this one carry sgl-project#1025 to sgl-project#1028 (draft layers moved, GQA/MLA attention components with the layer factory, residual kinds, visibility). This commit reconnects the MoE FFN to them: - layers/factory.py: build_ffn takes the layer index and delegates to the MoE package's switch, so moe_preset and ffn_layer_types work again. - layers/decoder.py: DecoderLayer.forward takes moe_token_mask and hands it to an MoELayer as its second argument; dense MLPs keep one argument. - dflash.py: pass the layer index and the mask; DFlash2 already calls _ffn. The one test that used the old build_ffn signature is updated. Temporary until sgl-project#1025 to sgl-project#1028 land on main; then this stack rebases onto them and only this commit is replayed.
A module that defines a per-head ``sinks`` parameter (HF's name, as in gpt-oss and DeepSeek-V4) now has it applied by the dense backend: - eager: transformers' DeepSeek-V4 eager_attention_forward, which puts the sink column into the softmax. The CPU reference path. - flex_attention on CUDA: request the logsumexp and scale the output by sigmoid(lse - sink), the same softmax over [keys, sink]; a query with no visible key keeps its zeros. The fast GPU path. - sdpa, flash_attention_* and CPU flex refuse sinks with a clear error. Tests compare the forward output and every gradient against a softmax reference with the extra column on both backends, and train the sinks alone back to hidden target values. First commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
``dflash_config.context_window`` limits every draft block to the last that many context keys before its anchor, the same keys for each query of the block, as DeepSeek-V4's DSpark does (get_dspark_topk_idxs: 128 cached context positions plus the whole block, for every block query). In-block visibility is unchanged and still follows ``is_causal``; the existing ``sliding_window`` keeps its per-query, FA2-style meaning, and the two intersect when both are set. The window reaches the two mask builders, the training masks of every DFlash-family layer type, the masks spec_generate builds itself (a full layer with a window now gets a mask instead of None) and the attention resume contract, which records the window only when it is set so that contracts of existing runs keep their shape. Existing configs are unaffected: the option defaults to null. Tests check the visible keys per query offset for both in-block directions, clipping at the sequence start, agreement of the dense and flex builders with the element-level reference across batches, blocks and a combined sliding window, the training-path plumbing, the config validation and the resume contract. Third commit of the DeepSeek-V4 DSpark draft (sgl-project#1057).
DeepSeek-V4.1's DSpark scores confidence on the readout after the final
mHC collapse and before the final norm (forward_head: ``logits =
head(norm(x))``, ``confidence = confidence_head(x, markov_embed)``).
SpecForge fed the normed hidden, so a V4.1 draft trained here would have
served a miscalibrated confidence.
``dflash_config.confidence_head_input`` ("post_norm", the default, or
"pre_norm") selects the input. With "pre_norm", the DSpark objective
asks the draft for the readout before the norm; the family model carries
it as one more aligned tensor through the chunked, checkpointed
objective and hands it to predict_confidence. Drafts that do not ask
keep the three-value draft-block output and never see the new keyword.
The confidence projection now computes in fp32 with autocast disabled,
as DeepSeek's does; the features may be an unnormalized readout.
Tests check that the requested readout norms to the regular output, that
the family model passes the pre-norm chunk matching each logits chunk
(and nothing extra in the default mode), the validation errors, and the
fp32 projection under bf16 autocast.
Unrelated to DeepSeek-V4: purely a speed change for any flex_attention draft with head_dim > 256 on Hopper and Blackwell. torch 2.13 picks 16x16 backward tiles past head_dim 256; 32x32 tiles with 8 warps and 3 stages are 3x faster on H200 and B200. On SM100 the default forward tiles fault (misaligned address) and an aliased key/value tensor makes the compiled backward 4.7x slower, so flex_kernel_options also sets 32x32 forward tiles and clones the value there. Caller kernel_options win.
Unrelated to DeepSeek-V4: purely a memory saving for FULL_SHARD with gradient accumulation. FSDP1 no_sync keeps the full unsharded gradient on every rank until the boundary micro-step (37 GiB for a 19.7B bf16 draft against 9 GiB sharded), which defeats FULL_SHARD and OOMs a 4-rank B200 run. Reduce-scattering every micro-step accumulates into the sharded gradient instead; SHARD_GRAD_OP and DDP are unchanged.
yushengsu-thu
force-pushed
the
v4.1-dspark
branch
from
November 4, 2026 01:11
52c096e to
aebfed7
Compare
``dflash_config.confidence_head_bias`` (default true, the existing behaviour) builds AcceptRatePredictor's projection without a bias when false. DeepSeek-V4.1's DSpark ships ``confidence_head.proj.weight`` and no bias, so its draft config will set the flag to false. Taken from the DeepSeek-V4 draft config commit of sgl-project#1057 without that config and its test; the fp32 confidence test already exercised the keyword.
build_moe assembles DeepSeek-V4's MoE layer from its config keys on the stack's parts: TopKRouter with sqrt(softplus) scores, routed_scaling_factor as route_scale and the expert bias, clamped SwiGLU (swiglu_limit) for the routed experts and the one ungated shared expert. As in transformers, the top-k weights are always renormalized. No new router class.
moe_preset: "deepseek_v4" now reaches the family through build_ffn.
The family's block against transformers' DeepseekV4SparseMoeBlock on the grouped_mm and loop paths (same weights, a non-zero expert bias; only two state-dict names differ), the CUDA bf16 case for both families, and the config reading: sqrtsoftplus scores, route_scale, the shared expert, the always-on renormalization, the clamp, the refusals, the preset.
DeepSeek-V4.1 keeps DeepSeek-V4's MoE layer (sqrtsoftplus router, the expert bias steering only the selection, clamped SwiGLU, one ungated shared expert); only the shapes differ, and those come from the draft config. ``moe_preset: "deepseek_v41"`` therefore maps to the deepseek_v4 family, so a V4.1 draft config names its target without a second copy of the code. The family test checks the alias through build_ffn.
A port of the Gate, Expert and MoE classes of DeepSeek-V4.1-Flash's inference/model.py (text tokens: gate_temp 1.0, no image bias, +1e-20 in the renormalization) runs on the family's own weights, with a non-zero expert bias; the forward output and every gradient match within rtol 1e-4 / atol 1e-6. This is why the V4.1 draft uses this family.
A new attention kind, ``deepseek_v41``, following DeepSeek-V4.1's inference DSparkAttention: q_a_proj -> q_a_norm -> q_b_proj; one 512-wide KV head shared by every query head, serving as both K and V (kv_proj -> kv_norm); interleaved RoPE on the trailing qk_rope_head_dim channels with the plain rope_theta (the draft layers have no YaRN); per-head sinks; and a grouped low-rank output, o_a_proj (block-diagonal, o_groups x o_lora_rank) then o_b_proj. Unlike DeepSeek-V4, V4.1 applies no parameter-free per-head RMS norm to q after q_b_proj. Because K is V, the rope slice of the attention output carries the keys' rotation; the conjugate rotation at the query positions turns it into the reference's relative rotation before the output projection. The class owns its forward for that step; DenseAttention and the dense backends are unchanged. The kind is registered in AttentionConfig/build_attention and accepted as dflash_config.attention_mode; the model-level rotary carries the partial-RoPE width as it does for MLA. Tests compare the forward output and every gradient against a plain-PyTorch port of V4.1's inference/model.py DSparkAttention on eager (CPU) and flex (CUDA), check the parameter names and the V4.1-Flash shapes against the checkpoint (hidden 5120, q rank 1280), and build DFlash and DSpark drafts with the kind. Adapted from the DeepSeek-V4 attention kind of sgl-project#1057.
configs/deepseek-v4.1-flash-moe-dspark.json is DeepSeek's own DSpark draft architecture for DeepSeek-V4.1-Flash, field for field from the official config and checkpoint: three stages of V4.1 attention (64 heads on one 512-wide KV head, q rank 1280, grouped output 8 x 1024, rope 64 with the plain theta, sinks, no per-head q norm), the block-anchored 128-token context window with a bidirectional block of 5, single-pass mHC (4 streams, 20 Sinkhorn iterations), the draft's own 128-expert MoE (2304 wide, top-3, one shared expert, sqrtsoftplus routing scaled by 1.5, clamped SwiGLU, expert bias), rms_norm_eps 1e-20, target layers 37-39 of 40, mask token 128799, a rank-256 Markov head and the bias-free confidence head on the pre-norm readout plus Markov embedding. The expert bias is kept fixed and the balance loss off, as the warm start from the official weights needs. The test builds the full-size model on the meta device and checks every state-dict name and shape against the official mtp.* layout, and that the element count equals the official tensors' (fp4 experts unpacked, without the three image-token routing biases a text draft does not carry); a small-scale run trains every parameter through the DSpark objective. The DSpark config guard admits attention_mode deepseek_v41 with one KV head.
specforge/export/deepseek_v41_dspark.py converts DeepSeek-V4.1's DSpark
draft between its official checkpoint and the SpecForge draft of
configs/deepseek-v4.1-flash-moe-dspark.json. The official layout keeps
the stages under mtp.{stage}.* with the inference code's names: FP8
linears (e4m3 with one e8m0 scale per 32x32 block; V4 used 128x128),
FP4 routed experts (two e2m1 codes per byte along the input, low nibble
first, one e8m0 scale per 32 inputs) and fp32 hyper-connection, sink and
gate-bias tensors.
Import dequantizes to bf16, renames to transformers' names (the Markov
head's embed/head become markov_w1/markov_w2), fuses the experts' w1/w3
into gate_up_proj and w2 into down_proj, drops the image-token routing
bias gate.bias_vl that a text draft does not carry, checks every tensor
name and shape against the draft built on the meta device and writes
Hub-style shards; only the shards holding mtp.* are downloaded. Export
reverses the mapping, unquantized, writing gate.bias_vl as a copy of
gate.bias so the official layout stays loadable, with the official
config minus its quantization fields.
scripts/convert_deepseek_v41_dspark.py is the CLI.
Verified on deepseek-ai/DeepSeek-V4.1-Flash: 2401 tensors from three
shards become the draft's 75 (14.225B parameters) that AutoDraftModel
loads without missing or unexpected keys and runs; export then import
reproduces them bit for bit. The official draft's routing biases are all
zero. Tests cover the block and nibble conventions, a round trip, exact
import of synthetic quantized tensors, both directory conversions and
the config.
Adapted from the DeepSeek-V4 converter of sgl-project#1057.
DeepSeek-V4.1-Flash ships a Python encoder instead of a Jinja template, like V4, but its prompt format differs: spaced DSML tag names, a numeric reasoning-effort prefix rendered behind <｜System｜> in thinking mode, and mid-conversation system messages. Register `deepseek-v4.1` and `deepseek-v4.1-thinking` with a `DeepSeekV41Parser` that reuses the V4 parser's message handling and swaps in the V4.1 encoder, vendored verbatim from deepseek-ai/DeepSeek-V4.1-Flash at revision dba1be0a. Tests pin the vendored file to that revision by hash, replay the five encoding cases shipped with the model, cover the V4.1 format changes and the loss-mask spans, and, when an SGLang with the V4.1 encoder is importable, check that training renders exactly what the server renders.
…system message, keep tool-call namespaces
examples/configs/online/disaggregated/external/deepseek-v4.1-flash-moe-dspark.yaml fine-tunes DeepSeek-V4.1's own DSpark draft, warm-started from the official mtp.* weights, for DeepSeek-V4.1-Flash on thinking-mode rows the target regenerated (deepseek-v4.1-thinking): one TP4/EP4 capture server on the spec-capture plugin and a two-rank FULL_SHARD trainer with the optimizer offloaded, global batch 128. Expert parallelism keeps the 2304-wide experts whole (the fp4 MoE kernels reject the 288-column TP split), and mem-fraction 0.80 with four running requests leaves room for the MoE kernels' workspace on a 16k-token prefill. docs/recipes/deepseek-v4.1-flash-moe-dspark-disaggregated.md is the runbook: the SGLang revision the plugin needs (V4.1 landed in 0.5.21, the extension points after 0.5.23), the warm-start conversion, data regeneration, the capture-server and trainer commands, the export back to the official layout, and the Mooncake wheel-version requirement. Validated on one 6x B200 node: 30 optimizer steps from the official weights with no capture failures, the checkpoint exported to HF and converted back to the official mtp.* layout.
… lint The sgl-project#887 conflict resolution left two GLMParser classes in parse.py: main's GLMParser(ThinkingParser) header with only its docstring, then an old GLMParser(GeneralParser) body that shadowed it, so GLM reasoning spans were no longer split (11 failures in tests/test_data/test_glm_parser.py, 27/27 on the base). Keep main's class and drop the dangling header. Also apply the isort/black fixes the lint job asked for in three V4.1 test files.
yushengsu-thu
marked this pull request as ready for review
November 5, 2026 23:44
yushengsu-thu
requested review from
FlamingoPg,
FrankLeeeee,
shuaills,
sleepcoo and
zyksir
as code owners
November 5, 2026 23:44
4 of 5 tasks
Contributor
Author
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 [cherry-pick some PRs - after those PRs merge, rebase this PR first]
DeepSeek's own DeepSeek-V4.1 DSpark draft architecture in SpecForge, trained against
deepseek-ai/DeepSeek-V4.1-Flash; stacked on #1035; increment: yushengsu-thu/SpecForge@dev-moe-aux...v4.1-dsparkWork in progress. So far this branch carries #1025 to #1028 (draft layers moved, GQA/MLA attention and the layer factory, residual hooks with DeepSeek-V4.1 single-pass mHC, attention visibility), cherry-picked from their own commits;
dflash.pyandmodeling/draft/__init__.pyconflicts with the MoE stack are resolved as in #1057.Next: the remaining shared commits from #1057, the
deepseek_v41MoE family, and the V4.1 DSpark draft (attention, config, checkpoint converter).Stacked: #1033 -> #1034 -> #1035 -> this PR.