feat: DSpark trainer (DFlash + Markov/confidence heads + L1 distillation)

[maocheng23]

Summary

Adds DSpark draft-model training to SpecForge — a port of TorchSpec PR #129 ("DeepSpec - DSpark trainer support"). DSpark extends DFlash's block-diffusion drafter with EAGLE-style Markov (low-rank learned bigram bias) and confidence (per-position accept-rate) heads, trained with a combined objective:

loss = ce_alpha·CE + l1_alpha·L1_distill + confidence_alpha·BCE

• CE — hard cross-entropy on the next tokens, on Markov-biased logits.
• L1 distillation — |softmax(draft) − softmax(target)|, where the target distribution is the frozen target LM head applied to the target's final hidden state at the aligned position.
• Confidence BCE — per-position accept-rate (1 − 0.5·L1) prediction, used for adaptive block length at inference.

Markov/confidence formulation adapted from DeepSeek's DeepSpec (MIT).

⚙️ Stacked on #652 (base dataflow-up-16-sgl0.5.14, which bumps up-16 to sglang 0.5.14). Originally on main; the DataFlow refactor moved the DFlash target code out of specforge/modeling/target/ into the specforge/inference/target_engine/ capture-policy architecture, so the last_hidden_states change is now expressed there (see Changed below). OnlineDSparkModel and DSparkDraftModel are unchanged — their base classes (OnlineDFlashModel, DFlashDraftModel) and every imported symbol are identical on this branch. Merge #652 first.

What's added

New files

• specforge/modeling/draft/dspark.py — DSparkConfig (extends Qwen3Config), VanillaMarkov, AcceptRatePredictor, build_markov_head, DSparkDraftModel (subclass of DFlashDraftModel).
• specforge/core/dspark.py — OnlineDSparkModel (subclass of OnlineDFlashModel): reuses DFlash anchor sampling / MASK-noise / block-mask machinery; adds Markov-biased logits, CE + L1 + confidence, within-block decay weighting, and a pooled global-mean loss.
• scripts/train_dspark.py — training driver (mirrors train_dflash.py); plumbs last_hidden_states, logs per-component losses, adds --max-steps.
• configs/qwen3-8b-dspark.json, examples/run_qwen3_8b_dspark_online.sh.
• tests/test_utils/test_dspark.py — 11 CPU unit tests (head math, config, internal loss identity, all-slots label convention, grad flow, CE-only path).

Changed — surface the target's final hidden state as DFlashTargetOutput.last_hidden_states (DFlash ignores it; DSpark's L1/confidence losses need it; CE-only training needs nothing new). On this branch that lives in the new capture-policy layer:

• specforge/inference/target_engine/target_capture_policy.py — add the optional last_hidden_states field to DFlashTargetOutput, populate it in DFlashCapturePolicy.hf_capture (post-norm final hidden) and .sglang_capture, and set the DFlash spec to build the sglang backend with wrap_eagle3_logits=True (so the final hidden can be surfaced).
• specforge/inference/target_engine/sglang_backend/capture.py — extend_dflash gains return_last_hidden_states and routes through the shared _forward_eagle3_reqs (return_logits=False): the aux concat is DFlash's context, last_hidden_states is the target final hidden. This is how sglang exposes the post-norm final hidden (via the eagle3 logits wrapper), and it unifies the dflash + eagle3 sglang capture path.
• specforge/modeling/draft/__init__.py, specforge/core/__init__.py — exports.

Design notes

• This is a port, not a cherry-pick. SpecForge structures DFlash differently from TorchSpec (no DFlashConfig/DFlashModel/DFlashTrainer class family; a monolithic OnlineDFlashModel + a flat training script; a different draft-model forward signature; DFlash drops anchor slot 0 while DSpark supervises all block slots with a cumprod eval mask). DSpark is re-expressed in SpecForge idioms — DSparkConfig extends Qwen3Config (not a DFlashConfig), the loss wrapper subclasses OnlineDFlashModel, and training is a script (no trainer class).
• Pooled global loss — local numerators over an all-reduced denominator, scaled by world_size to cancel FSDP's mean-grad reduction (correct for ZeRO-2 / SHARD_GRAD_OP, a single shard group). Verified by matching loss scale across 1-node vs 2-node runs.

Testing

