feat(speculative): D-PACE position weighting for DSpark and Domino

CHANGELOG.rst

@@ -37,6 +37,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, LiLiCorr, DSpark and Domino 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``, ``lilicorr_metrics``, ``dspark_metrics`` and ``domino_metrics``. DSpark's ``ce_loss``/``l1_loss``/``base_accuracy`` are now ``loss_ce``/``loss_tvd``/``backbone_accuracy``, and Domino's ``final_loss``/``base_loss``/``base_accuracy`` are now ``loss_final``/``loss_backbone``/``backbone_accuracy``.
+- DSpark and Domino training now honor ``dflash_loss_objective`` instead of always applying the static ``dflash_loss_decay_factor`` decay, so they train with D-PACE by default. Set ``dflash_loss_objective: decay`` to keep the previous weighting, as the shipped ``dspark.yaml`` and ``domino.yaml`` recipes now do, and finish runs started on an earlier version before upgrading: resuming one restores the objective saved in its checkpoint (``dpace`` unless the run set it) instead of the recipe's.
 
 **Deprecations**
 

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -895,11 +895,6 @@ def _draft_logits(self, batch):
         logits = logits.view(*batch.anchors.shape, self.dflash_block_size, -1)
         return DraftLogits(final=logits, backbone=logits)
 
-    @property
-    def _position_weighting(self):
-        """How block positions are weighted in the loss: ``"dpace"`` or ``"decay"``."""
-        return self.dflash_loss_objective
-
     def _loss_context(self, batch):
         """The step's labels, the draft's logits and the position weights."""
         targets = _block_targets(
@@ -915,7 +910,7 @@ def _loss_context(self, batch):
             self._draft_logits(batch),
             targets,
             weighting={
-                "scheme": self._position_weighting,
+                "scheme": self.dflash_loss_objective,
                 "first": 1 - self._label_shift,
                 "decay_factor": self.dflash_loss_decay_factor,
                 "dpace_alpha": self.dflash_dpace_alpha,
@@ -975,7 +970,7 @@ def forward(
         - Position weighting by D-PACE or exponential decay
 
         A variant changes the step only through its hooks: ``_online_base_outputs``,
-        ``_draft_logits`` (its head), ``_position_weighting`` and ``_compute_loss``.
+        ``_draft_logits`` (its head) and ``_compute_loss``.
         """
         if not self.training:
             if self.dflash_offline:

modelopt/torch/speculative/plugins/hf_domino.py

@@ -138,11 +138,6 @@ def _build_objective(self, config):
         """CE on the labels, for the corrected logits and for the backbone's."""
         return [objective_term("ce", 1.0)]
 
-    @property
-    def _position_weighting(self):
-        # dflash_loss_objective is not wired for Domino yet: it always decays.
-        return "decay"
-
     def _online_base_outputs(self, input_ids, attention_mask, **kwargs):
         """The target hidden states; Domino does not use target-logit KD."""
         with torch.no_grad():
@@ -181,7 +176,10 @@ def _draft_logits(self, batch):
         )
 
     def _compute_loss(self, ctx):
-        """``(1 - lambda_base) * objective(final) + lambda_base * objective(backbone)``."""
+        """``(1 - lambda_base) * objective(final) + lambda_base * objective(backbone)``.
+
+        Both share the step's position weights; D-PACE reads them off the served logits.
+        """
         out = super()._compute_loss(ctx)
         lambda_base = self._current_lambda_base()
         loss = backbone_loss = out.loss

modelopt/torch/speculative/plugins/hf_dspark.py

@@ -119,11 +119,6 @@ def _build_objective(self, config):
             objective_term("tvd", config.dflash_l1_loss_alpha),
         ]
 
-    @property
-    def _position_weighting(self):
-        # dflash_loss_objective is not wired for DSpark yet: it always decays.
-        return "decay"
-
     def _online_base_outputs(self, input_ids, attention_mask, **kwargs):
         """The target hidden states, and the base's final hidden for the loss to project."""
         # The inner model, not the CausalLM: the loss projects only the rows it reads.

modelopt_recipes/general/speculative_decoding/domino.yaml

@@ -64,7 +64,9 @@ dflash:
   dflash_use_torch_compile: false
   # Domino does not use target-logit KD; it trains its own base/final CE losses.
   dflash_self_logit_distillation: false
-  # gamma for exponential loss decay (block_size=16 -> 7).
+  # Static exponential position decay, gamma=7 for block_size=16: the weighting this recipe
+  # was tuned with. The default objective, 'dpace', would ignore the decay factor.
+  dflash_loss_objective: decay
   dflash_loss_decay_factor: 7.0
   # Qwen3 has no native mask token; 151669 is an unused id used by the reference.
   dflash_mask_token_id: 151669

modelopt_recipes/general/speculative_decoding/dspark.yaml

@@ -68,7 +68,9 @@ dflash:
   # DSpark computes the target distribution internally for its TVD/confidence
   # terms; the DFlash KD path is unused.
   dflash_self_logit_distillation: false
-  # gamma for exponential loss decay (block_size=16 -> 7).
+  # Static exponential position decay, gamma=7 for block_size=16: the weighting this recipe
+  # was tuned with. The default objective, 'dpace', would ignore the decay factor.
+  dflash_loss_objective: decay
   dflash_loss_decay_factor: 7.0
   # Qwen3 has no native mask token; 151669 is an unused id used by the reference.
   dflash_mask_token_id: 151669

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

@@ -52,7 +52,8 @@ Everything in `dflash.md` applies. Domino adds:
 `dflash_self_logit_distillation` is **false** — Domino trains its own base/final CE
 losses rather than distilling target logits. Recipe defaults also differ from
 DFlash's: `block_size` 16, `num_anchors` 256, `num_train_epochs` 6,
-`training_seq_len` 3072, `warmup_ratio` 0.04, `dflash_loss_decay_factor` 7.0.
+`training_seq_len` 3072, `warmup_ratio` 0.04, and the static position decay
+(`dflash_loss_objective: decay`, `dflash_loss_decay_factor` 7.0) instead of D-PACE.
 `max_grad_norm: 1.0` is stated explicitly in the recipe but is *not* a delta — it is the
 `transformers.TrainingArguments` default that `dflash.yaml` inherits by not setting it.
 
@@ -101,5 +102,6 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i
 | `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 |
+| `dflash_loss_decay_factor=7.0 is ignored when dflash_loss_objective='dpace'` (warning) when resuming a run started on an earlier ModelOpt version | That version saved the unset objective as `dpace` but trained with the static decay; resuming restores the saved `dpace`, not the recipe's `decay` pin | Finish the run on the version that started it |
 | 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 |

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

@@ -91,7 +91,8 @@ Total loss is `ce_alpha*CE + l1_alpha*TVD + conf_alpha*confidence_BCE`.
 `dflash_self_logit_distillation` is **false** for DSpark — it computes the target
 distribution internally for the TVD and confidence terms, so the DFlash KD path is
 unused. Recipe defaults also differ from DFlash's: `block_size` 16, `num_anchors` 256,
-`num_train_epochs` 6, `training_seq_len` 3072, `warmup_ratio` 0.04.
+`num_train_epochs` 6, `training_seq_len` 3072, `warmup_ratio` 0.04, and the static position
+decay (`dflash_loss_objective: decay`, `dflash_loss_decay_factor` 7.0) instead of D-PACE.
 
 ## Per-model adjustments
 
@@ -149,6 +150,7 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i
 | `shift_label=False is not supported: projector_type 'dspark' trains only next-token alignment` | `dflash_architecture_config.shift_label: false` | Remove it; DSpark only trains next-token alignment |
 | `dflash_confidence_head_alpha > 0 but the confidence head was not built` | Loss term enabled without the head | Set `dflash_architecture_config.use_confidence_head=true`, or set the alpha to 0 |
 | `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export |
+| `dflash_loss_decay_factor=7.0 is ignored when dflash_loss_objective='dpace'` (warning) when resuming a run started on an earlier ModelOpt version | That version saved the unset objective as `dpace` but trained with the static decay; resuming restores the saved `dpace`, not the recipe's `decay` pin | Finish the run on the version that started it |
 | Draft trains but acceptance length is poor | Draft dims left at defaults instead of matching the base | Set the GQA/FFN dims explicitly (see *Per-model adjustments*) |
 | Acceptance length capped despite clean training (streaming) | `EAGLE_CAPTURE_IDS` final layer wrong, or the vLLM aux-capture fix (vllm#46788) missing | Correct the ids; use a container with the fix |
 | Trainer init crash on a serve container | tensorboard absent | `training.report_to=none` |

tests/unit/torch/speculative/plugins/test_hf_dflash.py

@@ -1615,6 +1615,51 @@ def test_next_token_variants_refuse_shift_label_false_at_convert(head_fields):
         mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)])
 
 
