From d78031d1fedf094c32c1f0721041fc9447cdee7d Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:35:37 +0000 Subject: [PATCH 1/4] fix(speculative): one shift_label rule for DSpark and Domino DSpark and Domino train block position k on the token at anchor + k + 1 (HFDFlashModel._label_shift = 1). The same fact was also written as DSparkModule.shift_label = True, as DominoModule.shift_label read from the config, as a NotImplementedError in HFDominoModel.modify, and as the exporters' getattr(draft_config, "shift_label", True), which read the user's raw architecture config. A DSpark config with shift_label: false therefore trained with the shift but 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. HFDFlashModel.modify now refuses shift_label: false for every variant whose _label_shift is 1, before anything is built; unset or true converts as before. Both exporters write shift_label from _label_shift, and the vanilla DSpark export copies it into sample_from_anchor, vLLM's name for the same alignment, instead of writing a second literal True. The modules' shift_label attributes and Domino's own check go, since nothing reads them any more. Domino's error becomes a ValueError, like the other convert-time config checks, and the skill's Domino and DSpark sheets list it. The refusal applies at convert only. Restore replays convert, so restore_dflash_model drops shift_label from the saved architecture config, as it already drops the warm-start path. A DSpark checkpoint saved with shift_label: false trained next-token alignment all the same: it still loads, and now re-exports shift_label: true. Verified: a new test converts DSpark and Domino with shift_label unset, as the recipes leave it, restores a modelopt_state that carries shift_label: false, and expects shift_label: false to be refused at convert. On the parent commit DSpark does not raise and Domino raises NotImplementedError; without the restore change the restore raises. A tiny DSpark checkpoint saved by 0.47.0 with shift_label: false loads through from_pretrained and exports shift_label: true. The existing export tests still read shift_label: true and, for vanilla heads, sample_from_anchor: true. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> # Conflicts: # tests/unit/torch/speculative/plugins/test_hf_dflash.py --- .../torch/export/plugins/hf_spec_export.py | 6 ++-- .../torch/speculative/dflash/conversion.py | 5 ++++ .../torch/speculative/plugins/hf_dflash.py | 10 +++++++ .../torch/speculative/plugins/hf_domino.py | 4 --- .../speculative/plugins/modeling_domino.py | 1 - .../speculative/plugins/modeling_dspark.py | 4 --- .../references/algorithms/domino.md | 2 +- .../references/algorithms/dspark.md | 1 + .../speculative/plugins/test_hf_dflash.py | 30 +++++++++++++++++++ 9 files changed, 50 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 92ce9121b53..461c5c3bce4 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -530,7 +530,7 @@ def _export_config(self): config["dflash_config"].update( { "projector_type": getattr(draft_config, "projector_type", "domino"), - "shift_label": getattr(draft_config, "shift_label", True), + "shift_label": bool(self.model._label_shift), "pure_draft_prefix_len": getattr(draft_config, "pure_draft_prefix_len", 1), "gru_hidden_dim": gru_hidden_dim, "emb_dim": emb_dim, @@ -619,7 +619,7 @@ def _export_config(self): config["dflash_config"].update( { "projector_type": getattr(draft_config, "projector_type", "dspark"), - "shift_label": getattr(draft_config, "shift_label", True), + "shift_label": bool(self.model._label_shift), "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)), @@ -634,7 +634,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"], } diff --git a/modelopt/torch/speculative/dflash/conversion.py b/modelopt/torch/speculative/dflash/conversion.py index 859d960d0b2..c250ca48298 100644 --- a/modelopt/torch/speculative/dflash/conversion.py +++ b/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] diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 5d3fc293efc..f29d3171a31 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -393,6 +393,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 diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index efc22a47c46..a0cda038c7a 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -133,10 +133,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.""" diff --git a/modelopt/torch/speculative/plugins/modeling_domino.py b/modelopt/torch/speculative/plugins/modeling_domino.py index e6e9fc47c94..566967b3e7c 100644 --- a/modelopt/torch/speculative/plugins/modeling_domino.py +++ b/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). diff --git a/modelopt/torch/speculative/plugins/modeling_dspark.py b/modelopt/torch/speculative/plugins/modeling_dspark.py index e6464380d17..11f1ff66788 100644 --- a/modelopt/torch/speculative/plugins/modeling_dspark.py +++ b/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"): diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md index 3f8434d1000..ca35742b2d9 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md @@ -107,7 +107,7 @@ 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 | diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md index f857d193d3f..09ccbc70709 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md @@ -151,6 +151,7 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i | `DSpark (projector_type='dspark') requires 'markov_rank' (> 0) in dflash_architecture_config` | Markov head dimension missing | Set `dflash_architecture_config.markov_rank` | | `DSpark requires markov_rank > 0, got N` | Non-positive value | Set a positive rank | | `Unsupported markov_head_type: '...'. Expected 'vanilla', 'gated' or 'rnn'` | Typo or unsupported head | Use one of the three | +| `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 | | 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*) | diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index 55d2978a7e2..81c40c7da74 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -1574,6 +1574,36 @@ def keep_output(module, args, output): assert all(output.grad is not None for output in outputs) +@pytest.mark.parametrize( + "head_fields", + [ + {"projector_type": "dspark", "markov_rank": 4}, + {"projector_type": "domino", "emb_dim": 8, "gru_hidden_dim": 8}, + ], + ids=["dspark", "domino"], +) +def test_next_token_variants_refuse_shift_label_false_at_convert(head_fields): + """Left unset, as the recipes do, shift_label converts; set to false, it is refused. + + Restore does not refuse it: DSpark used to accept the field, ignore it in training and + save it into the checkpoint's modelopt_state. + """ + config = get_dflash_config() + config["dflash_architecture_config"].update(head_fields) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + + state = mto.modelopt_state(model) + saved_config = state["modelopt_state_dict"][0][1]["config"] + saved_config["dflash_architecture_config"]["shift_label"] = False + restored = mto.restore_from_modelopt_state(get_tiny_llama(num_hidden_layers=4), state) + assert restored._label_shift == 1 + + config["dflash_architecture_config"]["shift_label"] = False + with pytest.raises(ValueError, match="only next-token alignment"): + mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)]) + + class TestTeacherLogits: """The base distribution is projected only at the rows, and vocab entries, asked for. From da3b9ac4a4b3a4322b1d0a0f2fe93b0fa37542b6 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:14:33 +0000 Subject: [PATCH 2/4] fix(speculative): export the DSpark and Domino head fields the module was built with DSparkModule resolves its head fields from the architecture config as it builds them: it lower-cases markov_head_type, casts markov_rank to int and use_confidence_head to bool, and fills in defaults for the ones left out. #2902 made DSparkExporter read markov_head_type off the built module, but markov_rank and use_confidence_head were still read from the raw config dict, the latter with a second default. A markov_rank given as a float (64.0) therefore trained a rank-64 head but exported 64.0, and the vanilla export copies that value to the top level, where vLLM's Qwen3DSparkModel sizes its Markov table from it. DominoExporter likewise re-derived pure_draft_prefix_len with a second default. Both exporters now read every head field off the built module, as LiLiCorrExporter already does for its reranker, so the exported config describes the head that the saved weights belong to. Verified: the DSpark export-config test now converts with markov_rank given as a float and expects an int at both levels of config.json; all eight cases fail on the parent commit. Domino's exported values and DSpark's use_confidence_head are unchanged for every config, since the modules store them as given or cast them the same way. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/export/plugins/hf_spec_export.py | 23 +++++++++++-------- .../speculative/plugins/test_hf_dspark.py | 19 +++++++++++---- 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 461c5c3bce4..fdd76b37d8b 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/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": bool(self.model._label_shift), - "pure_draft_prefix_len": getattr(draft_config, "pure_draft_prefix_len", 1), - "gru_hidden_dim": gru_hidden_dim, - "emb_dim": emb_dim, + "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": bool(self.model._label_shift), - "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)), + "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": diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index c9000273988..4943da4c5e4 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -60,6 +60,7 @@ def _get_dspark_config( confidence_alpha=0.0, block_size=BLOCK_SIZE, num_layers=NUM_DRAFT_LAYERS, + markov_rank=MARKOV_RANK, ): """Create a DSpark config for testing (dflash mode + projector_type=dspark).""" config = deepcopy(DFLASH_DEFAULT_CFG["config"]) @@ -71,7 +72,7 @@ def _get_dspark_config( config["dflash_architecture_config"] = { "num_hidden_layers": num_layers, "projector_type": "dspark", - "markov_rank": MARKOV_RANK, + "markov_rank": markov_rank, "markov_head_type": head_type, "use_confidence_head": use_confidence_head, "pure_draft_prefix_len": 1, @@ -270,7 +271,9 @@ def spy(*args, **kwargs): class TestDSparkExporter: """Test vanilla serving exports and the native gated/recurrent formats.""" - def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False): + def _export( + self, tmp_path, head_type="vanilla", use_confidence_head=False, markov_rank=MARKOV_RANK + ): model = get_tiny_llama(num_hidden_layers=4) mtsp.convert( model, @@ -278,7 +281,9 @@ def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False): ( "dflash", _get_dspark_config( - head_type=head_type, use_confidence_head=use_confidence_head + head_type=head_type, + use_confidence_head=use_confidence_head, + markov_rank=markov_rank, ), ) ], @@ -312,8 +317,12 @@ def test_export_weight_keys_match_reference(self, tmp_path, head_type, use_confi @pytest.mark.parametrize("use_confidence_head", [False, True]) def test_export_config_has_dspark_fields(self, tmp_path, head_type, use_confidence_head): """Vanilla exports select the runtime and preserve the trained head settings.""" + # DSparkModule casts markov_rank to int; vLLM sizes its Markov table from the export. export_dir = self._export( - tmp_path, head_type=head_type, use_confidence_head=use_confidence_head + tmp_path, + head_type=head_type, + use_confidence_head=use_confidence_head, + markov_rank=float(MARKOV_RANK), ) with open(export_dir / "config.json") as f: cfg = json.load(f) @@ -323,6 +332,7 @@ def test_export_config_has_dspark_fields(self, tmp_path, head_type, use_confiden dc = cfg["dflash_config"] assert dc["projector_type"] == "dspark" assert dc["markov_rank"] == MARKOV_RANK + assert type(dc["markov_rank"]) is int assert dc["markov_head_type"] == head_type.lower() assert dc["use_confidence_head"] is use_confidence_head assert dc["shift_label"] is True @@ -331,6 +341,7 @@ def test_export_config_has_dspark_fields(self, tmp_path, head_type, use_confiden if vanilla: assert cfg["model_type"] == "qwen3" assert cfg["markov_rank"] == MARKOV_RANK + assert type(cfg["markov_rank"]) is int assert cfg["sample_from_anchor"] is True assert cfg["enable_confidence_head"] is use_confidence_head assert cfg["confidence_head_with_markov"] is use_confidence_head From 7ecbaafb4d70f9f87ac699ad329e78eee811d2a9 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:23:45 +0000 Subject: [PATCH 3/4] fix(speculative): refuse HF AR validation for Domino instead of drafting it one position off HFDominoModel inherited DFlash's pseudo_speculative_generate, which drafts block positions 1..bs-1 with the backbone alone. Domino trains position k on the token at anchor + k + 1 and corrects it with its GRU head, so every draft was scored one position off and without the head. It now raises NotImplementedError rather than gaining a decode of its own: HF AR validation is being retired in favour of measuring acceptance in a serving engine, and Domino's recipe already keeps ar_validate_steps at 0. The eval-time warning that called the reported acceptance rates backbone-only goes, along with the skill sheet's advice to treat it as expected. The sheet, the recipe comment and the launcher example now say AR validation is refused. Verified: a new test expects the refusal; without the override the inherited decode runs and the test fails. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/torch/speculative/plugins/hf_domino.py | 15 +++++++-------- .../general/speculative_decoding/domino.yaml | 3 +-- .../references/algorithms/domino.md | 9 ++++----- .../torch/speculative/plugins/test_hf_domino.py | 9 +++++++++ .../examples/Qwen/Qwen3-8B/hf_online_domino.yaml | 7 +++---- 5 files changed, 24 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index a0cda038c7a..128be406748 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -263,14 +263,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, @@ -362,6 +354,13 @@ def forward( return ModelOutput(loss=loss, logits=None, train_acc=[[accuracy]], domino_metrics=metrics) + def pseudo_speculative_generate(self, input_ids, steps=1): + """Refused: DFlash's decode would draft with the backbone alone, one block position off.""" + raise NotImplementedError( + "Domino has no HF AR validation (pseudo_speculative_generate). Set ar_validate_steps " + "to 0 and measure acceptance on the exported drafter in a serving engine." + ) + class DominoLambdaCallback(TrainerCallback): """Update the model's ``lambda_base`` from the HF Trainer global step. diff --git a/modelopt_recipes/general/speculative_decoding/domino.yaml b/modelopt_recipes/general/speculative_decoding/domino.yaml index e125381ce6e..0fe0ee3d060 100644 --- a/modelopt_recipes/general/speculative_decoding/domino.yaml +++ b/modelopt_recipes/general/speculative_decoding/domino.yaml @@ -37,9 +37,8 @@ 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 + # Keep off: Domino has no HF AR validation (pseudo_speculative_generate raises). ar_validate_steps: 0 answer_only_loss: true diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md index ca35742b2d9..1203cb33bf0 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md @@ -87,10 +87,9 @@ 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. +**Domino has no HF AR validation.** `pseudo_speculative_generate` raises for Domino, so +keep the recipe's `ar_validate_steps: 0` and measure acceptance on the exported drafter +in a serving engine. 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` @@ -110,6 +109,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 | -| `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 | +| `Domino has no HF AR validation (pseudo_speculative_generate)` | `ar_validate.py`, or `ar_validate_steps > 0` (training logs `AR validation not available.`) | Evaluate the exported drafter in a serving engine | | 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/tests/unit/torch/speculative/plugins/test_hf_domino.py b/tests/unit/torch/speculative/plugins/test_hf_domino.py index 0030abe6f33..3174b74e981 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_domino.py +++ b/tests/unit/torch/speculative/plugins/test_hf_domino.py @@ -24,6 +24,7 @@ import json from copy import deepcopy +import pytest import torch from _test_utils.torch.transformers_models import get_tiny_llama from safetensors.torch import load_file @@ -156,6 +157,14 @@ def test_lambda_zero_uses_final_only(self): out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) assert abs(out.loss.item() - out.domino_metrics["final_loss"]) < 1e-4 + def test_ar_validation_is_refused(self): + """DFlash's decode would draft Domino's block one position off and without its head.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", _get_domino_config())]) + input_ids, _, _ = self._make_batch(model.dflash_config.vocab_size) + with pytest.raises(NotImplementedError): + model.pseudo_speculative_generate(input_ids, steps=2) + class TestDominoSwa: """Domino honors dflash_swa_window_size (regression: the window was silently ignored).""" 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 f3095079a15..a29f57be028 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml @@ -10,10 +10,9 @@ # task_0: Build training conversations (Daring-Anteater multi-turn SFT, 50K) # task_1: Online Domino training + export of the drafter checkpoint # -# NOTE: the inference side (vLLM / AR evaluation) is intentionally not wired up -# yet — the Domino correction head is not applied in pseudo_speculative_generate -# or in the serving stack. Add the vLLM smoke-test / MT-Bench AR-eval steps -# (see hf_online_dflash.yaml task_1/task_2) once the inference path lands. +# NOTE: no evaluation step yet. Domino has no HF AR validation +# (pseudo_speculative_generate raises), and no serving stack applies its correction +# head. Add a vLLM smoke test (see hf_online_dflash.yaml task_1) once one does. # # Reference: SpecForge PR #571 (z-lab) | drafter format: # huggingface.co/Huang2020/Qwen3-8B-Domino-b16 From 63ab708f2652b380a381d3fc130e3de99c98550b Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:24:22 +0000 Subject: [PATCH 4/4] docs(speculative): note the DSpark and Domino fixes in the changelog Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d6397b993a0..308fe6f4e0f 100755 --- a/CHANGELOG.rst +++ b/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. +- DSpark and Domino now reject ``dflash_architecture_config.shift_label: false`` at conversion and export the ``shift_label``, ``markov_rank`` and other head fields they were trained with; checkpoints saved with ``shift_label: false`` still load. Domino's HF AR validation (``pseudo_speculative_generate``) now raises instead of scoring the backbone one block position off; evaluate Domino drafters in a serving engine. - 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.