From 99efe27b924719ac4ec21788b096629774b277a5 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:43:29 +0000 Subject: [PATCH 1/3] chore(speculative): pin the DSpark and Domino recipes to the static decay dspark.yaml and domino.yaml set dflash_loss_decay_factor: 7.0 without an objective. DSpark and Domino ignore dflash_loss_objective today and always decay; once they honor it, the default 'dpace' would switch both recipes to D-PACE and drop the decay factor they were tuned with. Pin dflash_loss_objective: decay in both, and list it among the knobs the Qwen3-8B Domino launcher repeats for visibility, as the LiLiCorr launcher does. The DSpark launchers (Kimi-K2.6, MiniMax-M3) inherit the pin from dspark.yaml. No behavior change at this commit. Verified that both recipes load through ModelOptDFlashRecipe with objective 'decay' and decay factor 7.0; pre-commit's recipe and launcher-reference validators pass. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt_recipes/general/speculative_decoding/domino.yaml | 4 +++- modelopt_recipes/general/speculative_decoding/dspark.yaml | 4 +++- tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml | 1 + 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/modelopt_recipes/general/speculative_decoding/domino.yaml b/modelopt_recipes/general/speculative_decoding/domino.yaml index e7237b11b8f..055f78508ad 100644 --- a/modelopt_recipes/general/speculative_decoding/domino.yaml +++ b/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 diff --git a/modelopt_recipes/general/speculative_decoding/dspark.yaml b/modelopt_recipes/general/speculative_decoding/dspark.yaml index db69f3c4b9e..084f7145b24 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark.yaml +++ b/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 diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml index 7adf3400cd7..8760c2047af 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml +++ b/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 From 7de62ebe6edaf804b5a51c19b5453aa645e84b9c Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:58:45 +0000 Subject: [PATCH 2/3] feat(speculative): weight DSpark and Domino positions by dflash_loss_objective DSpark and Domino overrode the shared forward's _position_weighting hook to return "decay", so dflash_loss_objective, whose default is "dpace", had no effect on them, and DFlashModel.modify's warning that the decay factor is ignored under D-PACE was false for both. They now follow the objective like the rest of the family, from their first predicted block position, 0. D-PACE reads the draft's confidence off the served (final) logits, and that one set of weights reduces every term of the step: DSpark's CE, TVD and confidence BCE, and both of Domino's objectives, final and backbone. With no override left, the hook goes and _loss_context reads dflash_loss_objective directly. The new test pins, for both variants and both objectives, the weights handed to the loss to _block_position_weights from position 0 on the final logits, and the reported CE (DSpark) and backbone CE (Domino) to their reduction under those weights. It fails on the parent for D-PACE, and it fails when the confidence is read off the backbone, when the first position is 1, and when Domino's backbone CE gets its own D-PACE weights. Verified against the parent with one training and one generation step per case, in fp32 and bf16. The 23 existing cases (every variant, offline and online, decay and D-PACE, zero-loss steps included) are bitwise equal in loss, gradients, metrics and generated tokens. The five new DSpark/Domino D-PACE training cases change loss and gradients, as intended, generate the same tokens, and equal bitwise the parent with the two overrides deleted at runtime; their two zero-loss cases are unchanged. The DFlash-family plugin tests and the recipe loader tests pass on transformers 5.5 and 5.18 (806, the 4 new cases included), the GPU plugin tests pass, and the FSDP2 DFlash probes pass on both ranks. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/speculative/plugins/hf_dflash.py | 9 +--- .../torch/speculative/plugins/hf_domino.py | 10 ++--- .../torch/speculative/plugins/hf_dspark.py | 5 --- .../speculative/plugins/test_hf_dflash.py | 45 +++++++++++++++++++ 4 files changed, 51 insertions(+), 18 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index a0c2cc34bb1..fa5c3e7f780 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/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: diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 0a4bc9fdbd4..658ca266c37 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/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 diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 69b5fdc3f1f..34a7391162f 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/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. diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ec86aa115c7..84793ae4442 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/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. From 9ed2dd2651a8b8aaa1e62d3ce03d74c4de56ee85 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:59:21 +0000 Subject: [PATCH 3/3] docs(speculative): note that DSpark and Domino follow dflash_loss_objective Users of a released DSpark or Domino who leave dflash_loss_objective unset move from the static decay to D-PACE, so the changelog lists it under backward breaking changes with the one-line way back. That way back does not reach a run resumed across the upgrade: its checkpoint stores the default 'dpace' the old code ignored, and main.py's HF-format resume restores it without re-reading the recipe, so the run continues with D-PACE. The entry says to finish such runs on the version that started them, and the DSpark and Domino skill sheets key the case on the warning the restore logs. The sheets also list the pinned static decay among the recipes' deltas from DFlash. Verified with pre-commit (rst, markdownlint, skill symlink sync), and by resuming DSpark and Domino checkpoints through main.py's HF-format loader: saved from the branch base's shipped recipes, they restore 'dpace' and weight the next step with the static decay on that base but with D-PACE here; saved from the pinned recipes, they restore 'decay' and keep it. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 1 + .../speculative-decoding/references/algorithms/domino.md | 4 +++- .../speculative-decoding/references/algorithms/dspark.md | 4 +++- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2907f626226..fc109c985e5 100755 --- a/CHANGELOG.rst +++ b/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** diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md index 8494f6f4b76..074b2420d4b 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md +++ b/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 | diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md index bb2c6992ac5..36f2e7849d9 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md +++ b/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` |