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

CHANGELOG.rst

@@ -46,6 +46,7 @@ Changelog
 - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root.
 - Fix GPT-OSS draft training with pre-normalization teacher hidden states by enabling the target model's final normalization.
 - Fix vanilla DSpark exports to use the Qwen3 DSpark architecture, head names and configuration expected by vLLM. Re-export existing training checkpoints to use this format.
+- Fix Domino AR validation (``pseudo_speculative_generate``), which scored the backbone alone at the wrong block positions, and a DSpark AR-validation crash when ``device_map="auto"`` places a tied ``lm_head`` on another GPU than the draft. DSpark and Domino now reject ``dflash_architecture_config.shift_label: false`` at conversion (DSpark trained with the shift regardless) and export the ``shift_label`` they were trained with; checkpoints saved with ``shift_label: false`` still load.
 - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``.
 - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source.
 - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU.

modelopt/torch/export/plugins/hf_spec_export.py

@@ -521,19 +521,19 @@ def _export_config(self):
         """Extend the DFlash config with the Domino head fields."""
         config = super()._export_config()
         draft_config = self.model.dflash_config
+        # Read off the built head rather than the config dict, so the export describes the
+        # head these weights belong to.
+        head = self.model.dflash_module
 
-        # Present because HFDominoModel.modify validates them at convert time.
-        emb_dim = draft_config.emb_dim
-        gru_hidden_dim = draft_config.gru_hidden_dim
         # Mirror the reference checkpoint: emb_dim also appears at the top level.
-        config["emb_dim"] = emb_dim
+        config["emb_dim"] = head.emb_dim
         config["dflash_config"].update(
             {
                 "projector_type": getattr(draft_config, "projector_type", "domino"),
-                "shift_label": getattr(draft_config, "shift_label", True),
-                "pure_draft_prefix_len": getattr(draft_config, "pure_draft_prefix_len", 1),
-                "gru_hidden_dim": gru_hidden_dim,
-                "emb_dim": emb_dim,
+                "shift_label": bool(self.model._label_shift),
+                "pure_draft_prefix_len": head.pure_draft_prefix_len,
+                "gru_hidden_dim": head.gru_hidden_dim,
+                "emb_dim": head.emb_dim,
             }
         )
         return config
@@ -615,14 +615,17 @@ def _export_config(self):
         """Extend the DFlash config with the DSpark head fields."""
         config = super()._export_config()
         draft_config = self.model.dflash_config
+        # Read off the built head, not the config dict: DSparkModule casts markov_rank to int
+        # and lower-cases markov_head_type, and the export must describe what was trained.
+        head = self.model.dflash_module
 
         config["dflash_config"].update(
             {
                 "projector_type": getattr(draft_config, "projector_type", "dspark"),
-                "shift_label": getattr(draft_config, "shift_label", True),
-                "markov_rank": draft_config.markov_rank,
-                "markov_head_type": self.model.dflash_module.markov_head_type,
-                "use_confidence_head": bool(getattr(draft_config, "use_confidence_head", False)),
+                "shift_label": bool(self.model._label_shift),
+                "markov_rank": head.markov_rank,
+                "markov_head_type": head.markov_head_type,
+                "use_confidence_head": head.use_confidence_head,
             }
         )
         if config["dflash_config"]["markov_head_type"] == "vanilla":
@@ -634,7 +637,7 @@ def _export_config(self):
                     "mask_token_id": config["dflash_config"]["mask_token_id"],
                     "target_layer_ids": config["dflash_config"]["target_layer_ids"],
                     "markov_rank": config["dflash_config"]["markov_rank"],
-                    "sample_from_anchor": True,
+                    "sample_from_anchor": config["dflash_config"]["shift_label"],
                     "enable_confidence_head": config["dflash_config"]["use_confidence_head"],
                     "confidence_head_with_markov": config["dflash_config"]["use_confidence_head"],
                 }

modelopt/torch/speculative/dflash/conversion.py

