Conversation
… DFlash Implement DeepSeek DeepSpec's DSpark draft-training objective as a "dspark" loss_objective on the existing DFlash drafter. DSpark and DFlash are the same block-parallel anchor drafter: DeepSpec realizes DSpark as that drafter [1] plus a low-rank Markov head [2], a per-position confidence head, and a distribution-matching loss [3]. TorchSpec already ships the drafter as DFlash, so DSpark reuses that code path (dual-source KV, block-causal FlexAttention, anchor/noise machinery) and is added as a loss_objective with two optional heads, rather than a separate model/trainer stack. A separate-stack port of DSpark also exists (lightseekorg#129); both are faithful to DeepSpec, and this one keeps a single code path for the shared DFlash backbone. [1] https://github.com/deepseek-ai/DeepSpec/blob/dd854392fadf053cfddcbc4dc0e6e32de46d1bd0/deepspec/modeling/dspark/qwen3/modeling.py [2] https://github.com/deepseek-ai/DeepSpec/blob/dd854392fadf053cfddcbc4dc0e6e32de46d1bd0/deepspec/modeling/dspark/markov_head.py [3] https://github.com/deepseek-ai/DeepSpec/blob/dd854392fadf053cfddcbc4dc0e6e32de46d1bd0/deepspec/modeling/dspark/loss.py dspark-scoped (the DFlash decay/dpace paths are unchanged): - next-token labels (slot j predicts anchor+j+1), all block_size slots supervised - contiguous-prefix eval mask (cumprod) - anchor-seeded Markov previous token - anchor sampling that requires both the anchor and its first target to be valid - global token-pooled loss normalization (cross-rank all-reduce x world_size) - low-rank vanilla Markov head + confidence head (optional Markov-feature fusion) - CE + TV/L1 distribution distillation (against the target LM head applied to its final hidden state) + confidence BCE Validated component-wise against the DeepSpec reference loss/forward and on the GPU FlexAttention path. Qwen3-8B recipe: block_size=7, gamma=4.0, markov_rank=256. Signed-off-by: Aaron Batilo <AaronBatilo@gmail.com>
e8a18fe to
0db7eac
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0db7eac045
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| "but none was found. Enable store_last_hidden_states on the inference " | ||
| "engine (captured by default for HF/vLLM; set it explicitly for SGLang)." | ||
| ) | ||
| target_last_hidden_states = tlhs.to(device, non_blocking=True) |
There was a problem hiding this comment.
When loss_objective='dspark' is run with inference_engine_type='vllm', the stored last_hidden_states are pre-final-norm by default (last_hidden_states_prenorm defaults true for vLLM), but this path forwards them directly and the DSpark loss projects them with lm_head_weight without loading/applying the verifier norm. That makes the L1 distillation and confidence targets come from the wrong target distribution for vLLM DSpark runs; mirror the Eagle3 path by loading the norm and applying it before passing these states onward.
Useful? React with 👍 / 👎.
| # DeepSpec next-token convention: slot j predicts anchor+j+1, all slots | ||
| # supervised; eval_mask = contiguous supervised prefix per block (cumprod). | ||
| label_offsets = torch.arange(1, self.block_size + 1, device=device).view(1, 1, -1) |
There was a problem hiding this comment.
This DSpark branch makes slot 0 the first next-token prediction, but DFlashTrainer._aggregate_metrics and _aggregate_eval_metrics still unconditionally drop index 0 as an anchor slot (pred_* = ...[1:]). In any DSpark run, the reported/eval average loss, accuracy, simulated accepted length, and acc_0/ploss_0 are therefore shifted and omit the first supervised token; keep slot 0 for DSpark or return a layout that the trainer does not slice.
Useful? React with 👍 / 👎.
|
@Dogacel that's right. I ran some small sample training jobs using your branch and this one and the results are almost exactly equivalent on the same data, etc. Numerically they AL that I got was <0.05 away from each other. I still want/need to compare baseline dflash and dspark though. This PR is kind of a check to see which design philosophy the upstream maintainers want in terms of treating dspark differently or treating it the same/similar as dflash |
Thanks for the contribution. We have choosen to go with reusing DFlash components but building DSpark on top of it as a separate module. Unlike DeepSpec, I think trying to add DSpark on top of the existing DFlash code doesn't feel natural here, at least for now. It brings many unused variables & codepaths for the DFlash, which might be confusing and harder to maintain. |
|
Closing! |
2 participants
Summary
Implement DeepSeek DeepSpec's DSpark draft-training objective as a
dsparkloss_objectiveon the existing DFlash drafter. DSpark and DFlash are the sameblock-parallel anchor drafter: DeepSpec realizes DSpark as that drafter [1] plus
a low-rank Markov head [2], a per-position confidence head, and a
distribution-matching loss [3]. TorchSpec already ships the drafter as DFlash, so
DSpark reuses that code path (dual-source KV, block-causal FlexAttention,
anchor/noise machinery) and is added as a
loss_objectivewith two optionalheads, rather than a separate model/trainer stack. A separate-stack port of
DSpark also exists (#129); both are faithful to DeepSpec, and this one keeps a
single code path for the shared DFlash backbone.
[1] https://github.com/deepseek-ai/DeepSpec/blob/dd854392fadf053cfddcbc4dc0e6e32de46d1bd0/deepspec/modeling/dspark/qwen3/modeling.py
[2] https://github.com/deepseek-ai/DeepSpec/blob/dd854392fadf053cfddcbc4dc0e6e32de46d1bd0/deepspec/modeling/dspark/markov_head.py
[3] https://github.com/deepseek-ai/DeepSpec/blob/dd854392fadf053cfddcbc4dc0e6e32de46d1bd0/deepspec/modeling/dspark/loss.py
Scope (dspark-only; the DFlash decay/dpace paths are unchanged)
block_sizeslots supervisedTesting
pytest tests/test_dspark.py— 8 CPU tests pass: forward shape/finiteness, loss decomposition, all-masked → 0, next-token supervises all block slots, grad flow + frozen embedding, Markov bias == W2(W1[prev]), linear confidence head, and numerical faithfulness vs the vendored DeepSpec reference loss.ruff checkandruff format --checkclean on all changed files.block_size=7, gamma=4.0, markov_rank=256.