[1/2] refactor(speculative): hand the DFlash-family loss one context per training step

.pre-commit-config.yaml

@@ -119,6 +119,7 @@ repos:
               modelopt/torch/quantization/plugins/attention.py|
               modelopt/torch/sparsity/attention_sparsity/methods/vsa_utils.py|
               modelopt/torch/speculative/eagle/utils.py|
+              modelopt/torch/speculative/plugins/dflash_losses.py|
               modelopt/torch/speculative/plugins/hf_domino.py|
               modelopt/torch/speculative/plugins/modeling_domino.py|
               modelopt/torch/speculative/plugins/hf_dflash.py|

CHANGELOG.rst

@@ -36,6 +36,7 @@ Changelog
 
 - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata.
+- The DFlash2 and LiLiCorr training forwards now return their metrics in one ``draft_metrics`` dict of detached 0-dim tensors (call ``.item()`` for a Python float), replacing ``selector_metrics`` and ``lilicorr_metrics``.
 
 **Deprecations**
 
@@ -46,6 +47,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/config.py

@@ -340,7 +340,7 @@ class DFlashConfig(ModeloptBaseConfig):
             "absolute — there is no outer multiplier scaling the three terms as a group — so "
             "each is the coefficient with which its term enters the total, and "
             "`loss == origin_loss + lilicorr_loss` holds exactly. Both halves and all three "
-            "weights are reported in the `lilicorr_metrics` dict the model attaches to its "
+            "weights are reported in the `draft_metrics` dict the model attaches to its "
             "forward output, so a consumer of those metrics can check the identity per step "
             "and read which composition produced a checkpoint. The three weights are validated "
             "all-or-nothing (a negative value means unset): a config that sets some but not "

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]