@@ -95,4 +95,9 @@ def restore_dflash_model(
     # Warm start is a one-time training-time init; the restored weights would overwrite it
     # anyway, and replaying it on the meta device would raise.
     config = config.model_copy(update={"dflash_init_checkpoint": None})
+    # shift_label: false is refused at convert only: DSpark drafters saved with it before the
+    # refusal trained next-token alignment regardless, and must still load.
+    config.dflash_architecture_config = {
+        k: v for k, v in config.dflash_architecture_config.items() if k != "shift_label"
+    }
     return convert_to_dflash_model(model, config)[0]

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -382,6 +382,16 @@ def modify(self, config):
         """Initialize DFlash draft module."""
         super().modify(config)
 
+        # shift_label names a next-token variant's alignment in its export; the alignment is
+        # fixed by _label_shift, so a config turning it off is refused rather than ignored.
+        arch_config = config.dflash_architecture_config
+        if self._label_shift and not arch_config.get("shift_label", True):
+            raise ValueError(
+                f"dflash_architecture_config.shift_label={arch_config['shift_label']!r} is not "
+                f"supported: projector_type {arch_config.get('projector_type')!r} trains only "
+                "next-token alignment (block position k predicts the token at anchor + k + 1)."
+            )
+
         self.dflash_fp32_master_weights = config.dflash_fp32_master_weights
 
         base_config = self._base_llm_config
@@ -1219,36 +1229,42 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
            The draft model's KV injection concatenates projected target_hidden as K/V
            with the block's own K/V, enabling the draft to "see" the target's context.
 
-        6. **Decode**: Apply ``lm_head`` to draft hidden states at positions 1..block_size-1
-           (skip position 0 which is the known anchor). Argmax gives draft tokens.
+        6. **Decode**: Apply ``lm_head`` to draft hidden states from position
+           ``1 - _label_shift`` on, the first one that predicts a token after the anchor
+           (position 1 for DFlash, whose position 0 is the known anchor). Argmax gives draft
+           tokens.
 
         7. **Return**: ``(base_token, draft_tokens[:steps])`` — base token is always
-           returned; draft tokens are truncated to ``steps`` (default: block_size-1).
+           returned; draft tokens are truncated to ``steps``.
 
         Note:
             This method re-runs the full target model from scratch on each call
             (no KV cache). For AR validation, it is called repeatedly with growing
             ``input_ids`` by ``AcceptanceRateValidation.validate()``. The ``steps``
-            parameter should be set to ``block_size - 1`` for full block evaluation.
+            parameter should be set to ``block_size - 1 + _label_shift`` for full block
+            evaluation.
 
         Args:
             input_ids: Input token IDs [B, seq_len].
-            steps: Number of draft tokens to return (capped at block_size-1).
+            steps: Number of draft tokens to return (capped at block_size - 1 + _label_shift).
 
         Returns:
             base_token: Next token from base model [B, 1].
-            draft_tokens: Draft tokens [B, min(steps, block_size-1)] or None if steps < 1.
+            draft_tokens: Draft tokens [B, min(steps, block_size - 1 + _label_shift)] or None
+                if steps < 1.
         """
         base_token, hidden_states = self._generation_base(input_ids)
         if steps < 1:
             return base_token, None
         draft_hidden = self._draft_block(base_token, hidden_states)
 
-        # Logits on positions 1..block_size-1 (skip anchor at position 0)
-        draft_logits = self._base_model_lm_head(draft_hidden[:, 1:, :])
-        draft_tokens = draft_logits.argmax(dim=-1)  # [B, block_size-1]
+        # Position k predicts the token at anchor + k + _label_shift, so this one drafts the
+        # first token after the anchor.
+        first = 1 - self._label_shift
+        draft_logits = self._base_model_lm_head(draft_hidden[:, first:, :])
+        draft_tokens = draft_logits.argmax(dim=-1)  # [B, block_size - first]
 
         # Return up to `steps` tokens. base_token already follows input_ids; draft_tokens is
         # produced on the draft's device, and callers cat both onto the running sequence.
-        num_tokens = min(steps, self.dflash_block_size - 1)
+        num_tokens = min(steps, self.dflash_block_size - first)
         return base_token, draft_tokens[:, :num_tokens].to(input_ids.device)

modelopt/torch/speculative/plugins/hf_domino.py

@@ -132,10 +132,6 @@ def modify(self, config):
         self.dflash_lambda_base_start = config.dflash_lambda_base_start
         self.dflash_lambda_base_decay_ratio = config.dflash_lambda_base_decay_ratio
         self._lambda_base = self.dflash_lambda_base_start
-        if not getattr(self.dflash_module, "shift_label", True):
-            raise NotImplementedError(
-                "Domino currently supports shift_label=True (next-token alignment) only."
-            )
 
     def get_exporter(self):
         """Get the exporter for the Domino draft model."""
@@ -266,14 +262,6 @@ def forward(
         two-term loss. Eval/offline-eval is delegated to the DFlash parent.
         """
         if not self.training:
-            # Eval delegates to the DFlash backbone; the correction head is not
-            # applied yet, so warn once that acceptance rates are backbone-only.
-            if not getattr(self, "_warned_eval_head_bypass", False):
-                logger.warning(
-                    "Domino eval uses the DFlash backbone only (correction head not "
-                    "applied yet); reported acceptance rates are backbone-only."
-                )
-                self._warned_eval_head_bypass = True
             return super().forward(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
@@ -365,6 +353,38 @@ def forward(
 
         return ModelOutput(loss=loss, logits=None, train_acc=[[accuracy]], domino_metrics=metrics)
 
+    @torch.no_grad()
+    def pseudo_speculative_generate(self, input_ids, steps=1):
+        """Draft the block left to right, adding the correction head as training does.
+
+        Block position ``k`` predicts the token after the anchor and the ``k`` tokens drafted
+        before it. From ``pure_draft_prefix_len`` on, its backbone logits get the correction
+        from the GRU state over exactly those tokens.
+        """
+        base_token, hidden_states = self._generation_base(input_ids)
+        if steps < 1:
+            return base_token, None
+        draft_hidden = self._draft_block(base_token, hidden_states)
+        device = draft_hidden.device
+        # Under device_map="auto" the embedding table and a tied lm_head sit on the first GPU,
+        # not the draft's, and their outputs stay there.
+        backbone_logits = self._base_model_lm_head(draft_hidden).to(device)  # [B, block_size, V]
+
+        m = self.dflash_module
+        prev_token = base_token.squeeze(-1).to(device)
+        gru_state = None
+        draft_tokens = []
+        for k in range(min(steps, self.dflash_block_size)):
+            prev_emb = self._base_model_embeddings(prev_token).to(device)
+            gru_out, gru_state = m.prefix_gru(prev_emb.unsqueeze(1), gru_state)
+            logits = backbone_logits[:, k, :]
+            if k >= m.pure_draft_prefix_len:
+                correction = m.embed_proj(torch.cat([draft_hidden[:, k, :], gru_out[:, 0, :]], -1))
+                logits = logits + correction
+            prev_token = logits.argmax(dim=-1)
+            draft_tokens.append(prev_token)
+        return base_token, torch.stack(draft_tokens, dim=1).to(input_ids.device)
+
 
 class DominoLambdaCallback(TrainerCallback):
     """Update the model's ``lambda_base`` from the HF Trainer global step.

modelopt/torch/speculative/plugins/hf_dspark.py

@@ -419,7 +419,8 @@ def pseudo_speculative_generate(self, input_ids, steps=1):
         if steps < 1:
             return base_token, None
         draft_hidden = self._draft_block(base_token, hidden_states)
-        backbone_logits = self._base_model_lm_head(draft_hidden)  # [B, block_size, V]
+        # Under device_map="auto" a tied lm_head sits with the embedding table, not the draft.
+        backbone_logits = self._base_model_lm_head(draft_hidden).to(draft_hidden.device)
 
         # Autoregressive Markov sampling over the block.
         m = self.dflash_module

modelopt/torch/speculative/plugins/modeling_domino.py

@@ -82,7 +82,6 @@ def __init__(self, config):
                 f"pure_draft_prefix_len must be in [0, {self.block_size - 1}] "
                 f"(block_size={self.block_size}), got {self.pure_draft_prefix_len}."
             )
-        self.shift_label = getattr(config, "shift_label", True)
 
         # Causal state over the block's token embeddings. bias=False matches the
         # reference checkpoint (only weight_ih_l0 / weight_hh_l0 are stored).

modelopt/torch/speculative/plugins/modeling_dspark.py

@@ -89,10 +89,6 @@ def __init__(self, config):
         self.projector_type = getattr(config, "projector_type", "dspark")
         self.markov_rank = int(config.markov_rank)
         self.markov_head_type = str(getattr(config, "markov_head_type", "vanilla")).lower()
-        # DSpark treats the anchor as the first prediction position (next-token
-        # alignment), matching Domino's shift_label path; kept as an attribute so
-        # the wrapper and exporter can read it uniformly.
-        self.shift_label = True
         if self.markov_rank <= 0:
             raise ValueError(f"DSpark requires markov_rank > 0, got {self.markov_rank}.")
         if self.markov_head_type not in ("vanilla", "gated", "rnn"):

modelopt_recipes/general/speculative_decoding/domino.yaml

@@ -37,8 +37,6 @@ training:
   cp_size: 1
   dp_shard_size: 1
   disable_tqdm: true
-  # Keep off: eval runs the DFlash backbone only (correction head not applied
-  # yet), so AR would reflect the backbone alone, not the trained model.
   estimate_ar: false
   ar_validate_steps: 0
   answer_only_loss: true

modelopt_recipes/general/speculative_decoding/dspark.yaml

@@ -40,9 +40,6 @@ training:
   cp_size: 1
   dp_shard_size: 1
   disable_tqdm: true
-  # Keep off: eval runs the DFlash backbone only (Markov head not applied yet),
-  # so AR would reflect the backbone alone, not the trained model. Compare via
-  # export + the offline acceptance-length harness instead.
   estimate_ar: false
   ar_validate_steps: 0
   answer_only_loss: true

modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml

@@ -54,7 +54,6 @@ training:
   cp_size: 1
   dp_shard_size: 1
   disable_tqdm: true
-  # Eval runs the DFlash backbone only (no Markov head), so AR would understate the model.
   estimate_ar: false
   ar_validate_steps: 0
   answer_only_loss: true

plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md

@@ -28,16 +28,12 @@ The committed example is **online**, 2 tasks:
 `--full-conversations` matters: it keeps real assistant completions so
 `answer_only_loss` has assistant spans to mask.
 
-**The inference side is intentionally not wired up yet.** The Domino correction head
-is not applied in `pseudo_speculative_generate` or in the serving stack, so the
-example ships no vLLM smoke test and no MT-Bench AR eval. Do not treat their absence
-as a broken config.
-
-When that path lands, copy the smoke-test and AR-eval steps from
-`hf_online_dflash.yaml` (its `task_1` and `task_2`) and append them to the Domino
-config as **`task_2` and `task_3`** — Domino's `task_0`/`task_1` are already the
-dataset build and training, so keep the source file's numbering and the destination's
-distinct.
+The example ships no vLLM smoke test or MT-Bench AR eval; do not treat their absence
+as a broken config. Serving is not wired up yet, but HF AR validation
+(`pseudo_speculative_generate`) applies the correction head: to measure AR, append
+`hf_online_dflash.yaml`'s `task_2` as Domino's **`task_2`** with `--steps` equal to the
+block size (Domino drafts the whole block: 16, not DFlash's 15). Add the smoke test
+(`task_1` there) as `task_3` once serving lands.
 
 ## Recipe and training knobs
 
@@ -87,11 +83,6 @@ export is the whole in-pipeline signal.
 
 ## Quality gate
 
-**Do not trust in-training AR for Domino.** The recipe pins `estimate_ar: false` and
-`ar_validate_steps: 0` deliberately: eval delegates to the DFlash backbone with the
-correction head not applied, so reported acceptance rates are backbone-only. The code
-logs this once as a warning — treat that warning as expected, not as a defect.
-
 The training-regression gate from `dflash.md` applies and the Qwen3-8B example sets
 it: `MAX_FINAL_LOSS=5.0`, `MIN_FINAL_ACC=0.15`, checked by `check_regression.py`
 against `trainer_state.json`. Since no inference metric is produced, this is currently
@@ -107,9 +98,8 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i
 | Error pattern | Root cause | Fix |
 | --- | --- | --- |
 | `Domino (projector_type='domino') requires ['emb_dim', 'gru_hidden_dim'] in dflash_architecture_config` | GRU head dims missing | Set both in `dflash_architecture_config` |
-| `Domino currently supports shift_label=True (next-token alignment) only` | `shift_label=false` | Leave it at `true` |
+| `shift_label=False is not supported: projector_type 'domino' trains only next-token alignment` | `shift_label: false` | Leave it at `true` |
 | `pure_draft_prefix_len must be in [0, N] (block_size=B), got X` | Prefix length >= block size | Lower it below `block_size` |
 | `DominoLambdaCallback: state.max_steps unset (<=0); lambda_base curriculum disabled` (warning) | `training.max_steps` not set | Set `training.max_steps`, else the curriculum never runs |
-| `Domino eval uses the DFlash backbone only ...` (warning) | Correction head not applied at eval | Expected — do not chase it; evaluate after export once the inference path lands |
 | DDP error about unused parameters | `ddp_find_unused_parameters` turned off | Keep it `true` |
 | Draft trains but quality is poor | Draft dims left at `Qwen3Config` defaults instead of matching the base | Set the GQA/FFN dims explicitly |