fix(speculative): DSpark and Domino AR validation and label alignment

[h-guo18]

What does this PR do?

Type of change: Bug fix

Stacked on #2757 (its branch is the base). DSpark and Domino train block position k on the token at anchor + k + 1, and DFlash, DFlash2 and LiLiCorr train it on anchor + k. #2757 made the class constant _label_shift drive the training targets. The same fact was still written in four more places:

• DSparkModule.shift_label = True
• DominoModule.shift_label, read from the config
• a NotImplementedError in HFDominoModel.modify
• the DSpark and Domino exporters' getattr(draft_config, "shift_label", True), which re-read the user's raw architecture config

Those copies, and the AR-validation decode, caused four bugs:

1. Domino's AR validation drafted one position off, without its head. HFDominoModel inherited DFlash's pseudo_speculative_generate, which decodes block positions 1..bs-1 with the backbone alone.
   • For Domino, those positions predict the token one step further on. Every draft was therefore checked against the wrong target, so drafts were almost never accepted.
   • The GRU correction head, which training optimizes and the exporter ships, was never applied.
   • domino.yaml keeps AR validation off, so shipped training was not affected.
2. DSpark's AR validation crashed on a sharded base with tied embeddings. scripts/ar_validate.py loads with device_map="auto". accelerate keeps a tied lm_head with embed_tokens on the first GPU, while the draft can land on a later one. DSpark added its Markov bias, computed on the draft's GPU, to backbone logits left on the first GPU, so every sample raised a device error. This dates from DSpark's addition, and Domino's new decode needs the same handling.
3. DSpark ignored shift_label: false. It trained with the shift anyway and exported shift_label: false, so a loader that honors the field would read every drafted position one token off. Domino refused the same config, but only after building the draft.
4. The exporters re-derived head fields from the raw config. DSparkModule lower-cases markov_head_type, so Gated trained the gated head but exported "Gated". markov_rank, use_confidence_head and Domino's pure_draft_prefix_len were also re-read, each with its own default.

The fixes, one commit each:

• DSpark's decode keeps its logits on the draft's device. The backbone logits move to the draft's device before the Markov loop. On a single device this is a no-op.
• Domino generation.
  • DFlash's default decode starts at block position 1 - _label_shift and returns at most bs - 1 + _label_shift tokens. For shift 0 this is the same slice as before.
  • Domino overrides it the way DSpark does for its Markov head. It decodes the block left to right and steps the GRU over the anchor and the tokens drafted so far. From pure_draft_prefix_len on, it adds embed_proj([h_k; GRU state]) to the backbone logits, as _apply_domino_head does in training. Like DSpark's decode, it brings the backbone logits and the token embeddings to the draft's device.
  • The misplaced eval warning ("acceptance rates are backbone-only") is removed, along with every other copy of that claim: the comments in domino.yaml, dspark.yaml and the Nemotron-3.5 DSpark warm-start recipe, the Domino launcher example, and the speculative-decoding skill's Domino, DSpark and validation sheets. DSpark's generation has applied its Markov head since DSpark was added.
  • The Domino sheet and example now say how to add the HF AR eval: DFlash's task_2 with --steps equal to the block size, since Domino drafts the whole block. The flag values do not change.
• One alignment rule.
  • HFDFlashModel.modify refuses shift_label: false for every variant whose _label_shift is 1, before anything is built. Unset or true converts as before.
  • The refusal applies at convert only. restore_dflash_model drops shift_label before it replays convert, as it already drops the warm-start path. DSpark checkpoints saved with shift_label: false therefore still load, and now re-export shift_label: true.
  • Both exporters write shift_label from _label_shift. The modules' shift_label attributes and Domino's own check are removed, since nothing reads them any more.
• Exporters read the built head. DSparkExporter and DominoExporter read markov_rank, markov_head_type, use_confidence_head, pure_draft_prefix_len, gru_hidden_dim and emb_dim off the module, as LiLiCorrExporter already does for its reranker.
• Changelog. One 0.48.0 Bug Fixes entry.

Usage

N/A. Bug fix with no API change.

Testing