• CPU unit tests: 11/11 pass on this branch (python -m pytest tests/test_utils/test_dspark.py).
• GPU on sci-h200 (8×H200), against real sglang 0.5.14 (with fix(sglang): port the DataFlow capture backend to sglang 0.5.14 (up-16) #652):
  • tests/test_runtime: 296 passed / 2 skipped / 1 xfailed — the up-16 port regresses nothing.
  • HF last_hidden_states capture: ✅ final hidden surfaced on a real Qwen forward ((2,16,896), width = context/3, finite). The L1/confidence path works on HF.
  • sglang DFlash context capture: ✅ works on 0.5.14 (hidden_states=(2,16,2688)).
  • sglang last_hidden_states: ✅ (2,16,896) — surfaced via the eagle3 wrapper (see Notes). L1/confidence now train on sglang too, not just HF.
  • eagle3 sglang capture: ✅ unaffected (hidden=(2,16,2688), target=(2,16,151936)).
• GPU (2× H200 nodes, Qwen3-8B target, flex_attention) — from the pre-rebase (main-based) revision; OnlineDSparkModel / DSparkDraftModel are byte-identical here, so these results carry over. The HF L1/confidence path is unchanged; the sglang last_hidden_states path was re-expressed against the new capture policy and should be re-smoked on GPU (see follow-ups).
  • 1-node 8-GPU and 2-node 16-GPU smoke: clean, all three loss components healthy, checkpoint saved. Loss scale matches across 1-node vs 2-node (validates the world_size correction).
  • 300-step real ShareGPT run (2-node): loss 3.09→2.49, draft acc 1.5%→16.2%, ce 11.46→6.20, l1 1.97→1.66.
  • Full 2-epoch run on the complete ShareGPT set (120,675 convs, 7,276 steps, 2-node): combined loss 3.09→1.31, draft greedy accuracy 1.5%→~52%, ce_loss→~2.0, l1_loss 1.97→0.67. confidence_loss settles/decreases (0.51→0.45) once the accept-rate target stabilizes. All three components healthy throughout; no NaN/OOM/instability.

Notes / follow-ups

• sglang last_hidden_states now works (GPU-validated on 0.5.14). sglang exposes the post-norm final hidden only through the eagle3 logits-processor wrapper (when DFlash captures aux layers, LogitsProcessorOutput.hidden_states is the aux concat, with no separate final stream). So the DFlash sglang backend is now built with wrap_eagle3_logits=True and extend_dflash routes through the shared _forward_eagle3_reqs with return_logits=False — no vocab logits, hence no eagle3 vocab mapping needed. DSpark's L1/confidence losses train on both HF and sglang backends. This also unifies the dflash and eagle3 sglang capture on one _forward_eagle3_reqs primitive. CE-only (--l1-loss-alpha 0 --no-confidence-head) remains available on either backend.
• Heads use default init (backbone keeps post_init HF init) — faithful to the reference.
• Optional next step: wire a DSparkTrainStrategy into the runtime path for parity with EAGLE3.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces the DSpark draft model and its online training pipeline, which builds upon the DFlash backbone by adding EAGLE-style Markov and confidence heads along with L1 distribution distillation. The changes include new configurations, training scripts, model definitions, target model updates to surface final hidden states, and unit tests. The review feedback highlights two key issues: a bug in the dataset filtering step in train_dspark.py where calling .sum() on a Python list will cause a crash, and an incorrect loss scaling calculation in dspark.py when Tensor Parallelism is enabled, which should use the data-parallel group instead of the global world size.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The Hugging Face Dataset.filter method passes a dictionary of Python types (where sequence features like loss_mask are standard Python list objects) to the filter function when batched=False (the default). Calling .sum() on a Python list will raise an AttributeError: 'list' object has no attribute 'sum' and crash the training script. Use the built-in sum() function instead.

Suggested change

	train_eagle3_dataset = train_eagle3_dataset.filter(
	lambda x: x["loss_mask"].sum() >= min_loss_tokens
	)
	train_eagle3_dataset = train_eagle3_dataset.filter(
	lambda x: sum(x["loss_mask"]) >= min_loss_tokens
	)

[gemini-code-assist]

When Tensor Parallelism (TP) is enabled (e.g., via --tp-size), dist.get_world_size() returns the total number of GPUs across both DP and TP groups. However, the loss denominator should only be summed across the data parallel (DP) group, and the loss scaling factor should be dp_size rather than world_size to ensure correct gradient scaling under FSDP. Using the global world_size will scale gradients incorrectly by a factor of tp_size.

Consider using get_dp_group() from specforge.distributed to perform the all_reduce and scale the loss correctly.

Suggested change

	world_size = dist.get_world_size() if dist.is_initialized() else 1
	global_den = local_den.detach().clone()
	if world_size > 1:
	dist.all_reduce(global_den, op=dist.ReduceOp.SUM)
	global_den = global_den + 1e-6
	loss = (
	self.ce_loss_alpha * ce_num / global_den
	+ self.l1_loss_alpha * l1_num / global_den
	+ self.confidence_head_alpha * conf_num / global_den
	) * world_size
	from specforge.distributed import get_dp_group
	dp_group = get_dp_group() if dist.is_initialized() else None
	dp_size = dist.get_world_size(dp_group) if dp_group is not None else 1
	global_den = local_den.detach().clone()
	if dp_size > 1:
	dist.all_reduce(global_den, op=dist.ReduceOp.SUM, group=dp_group)
	global_den = global_den + 1e-6
	loss = (
	self.ce_loss_alpha * ce_num / global_den
	+ self.l1_loss_alpha * l1_num / global_den
	+ self.confidence_head_alpha * conf_num / global_den
	) * dp_size

[Dogacel]

This sounds right.

[Dogacel]

I think overall LGTM, thank you!

[Dogacel]

What does this mean? Does SGLang not return the final hidden state? It should have.

[Dogacel]

Will this work? Isn't there other local imports that those files depend on?

[Dogacel]

Let's add "ce_loss", "l1_loss", "confidence_loss" ? I think they give us some insight about the training overall.

[Dogacel]

This sounds right.

[Dogacel]

Report indiviudal losses here.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[maocheng23]

🔬 GPU retest on sci-h200 (8×H200), real sglang 0.5.14 (stacked on #652):

check	result
tests/test_runtime	✅ 296 passed / 2 skip / 1 xfail
DSpark CPU unit tests	✅ 11/11 (sglang 0.5.14 + transformers 5.8.1)
HF last_hidden_states capture	✅ final hidden (2,16,896), width = context/3, finite → L1/confidence works on HF
sglang DFlash context capture	✅ (2,16,2688) = 3 layers × 896 → CE-only DSpark trains on sglang
sglang last_hidden_states	⚠️ None on 0.5.14

Root cause for the sglang case: when DFlash captures aux hidden states, sglang 0.5.14's LogitsProcessorOutput.hidden_states is the aux concat — there's no separate final-hidden stream on the raw output (eagle3 only gets it via its logits-processor wrap_eagle3_logits wrapper, which DFlash doesn't use). So DSpark's L1/confidence losses correctly raise on sglang; use --target-model-backend hf (validated) or CE-only (--l1-loss-alpha 0 --no-confidence-head). Follow-up to enable it on sglang: route DFlash capture through the eagle3 return_last_hidden_states wrapper path.

[maocheng23]

🔁 Rebased onto the updated #652 (f59f014) and re-validated on sglang 0.5.14: DSpark CPU tests 11/11, dflash sglang context capture OK ((2,16,2688), CE-only path works). sglang last_hidden_states unchanged (still None — #652's update was eagle3-only; the dflash-on-sglang last_hidden path remains the documented HF-only / follow-up item).

[maocheng23]

✅ Fixed: sglang last_hidden_states now works on 0.5.14 (commit b2ba9aa).

sglang only exposes the post-norm final hidden through the eagle3 logits-processor wrapper, so the DFlash sglang backend is now built with wrap_eagle3_logits=True and extend_dflash routes through the shared _forward_eagle3_reqs with return_logits=False (no vocab logits → no eagle3 vocab mapping). This also unifies the dflash + eagle3 sglang capture path.

GPU-validated on sglang 0.5.14 (sci-h200, 8×H200):

check	before	after
sglang last_hidden_states	⚠️ None	✅ (2,16,896)
sglang DFlash context	✅ (2,16,2688)	✅ (2,16,2688)
eagle3 sglang capture	✅	✅ (unaffected)
tests/test_runtime	296	✅ 296 passed (updated the spec-flags assertion for the new dflash wrap_eagle3_logits)
DSpark CPU tests	11/11	✅ 11/11

DSpark's L1/confidence losses now train on both HF and sglang backends.