+@pytest.mark.parametrize("objective", ["dpace", "decay"])
+@pytest.mark.parametrize(
+    ("head_fields", "metric", "source"),
+    [
+        ({"projector_type": "dspark", "markov_rank": 4}, "loss_ce", "final"),
+        (
+            {"projector_type": "domino", "emb_dim": 8, "gru_hidden_dim": 8},
+            "loss_backbone",
+            "backbone",
+        ),
+    ],
+    ids=["dspark", "domino"],
+)
+def test_next_token_variants_weight_positions_by_the_objective(
+    objective, head_fields, metric, source
+):
+    """Weights follow dflash_loss_objective from block position 0; D-PACE reads the served logits.
+
+    Every objective term is reduced under them, Domino's backbone CE included.
+    """
+    config = get_dflash_config()
+    config["dflash_architecture_config"].update(head_fields)
+    config.update(dflash_loss_objective=objective, dflash_loss_decay_factor=4.0)
+    model = get_tiny_llama(num_hidden_layers=4)
+    mtsp.convert(model, [("dflash", config)])
+    model.float().train()
+    with patch.object(model, "_compute_loss", wraps=model._compute_loss) as compute_loss:
+        out = model(**_dflash_batch(model.config.vocab_size))
+    ctx = compute_loss.call_args.args[0]
+    assert not torch.equal(ctx.logits.final, ctx.logits.backbone)
+
+    def ce(logits):
+        return F.cross_entropy(
+            logits.flatten(0, 2), ctx.targets.token_ids.flatten(), reduction="none"
+        )
+
+    mask = ctx.targets.mask
+    weights = _block_position_weights(
+        mask, objective, 0, 4.0, 0.5, target_ce=lambda: ce(ctx.logits.final).view_as(mask)
+    )
+    assert torch.equal(ctx.weights, weights)
+    want = (ce(getattr(ctx.logits, source)) * weights.flatten()).sum() / (weights.sum() + 1e-6)
+    assert torch.equal(out.draft_metrics[metric], want)
+
+
 class TestTeacherLogits:
     """The base distribution is projected only at the rows, and vocab entries, asked for.
 

tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml

@@ -63,6 +63,7 @@ pipeline:
       # Domino knobs (also set in the recipe; repeated here for visibility).
       - dflash.dflash_block_size=16
       - dflash.dflash_num_anchors=512
+      - dflash.dflash_loss_objective=decay
       - dflash.dflash_loss_decay_factor=7
       - dflash.dflash_mask_token_id=151669
       - dflash.dflash_self_logit_distillation=false