• New and changed tests:
  • TestDominoGenerate::test_drafts_are_the_training_head_argmax[0|1] checks that Domino's greedy drafts equal the argmax of _apply_domino_head teacher-forced on those drafts, for pure_draft_prefix_len 0 and 1.
    • The head weights are large enough to overturn the backbone's argmax.
    • The test fails on refactor(speculative): one home each for the DFlash family's shared rules #2757.
    • In a scratch run, it also fails when the head is skipped, when its GRU state lags one step, or when the prefix boundary is off by one.
  • test_next_token_variants_refuse_shift_label_false_at_convert[dspark|domino] converts both variants with shift_label unset, as the recipes leave it. It then restores a modelopt_state that carries shift_label: false, and expects false to be refused at convert. On refactor(speculative): one home each for the DFlash family's shared rules #2757 DSpark does not raise and Domino raises NotImplementedError; without the restore change, the restore raises.
  • test_export_config_has_dspark_fields now converts with markov_head_type: Gated and expects gated in config.json. It fails on refactor(speculative): one home each for the DFlash family's shared rules #2757.
  • GPU TestLmHeadOffTheDraftDevice::test_generate[dflash|dspark|domino] uses accelerate's dispatch_model to put embed_tokens and lm_head on the CPU, standing in for the first GPU, and the layers and draft on cuda:0. Generation must return on the caller's device. DSpark fails on refactor(speculative): one home each for the DFlash family's shared rules #2757, and Domino fails without its device move.
  • Each new unit test takes at most 0.12 s on 4 CPUs (taskset -c 0-3).
• tests/unit/torch/speculative and tests/unit/torch/export/test_hf_spec_rope_export.py on CPU: 415 passed on each of transformers 5.5.4, 5.12.1 and 5.18.0 (411 on refactor(speculative): one home each for the DFlash family's shared rules #2757). 5.5.4 and 5.18.0 fall within CI's tf_min and tf_latest pins.
• tests/gpu/torch/speculative/plugins on 1 GPU: 17 passed, 1 skipped. The skipped test needs 2 GPUs to shard the base.
• Scratch checks on 1 GPU, with the CPU standing in for the first GPU:
  • A tiny tied-embedding Qwen3 + Domino checkpoint, loaded through ar_validate.py's load_vlm_or_llm with real from_pretrained and accelerate dispatch, in both layouts transformers picks for 2 GPUs. The drafts equal the single-device drafts, and validate_online completes with the inputs on the first device, where ar_validate.py puts them. Without the device move in Domino's decode, the same load crashes.
  • The base's early hidden states moved to the CPU while the model runs on the GPU: DFlash, DSpark and Domino generation all gather them and return on the caller's device.
  • A tiny DSpark checkpoint saved by 0.47.0 with shift_label: false loads through from_pretrained and exports shift_label: true.
• Bitwise A/B against refactor(speculative): one home each for the DFlash family's shared rules #2757: one training step for each of 23 cases, plus generation at steps 0, 2 and 4 in the online ones, in fp32 and bf16.
  • 44 of 46 runs are identical: loss, every gradient, the variants' metrics and every generated token, including every DFlash, DFlash2, LiLiCorr and DSpark generation.
  • The other 2 are Domino's online case, in fp32 and bf16. Only its drafted tokens differ, as intended; its loss, gradients, metrics and base tokens are identical.
  • In fp32 at steps=4, the drafts go from [[22, 8, 8], [17, 17, 17]] (3 tokens, positions 1 to 3, backbone only) to [[8, 22, 8, 8], [21, 17, 17, 17]] (4 tokens, positions 0 to 3, head applied). The old tokens reappear one position later: this is the off-by-one.
• FSDP2 probe (2 gloo ranks, the draft as its own fully_shard unit, one rank on the zero-loss branch): DFlash finishes on both ranks, with the same losses as on refactor(speculative): one home each for the DFlash family's shared rules #2757. This PR does not change training code.
• pre-commit run --files <changed files>: clean.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅ Shipped recipes and existing checkpoints are unaffected. The recipes leave shift_label unset (or true, in domino.yaml), and checkpoints saved with shift_label: false still restore. Only a new conversion behaves differently: DSpark with shift_label: false now raises instead of training next-token alignment anyway and exporting false, and Domino raises ValueError for that config instead of NotImplementedError.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ One Bug Fixes entry, since DSpark and Domino have shipped since 0.46.0.
• Did you get Claude approval on this PR?: ❌ Not run yet.

Additional Information

Follow-up to #2757; review after it.

Found while testing, not fixed here. Both fail the same way on #2757:

• With DFlashModule as its own FSDP2 unit, DSpark and Domino training fails with got mixed torch.Tensor and DTensor.
  • The heads' layers (DSpark's Markov and confidence layers, Domino's prefix_gru and embed_proj) run outside DFlashModule.forward, so FSDP2 never unshards their parameters.
  • With the draft in the root unit, both train on both ranks.
  • The only launcher that makes DFlashModule its own unit is the MiniMax-M2.7 DFlash one, so no shipped DSpark or Domino setup hits this.
• AcceptanceRateValidation.validate_online compares the base lm_head's predictions with drafts returned on the caller's device. When device_map="auto" puts an untied lm_head on another GPU than the inputs, it raises a device error for every DFlash-family variant. This was reproduced on 1 GPU with the CPU as the first device, not on a 2-GPU host.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}