From b6ebd4b453a5aab639ec09443b11ba8110ba1e36 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:22:45 +0000 Subject: [PATCH 1/9] fix(speculative): keep DSpark's AR-validation logits on the draft's device DSpark's pseudo_speculative_generate adds its Markov bias, computed on the draft's device, to the backbone logits that the base model's lm_head returns. scripts/ar_validate.py loads with device_map="auto", and accelerate keeps a tied lm_head with embed_tokens on the first GPU while the draft can land on a later one. The logits then stay on the first GPU, so the sum raises "Expected all tensors to be on the same device" for every sample and AR validation reports nothing. This has been the case since DSpark was added. The backbone logits now move to the draft's device before the loop. On a single device the move is a no-op. Verified: a new GPU test dispatches the model with accelerate's dispatch_model, embed_tokens and lm_head on the CPU and everything else on cuda:0, and generates with DFlash and DSpark. DSpark fails without the fix and passes with it. The bitwise A/B against #2757 shows DSpark's drafted tokens unchanged. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/speculative/plugins/hf_dspark.py | 3 +- .../speculative/plugins/test_hf_dflash.py | 39 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index b29d11584ce..bacabdf1790 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/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 diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 5f2e8adc045..85558ecb0a1 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -22,6 +22,7 @@ import torch from _test_utils.torch.speculative.dflash import get_dflash_config from _test_utils.torch.transformers_models import get_tiny_llama +from accelerate import dispatch_model import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.plugins.dflash_flex_attention import is_block_mask @@ -296,6 +297,44 @@ def test_generate_gathers_sharded_hidden_states(self): assert draft_tokens.shape == (1, 2) +class TestLmHeadOffTheDraftDevice: + """``device_map="auto"`` keeps a tied ``lm_head`` with ``embed_tokens`` on the first GPU. + + The draft can land on another GPU, so a decode that adds the draft's own outputs to the + backbone logits must bring those logits to the draft first. + """ + + @pytest.mark.parametrize( + "head_fields", + [{}, {"projector_type": "dspark", "markov_rank": 4}], + ids=["dflash", "dspark"], + ) + def test_generate(self, head_fields): + config = get_dflash_config() + config["dflash_architecture_config"].update(head_fields) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + device_map = { + "model.embed_tokens": "cpu", + "lm_head": "cpu", + "model.layers": 0, + "model.norm": 0, + "model.rotary_emb": 0, + "dflash_module": 0, + } + # main_device="cpu" runs the "cpu" entries there rather than offloading them to the + # GPU, so the CPU stands in for the first GPU. + dispatch_model(model.eval(), device_map, main_device="cpu") + + input_ids = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + base_token, draft_tokens = model.pseudo_speculative_generate(input_ids, steps=2) + + assert base_token.device == input_ids.device + assert draft_tokens.device == input_ids.device + assert draft_tokens.shape == (1, 2) + + class TestDFlashFlexAttentionGPU: """FlexAttention path must match the dense-mask SDPA path it replaces. From 45cc8260d7958dfb2e0b3de8c13ec253c44b9bf9 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:26:05 +0000 Subject: [PATCH 2/9] fix(speculative): draft Domino's AR-validation block with its alignment and head HFDominoModel did not override pseudo_speculative_generate, so AR validation ran DFlash's, which decodes block positions 1..bs-1 with the backbone alone. DFlash trains position k on the token at anchor + k; Domino trains it on anchor + k + 1 (_label_shift = 1). Every drafted token was therefore compared with the target one position too early, so drafts were almost never accepted, and the correction head that training optimizes and the exporter ships was never applied. domino.yaml keeps AR validation off, so shipped training did not hit this. DFlash's default decode now starts at block position 1 - _label_shift and returns at most bs - 1 + _label_shift tokens. For DFlash, DFlash2 and LiLiCorr (shift 0) that is the same slice as before. Domino overrides it the way DSpark does for its Markov head: it decodes the block left to right, stepping the GRU over the anchor and the tokens drafted so far, and from pure_draft_prefix_len on adds embed_proj([hidden_k; GRU state]) to the backbone logits, as _apply_domino_head does in training. As in DSpark's decode, the backbone logits and the drafted tokens' embeddings are moved to the draft's device: under device_map="auto" the embedding table, and a tied lm_head with it, can sit on another GPU. The eval forward warned that Domino's acceptance rates were backbone-only, but that method computes no acceptance rate, and AR validation now runs the head. The warning goes, and so do the other copies of the claim: the comments in domino.yaml, dspark.yaml and the Nemotron-3.5 DSpark warm-start recipe, the note in the Domino launcher example, and the speculative-decoding skill's Domino, DSpark and validation sheets. DSpark's generation has applied its Markov head since DSpark was added. The validation sheet keeps its rule against reporting such an AR as the model's result, now for the heads that AR validation still skips (LiLiCorr's reranker, DFlash2's selector), and adds the LiLiCorr and online DFlash2 examples to the ones it lists as shipping no benchmark task. The Domino sheet and example now say how to add the HF AR eval: DFlash's task_2 with --steps equal to the block size, since Domino drafts the whole block. The flag values do not change. Verified: a new test checks that Domino's greedy drafts equal the argmax of _apply_domino_head teacher-forced on them, with head weights large enough to overturn the backbone. It fails on the parent commit, and also when the head is skipped, its GRU state lags a step, or the prefix boundary is off by one. The GPU test that puts lm_head on another device now covers Domino as well, and fails without the move. Bitwise A/B against #2757: DFlash and DSpark loss, gradients and drafted tokens are unchanged; only Domino's drafted tokens change. 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 | 26 +++++++----- .../torch/speculative/plugins/hf_domino.py | 40 +++++++++++++++---- .../general/speculative_decoding/domino.yaml | 2 - .../general/speculative_decoding/dspark.yaml | 3 -- .../dspark_warmstart.yaml | 1 - .../references/algorithms/domino.md | 22 +++------- .../references/algorithms/dspark.md | 13 ++---- .../references/stages/validate.md | 13 +++--- .../speculative/plugins/test_hf_dflash.py | 8 +++- .../speculative/plugins/test_hf_domino.py | 36 +++++++++++++++++ .../Qwen/Qwen3-8B/hf_online_domino.yaml | 7 ++-- 11 files changed, 110 insertions(+), 61 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index d4d76d7e550..96d41932651 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -1219,36 +1219,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) diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index fc33a5155b6..31b9b4de392 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -266,14 +266,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 +357,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. diff --git a/modelopt_recipes/general/speculative_decoding/domino.yaml b/modelopt_recipes/general/speculative_decoding/domino.yaml index e125381ce6e..e7237b11b8f 100644 --- a/modelopt_recipes/general/speculative_decoding/domino.yaml +++ b/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 diff --git a/modelopt_recipes/general/speculative_decoding/dspark.yaml b/modelopt_recipes/general/speculative_decoding/dspark.yaml index a071b071372..db69f3c4b9e 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark.yaml +++ b/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 diff --git a/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml b/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml index eb78e80c466..d96780e05d2 100644 --- a/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/speculative_decoding/dspark_warmstart.yaml +++ b/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 diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md index 3f8434d1000..1b641e11927 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md +++ b/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 @@ -110,6 +101,5 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i | `Domino currently supports shift_label=True (next-token alignment) only` | `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 | diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md index f857d193d3f..793b5c0960c 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md @@ -117,15 +117,10 @@ the only in-pipeline evidence is training progress plus the export landing in ## Quality gate -**Do not trust in-training AR for DSpark.** The recipe pins `estimate_ar: false` and -`ar_validate_steps: 0` deliberately: eval runs the DFlash backbone only, with the -Markov head not applied, so any reported AR reflects the backbone alone rather than -the trained model. - -`pseudo_speculative_generate` *is* overridden for DSpark (unlike Domino), so a -non-offline model can generate correctly — but the offline/streaming path deletes base -layers and refuses. Evaluate by exporting and running the offline acceptance-length -harness separately. +`pseudo_speculative_generate` applies the Markov head, so AR validation measures the +trained drafter on an online model — but the offline/streaming path deletes base +layers and refuses. Evaluate those by exporting and running the offline +acceptance-length harness separately. Otherwise the training-regression gate from `dflash.md` (`MAX_FINAL_LOSS`, `MIN_FINAL_ACC` via `check_regression.py`) applies; none of the streaming examples set diff --git a/plugins/modelopt/skills/speculative-decoding/references/stages/validate.md b/plugins/modelopt/skills/speculative-decoding/references/stages/validate.md index 8bde138c1b0..be7bba191f7 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/stages/validate.md +++ b/plugins/modelopt/skills/speculative-decoding/references/stages/validate.md @@ -63,12 +63,13 @@ produces an in-pipeline metric**, so what you check depends on the sheet: metric, and no launcher example sets one. A green benchmark task therefore says the run finished, not that the metric passed. Reporting PASS without having read the number off the log is a false PASS on the run's headline metric. -- **Sheet defines no inference metric** (currently Domino and DSpark — their eval path - runs the DFlash backbone with the new head bypassed, and Domino ships no benchmark - task at all) — do **not** go looking for a benchmark log. Report the training - regression gate instead, and state plainly that acceptance quality requires a - separate evaluation of the exported checkpoint. Never report a backbone-only - acceptance rate as the model's result. +- **Sheet defines no inference metric** (currently the Domino, LiLiCorr and DSpark + training examples and the online DFlash2 one, which ship no benchmark task) — do + **not** go looking for a benchmark log. Report the training regression gate instead, + and state plainly that acceptance quality requires a separate evaluation of the + exported checkpoint. Never report an AR measured by HF AR validation + (`pseudo_speculative_generate`) as the model's result for a variant whose head it + skips (LiLiCorr's reranker, DFlash2's selector). Where the gate is the training regression check (`check_regression.py` against `trainer_state.json`), confirm the `=== Regression Check ===` block is actually diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 85558ecb0a1..a0fc7301f39 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -306,8 +306,12 @@ class TestLmHeadOffTheDraftDevice: @pytest.mark.parametrize( "head_fields", - [{}, {"projector_type": "dspark", "markov_rank": 4}], - ids=["dflash", "dspark"], + [ + {}, + {"projector_type": "dspark", "markov_rank": 4}, + {"projector_type": "domino", "emb_dim": 8, "gru_hidden_dim": 8}, + ], + ids=["dflash", "dspark", "domino"], ) def test_generate(self, head_fields): config = get_dflash_config() diff --git a/tests/unit/torch/speculative/plugins/test_hf_domino.py b/tests/unit/torch/speculative/plugins/test_hf_domino.py index 0030abe6f33..5f0fdf996d3 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 @@ -188,6 +189,41 @@ def spy(*args, **kwargs): assert windows == [6] +class TestDominoGenerate: + """AR validation drafts with the model that is trained and exported.""" + + @pytest.mark.parametrize("pure_draft_prefix_len", [0, 1]) + def test_drafts_are_the_training_head_argmax(self, pure_draft_prefix_len): + """Greedy drafts equal the training head's argmax when it is teacher-forced on them.""" + config = _get_domino_config() + config["dflash_architecture_config"]["pure_draft_prefix_len"] = pure_draft_prefix_len + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + model.float().eval() + head = model.dflash_module + torch.manual_seed(0) + with torch.no_grad(): + # Large enough to overturn the backbone's argmax, so a draft that skips the head + # or feeds it the wrong prefix disagrees. + for p in [*head.prefix_gru.parameters(), *head.embed_proj.parameters()]: + p.normal_() + input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN)) + + base_token, drafts = model.pseudo_speculative_generate(input_ids, steps=BLOCK_SIZE) + + with torch.no_grad(): + _, hidden_states = model._generation_base(input_ids) + hidden = model._draft_block(base_token, hidden_states) + final_logits = model._apply_domino_head( + hidden, + model._base_model_lm_head(hidden), + torch.cat([base_token, drafts], dim=1), # the anchor, then the drafts + torch.zeros(2, 1, dtype=torch.long), + 1, + ) + assert torch.equal(final_logits.argmax(dim=-1), drafts) + + class TestLambdaSchedule: """Test the lambda_base curriculum schedule.""" 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..7adf3400cd7 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: serving (vLLM) is not wired up for Domino yet, so there is no smoke test. +# HF AR validation applies the correction head: to measure AR, append +# hf_online_dflash.yaml's task_2 with --steps 16 (Domino drafts the whole block). # # Reference: SpecForge PR #571 (z-lab) | drafter format: # huggingface.co/Huang2020/Qwen3-8B-Domino-b16 From cf46bab2e5d231e3f3efcb79b61ee6ae55cb2b5b 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 3/9] 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> --- .../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 96d41932651..d1c5d36eb60 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/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 diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 31b9b4de392..8d739f31253 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/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.""" 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 1b641e11927..8494f6f4b76 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md @@ -98,7 +98,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 | | DDP error about unused parameters | `ddp_find_unused_parameters` turned off | Keep it `true` | diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md index 793b5c0960c..bb2c6992ac5 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md @@ -146,6 +146,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 d100b4a8f79..306eb2679ea 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -1529,6 +1529,36 @@ def test_block_targets_follow_the_label_shift(model_cls, expected): assert torch.equal(targets.mask, torch.tensor([mask], dtype=torch.float32)) +@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 456c4341f6dfad6db862be283317bd471f02d8cb 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 4/9] 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 2dabbe37d57..ad8db23babd 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 403439a11fdf1a5a80726c500266a1647ca4d361 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:37:54 +0000 Subject: [PATCH 5/9] docs(speculative): note the DSpark and Domino fixes in the changelog DSpark and Domino have shipped since 0.46.0, so these are fixes to released features, and two of them ask something of users: drop shift_label: false from DSpark configs, which now fail to convert, and re-export DSpark drafters whose export carried shift_label: false. The entry leaves out markov_head_type: #2902 already exports it as built. 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..815cae82aac 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. +- 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. From 1770131347b6ceecb584729ba7bc02eac8b6e360 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:13:54 +0000 Subject: [PATCH 6/9] refactor(speculative): hand the DFlash loss one context per training step HFDFlashModel.forward built the batch, ran the draft and computed the loss inline, and the variants built on it got at the loss's pieces piecemeal: DFlash2 asked _compute_loss for its unreduced terms (return_terms, _DFlashLossTerms) to recompose the block objective, and DFlash2 and LiLiCorr overrode forward() only to carry out the metrics their _compute_loss had left on the model (_selector_metrics, _lilicorr_metrics). The forward now hands the loss one object per step: - DraftBatch: the batch, its blocks and the draft's hidden states. - LossContext: the labels (_block_targets), the draft's logits (DraftLogits: the distribution the draft serves, and the backbone's before any head) and the position weights, with the per-token cross-entropy and the teacher's logits computed lazily, once. - _compute_loss(ctx) returns a LossOutput: the loss, the accuracy train_acc reports and detached 0-dim metrics, which the forward returns as draft_metrics. What a variant may change becomes a hook: _online_base_outputs (the base model's outputs), _draft_logits (a head on the backbone's logits), _position_weighting and _compute_loss. The loss mask comes from _build_loss_mask. The DSpark/Domino padding fix, a separate PR against main, adds the same helper, with the same text, at the same place (right after get_exporter), so a merge of the two keeps one copy of it. DFlash2 and LiLiCorr take the context: DFlash2's tv/lambda blend reads the cross-entropy and weights from it, and both return their metrics instead of overriding forward(). DSpark and Domino keep their own forwards for now. Verified against the previous commit with one train step and one generation step per case, in fp32 and bf16: DFlash (KD and CE), DFlash2 (ce/tv/lambda with the selector) and LiLiCorr, offline and online, D-PACE and decay, including zero-loss steps, give bitwise-equal loss, gradients, train_acc and generated tokens (32/32). The selector and LiLiCorr metrics keep their names and values under draft_metrics. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/torch/speculative/config.py | 2 +- .../torch/speculative/plugins/hf_dflash.py | 383 ++++++++++-------- .../torch/speculative/plugins/hf_dflash2.py | 92 ++--- .../torch/speculative/plugins/hf_lilicorr.py | 107 ++--- 4 files changed, 271 insertions(+), 313 deletions(-) diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index f49c38254d6..1a53b21e4d4 100644 --- a/modelopt/torch/speculative/config.py +++ b/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 " diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index d1c5d36eb60..f3d915be237 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -71,6 +71,7 @@ off meta; ``_maybe_init_rotary_emb`` stays idempotent for the lazy pattern. """ +import functools import logging from pathlib import Path from typing import NamedTuple @@ -296,19 +297,89 @@ def _dpace_position_weights( return weights.to(dtype=confidences.dtype) -class _DFlashLossTerms(NamedTuple): - """The unreduced pieces behind the block loss, for variants that re-derive it. +class DraftLogits(NamedTuple): + """The draft's logits per block position: [B, N, bs, V].""" - ``ce_per_token`` is ``None`` on the KD path, which never forms a per-position - cross-entropy. ``weights`` carries the position weighting (decay or D-PACE) and - normalizes by ``weight_sum``; ``supervised_mask`` is the unweighted mask the - reported accuracy uses. + final: torch.Tensor # the distribution the draft serves + backbone: torch.Tensor # before any head; ``final`` itself for a head-less variant + + +class LossOutput(NamedTuple): + """A training step's loss, the accuracy ``train_acc`` reports, and its metrics.""" + + loss: torch.Tensor + accuracy: float + metrics: dict # name -> detached 0-dim tensor + + +class LossContext: + """One training step's supervision; costly per-token quantities are computed lazily, once. + + ``weights`` are the per-position weights after decay or D-PACE, which reads the draft's + confidence off ``ce("final")``. ``ce`` keeps its gradient only when ``ce_grad`` says the + objective backpropagates through it; otherwise it runs under ``no_grad``. """ - ce_per_token: torch.Tensor | None - weights: torch.Tensor - weight_sum: torch.Tensor - supervised_mask: torch.Tensor + def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad): + """Weigh the step's block positions. + + Args: + batch: The forward's ``DraftBatch``. + logits: The draft's ``DraftLogits``. + targets: The labels and their mask, from ``_block_targets``. + weighting: ``_block_position_weights``' scheme, first, decay_factor and dpace_alpha. + teacher: Returns the base model's logits for every label position. + ce_grad: Whether the objective backpropagates through ``ce``. + """ + self.batch = batch + self.logits = logits + self.targets = targets + self._teacher_fn = teacher + self._ce_grad = ce_grad and torch.is_grad_enabled() + self._ce = {} + self._teacher = None + self.weights = _block_position_weights( + targets.mask, + **weighting, + target_ce=lambda: self.ce("final").detach().view(targets.mask.shape), + ) + self.weight_sum = self.weights.sum() + 1e-6 + + @functools.cached_property + def trainable(self) -> bool: + """Whether the weights leave anything to train on (reads the device once).""" + return bool(self.weight_sum > 1.0) + + def flat_logits(self, source): + """``logits.`` as [B * N * bs, V].""" + logits = getattr(self.logits, source) + return logits.reshape(-1, logits.shape[-1]) + + def ce(self, source): + """Per-token cross-entropy of ``logits.`` on the labels: [B * N * bs].""" + if source not in self._ce: + with torch.set_grad_enabled(self._ce_grad): + self._ce[source] = F.cross_entropy( + self.flat_logits(source), self.targets.token_ids.view(-1), reduction="none" + ) + return self._ce[source] + + def teacher(self): + """The base model's logits for every label position: [B, N, bs, V].""" + if self._teacher is None: + self._teacher = self._teacher_fn() + return self._teacher + + +class DraftBatch(NamedTuple): + """What a training step hands its loss: the batch, its blocks and the draft's output.""" + + input_ids: torch.Tensor # [B, seq] + loss_mask: torch.Tensor # [B, seq] 1.0 where a token may be a label + anchors: torch.Tensor # [B, N] each block's anchor position + keep: torch.Tensor # [B, N] whether each block is trained + base_outputs: DFlashBaseModelOutput + hidden: torch.Tensor # [B, N * bs, H] the draft's hidden states @DFlashDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) @@ -677,6 +748,24 @@ def get_exporter(self): return DFlashExporter(self) + @staticmethod + def _build_loss_mask(input_ids, attention_mask, labels, loss_mask): + """Return the float ``[B, seq_len]`` mask of positions anchor sampling and the loss may use. + + Every DFlash-family training forward calls this so the variants mask the same positions. + """ + mask = torch.ones(input_ids.shape, device=input_ids.device) + if labels is not None: + mask = mask * (labels != LabelSmoother.ignore_index).float() + # Labels carry optional answer-only masking but do not in general mark padding with -100 + # (collators copy them from padded input_ids), so always intersect with attention_mask. + if attention_mask is not None: + mask = mask * attention_mask.float() + # In offline training, assistant mask is dumped and passed as kwarg. + if loss_mask is not None: + mask = mask * loss_mask + return mask + def _sample_anchor_positions(self, seq_len, loss_mask, device): """Randomly sample anchor positions per sample. @@ -927,104 +1016,112 @@ def _teacher_logits(self, base_outputs, label_positions, token_ids=None): bias = getattr(lm_head, "bias", None) return logits if bias is None else logits + bias[token_ids] - def _compute_loss( + def _online_base_outputs( self, - logits, input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - *, - draft_hidden=None, - base_outputs=None, - return_terms=False, + attention_mask, + position_ids=None, + past_key_values=None, + inputs_embeds=None, + output_attentions=None, + cache_position=None, + **kwargs, ): - """Compute weighted cross-entropy (or KD) loss and accuracy. - - Args: - logits: Draft model output [B, N*block_size, vocab]. - input_ids: Original input token IDs [B, seq_len]. - anchor_positions: Anchor positions per block [B, N]. - block_keep_mask: Valid block mask [B, N]. - loss_mask: Token-level loss mask [B, seq_len]. - draft_hidden: Draft hidden states [B, N*block_size, H] behind ``logits``. - Unused here; passed for variants whose head consumes them. - base_outputs: The step's ``DFlashBaseModelOutput``: the KD teacher under - ``dflash_self_logit_distillation``, and read by variants. - return_terms: Also return the unreduced pieces behind the loss, so a variant - can recompose the block objective from a different divergence without - rebuilding the target alignment and position weighting. - TODO: promote this into a shared divergence seam when the DFlash-family - loss code is refactored; DFlash2 is the only consumer today. + """Run the frozen base model: the target hidden states, and its logits for KD.""" + # Multimodal models need the top-level conditional-generation forward so their + # image/video features are inserted before the language model runs. Keep the + # long-standing narrow call for text-only models. + base_forward_kwargs = _multimodal_forward_kwargs(kwargs) + with torch.no_grad(): + if base_forward_kwargs: + raw_outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=False, + output_attentions=output_attentions, + output_hidden_states=True, + cache_position=cache_position, + return_dict=True, + **base_forward_kwargs, + ) + else: + raw_outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) - Returns: - (loss, accuracy) tuple. - """ - block_size = self.dflash_block_size - targets = _block_targets( - input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift + if not getattr(raw_outputs, "hidden_states", None): + raise RuntimeError( + "The base model did not return hidden states required for DFlash training. " + "Ensure its top-level multimodal forward supports output_hidden_states=True." + ) + return DFlashBaseModelOutput( + target_hidden=self._target_hidden(raw_outputs.hidden_states), logits=raw_outputs.logits ) - binary_eval_mask = targets.mask.view(-1) - flat_logits = logits.view(-1, logits.size(-1)) - flat_targets = targets.token_ids.view(-1) + def _draft_logits(self, batch): + """The draft's logits per block position; a variant with a head applies it here.""" + logits = self._base_model_lm_head(batch.hidden) + logits = logits.view(*batch.anchors.shape, self.dflash_block_size, -1) + return DraftLogits(final=logits, backbone=logits) - kd = self.dflash_self_logit_distillation - if kd and base_outputs is None: - raise ValueError( - "dflash_self_logit_distillation distills from base_outputs, but none was " - "passed; an override of _compute_loss must forward it." - ) - # Non-KD loss is per-token cross-entropy; compute it once (grad enabled) so the - # D-PACE confidences below can reuse it instead of a second CE pass. The KD path - # optimizes KL, so its confidences need a dedicated no_grad CE pass. - loss_per_token = None - if not kd: - loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") - - weight_mask = _block_position_weights( - targets.mask, - self.dflash_loss_objective, - 1 - self._label_shift, - self.dflash_loss_decay_factor, - self.dflash_dpace_alpha, - target_ce=lambda: ( - loss_per_token.detach() - if loss_per_token is not None - else F.cross_entropy(flat_logits, flat_targets, reduction="none") - ).view(targets.mask.shape), + @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( + batch.input_ids, + batch.anchors, + batch.keep, + batch.loss_mask, + self.dflash_block_size, + self._label_shift, + ) + return LossContext( + batch, + self._draft_logits(batch), + targets, + weighting={ + "scheme": self._position_weighting, + "first": 1 - self._label_shift, + "decay_factor": self.dflash_loss_decay_factor, + "dpace_alpha": self.dflash_dpace_alpha, + }, + teacher=functools.partial(self._teacher_logits, batch.base_outputs, targets.positions), + # The KD path optimizes the soft target; its D-PACE confidences need no gradient. + ce_grad=not self.dflash_self_logit_distillation, ) - flat_weights = weight_mask.view(-1) - valid_count = flat_weights.sum() + 1e-6 - - if valid_count > 1.0: - if kd: - flat_teacher = self._teacher_logits(base_outputs, targets.positions) - flat_teacher = flat_teacher.reshape(flat_logits.shape) - target_soft = torch.softmax(flat_teacher, dim=-1) - draft_logsoft = torch.log_softmax(flat_logits, dim=-1) - kd_loss = -(target_soft * draft_logsoft).sum(dim=-1) - loss = (kd_loss * flat_weights).sum() / valid_count - else: - loss = (loss_per_token * flat_weights).sum() / valid_count + def _compute_loss(self, ctx): + """The weighted cross-entropy (or KD) loss on the served logits, and its accuracy. - with torch.no_grad(): - accuracy = _masked_top1(flat_logits, flat_targets, binary_eval_mask).item() - else: + A step whose position weights leave nothing to train returns a zero loss and + accuracy 0.0. + """ + final = ctx.logits.final + if not ctx.trainable: # Through the draft's output, so an FSDP unit around the draft still runs its backward. - loss = flat_logits.sum() * 0.0 - accuracy = 0.0 - - if return_terms: - terms = _DFlashLossTerms( - ce_per_token=loss_per_token, - weights=flat_weights, - weight_sum=valid_count, - supervised_mask=binary_eval_mask, - ) - return loss, accuracy, terms - return loss, accuracy + return LossOutput(final.sum() * 0.0, 0.0, {}) + + if self.dflash_self_logit_distillation: + flat_logits = ctx.flat_logits("final") + target_soft = torch.softmax(ctx.teacher().reshape(flat_logits.shape), dim=-1) + draft_logsoft = torch.log_softmax(flat_logits, dim=-1) + per_token = -(target_soft * draft_logsoft).sum(dim=-1) + else: + per_token = ctx.ce("final") + loss = (per_token * ctx.weights.view(-1)).sum() / ctx.weight_sum + + with torch.no_grad(): + accuracy = _masked_top1(final, ctx.targets.token_ids, ctx.targets.mask) + return LossOutput(loss, accuracy.item(), {}) def forward( self, @@ -1040,13 +1137,16 @@ def forward( cache_position=None, **kwargs, ): - """Training forward with random anchor sampling. + """Training forward with random anchor sampling, shared by every DFlash-family variant. - Random anchor sampling instead of uniform block division - Bidirectional intra-block attention (no causal constraint) - Context sees strictly before anchor position - - Label alignment: position k predicts token at anchor+k - - Optional loss decay weighting + - Label alignment: position k predicts the token at anchor + k + ``_label_shift`` + - 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``. """ if not self.training: if self.dflash_offline: @@ -1070,7 +1170,7 @@ def forward( **kwargs, ) - bsz, seq_len = input_ids.shape + seq_len = input_ids.shape[1] block_size = self.dflash_block_size device = input_ids.device @@ -1091,58 +1191,23 @@ def forward( and base_outputs.base_hidden is None ): raise KeyError("base_model_hidden_states") - target_hidden = base_outputs.target_hidden else: - # Multimodal models need the top-level conditional-generation forward so their - # image/video features are inserted before the language model runs. Keep the - # long-standing narrow call for text-only models. - base_forward_kwargs = _multimodal_forward_kwargs(kwargs) - use_top_level_forward = bool(base_forward_kwargs) - with torch.no_grad(): - if use_top_level_forward: - raw_outputs = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=False, - output_attentions=output_attentions, - output_hidden_states=True, - cache_position=cache_position, - return_dict=True, - **base_forward_kwargs, - ) - else: - raw_outputs = super().forward( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - ) - - if not getattr(raw_outputs, "hidden_states", None): - raise RuntimeError( - "The base model did not return hidden states required for DFlash training. " - "Ensure its top-level multimodal forward supports output_hidden_states=True." - ) - target_hidden = self._target_hidden(raw_outputs.hidden_states) - base_outputs = DFlashBaseModelOutput( - target_hidden=target_hidden, logits=raw_outputs.logits + base_outputs = self._online_base_outputs( + input_ids, + attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + cache_position=cache_position, + **kwargs, ) + target_hidden = base_outputs.target_hidden - # 2. Build loss mask. Labels carry optional answer-only masking, but do - # not in general mark padded tokens with -100 (the VLM collator creates - # them from padded input_ids). Always intersect with attention_mask so - # anchor sampling and loss never include the padded tail. - loss_mask = torch.ones(bsz, seq_len, device=device) - if labels is not None: - loss_mask = loss_mask * (labels != LabelSmoother.ignore_index).float() - if attention_mask is not None: - loss_mask = loss_mask * attention_mask.float() - - # In offline training, assistant mask is dumped and passed as kwarg. - if kwargs.get("loss_mask") is not None: - loss_mask = loss_mask * kwargs["loss_mask"] + # 2. Build loss mask. + loss_mask = self._build_loss_mask( + input_ids, attention_mask, labels, kwargs.get("loss_mask") + ) # 3. Random anchor sampling anchor_positions, block_keep_mask = self._sample_anchor_positions( @@ -1180,21 +1245,15 @@ def forward( ) # 6. Compute loss and accuracy - logits = self._base_model_lm_head(hidden) - loss, accuracy = self._compute_loss( - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - draft_hidden=hidden, - base_outputs=base_outputs, + batch = DraftBatch( + input_ids, loss_mask, anchor_positions, block_keep_mask, base_outputs, hidden ) - + out = self._compute_loss(self._loss_context(batch)) return ModelOutput( - loss=loss, + loss=out.loss, logits=base_outputs.logits, - train_acc=[[accuracy]], + train_acc=[[out.accuracy]], + draft_metrics=out.metrics, ) @torch.no_grad() diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py index 6e4171c9210..ebba52984a1 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash2.py +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -39,7 +39,7 @@ from transformers import PreTrainedModel from ..dflash.conversion import DFlash2DMRegistry -from .hf_dflash import HFDFlashModel, _block_targets, _block_tokens +from .hf_dflash import HFDFlashModel, LossOutput, _block_tokens from .modeling_dflash2 import DFlash2Module __all__ = ["HFDFlash2Model"] @@ -83,22 +83,6 @@ def modify(self, config): "optimizes a soft target instead. Set dflash_self_logit_distillation=false, " "or dflash_lk_loss_type='ce' to keep distillation." ) - self._selector_metrics = None - - def forward(self, *args, **kwargs): - """Run the DFlash training forward and attach the candidate-selector metrics. - - The variant is the objective, so the pipeline above it is inherited verbatim; - this only carries out what ``_compute_loss`` produced. Without it - ``selector_coverage`` -- the only signal distinguishing a selector that is - choosing from one being handed the gold token -- never reaches the logs. - """ - self._selector_metrics = None - outputs = super().forward(*args, **kwargs) - if self._selector_metrics is not None: - outputs["selector_metrics"] = self._selector_metrics - self._selector_metrics = None - return outputs def get_exporter(self): """Get the exporter for the DFlash2 draft model.""" @@ -163,7 +147,7 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas # every training step. The trainer converts them at the logging boundary. return loss, accuracy.detach(), coverage.detach() - def _lk_loss(self, terms): + def _lk_loss(self, ctx): """Re-weight the block objective between cross-entropy and acceptance. Both terms are read off the same per-position cross-entropy the backbone loss @@ -181,66 +165,34 @@ def _lk_loss(self, terms): The blend weight is detached, so it reshapes the objective without adding a gradient path of its own. """ - ce, weights, weight_sum = terms.ce_per_token, terms.weights, terms.weight_sum - assert ce is not None, ( - "the KD path produced no per-position cross-entropy, so q(gold) is " - "unavailable and the blend would silently fall back to it; modify() is " - "supposed to have rejected this combination at convert time" - ) + ce, weights, weight_sum = ctx.ce("final"), ctx.weights.view(-1), ctx.weight_sum gold_probability = torch.exp(-ce) tv_loss = ((1.0 - gold_probability) * weights).sum() / weight_sum if self.dflash_lk_loss_type == "tv": return tv_loss ce_loss = (ce * weights).sum() / weight_sum - mask = terms.supervised_mask + mask = ctx.targets.mask.view(-1) acceptance = (gold_probability.detach() * mask).sum() / (mask.sum() + 1e-6) ce_share = self.dflash_lk_ce_scale * torch.exp(-self.dflash_lk_ce_decay * acceptance) return ce_share * ce_loss + (1.0 - ce_share) * tv_loss - def _compute_loss( - self, - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - *, - draft_hidden=None, - base_outputs=None, - ): + def _compute_loss(self, ctx): """Backbone DFlash loss plus the candidate-selector cross-entropy. - Reuses ``HFDFlashModel._compute_loss`` for the backbone term, then rebuilds - the same target/weight alignment for the selector term. Reported accuracy - stays the backbone's top-1, so DFlash and DFlash2 runs remain comparable; - the selector's own accuracy is logged separately. + Reported accuracy stays the backbone's top-1, so DFlash and DFlash2 runs remain + comparable; the selector's own accuracy and coverage are added to the metrics. """ - loss, accuracy, terms = super()._compute_loss( - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - draft_hidden=draft_hidden, - base_outputs=base_outputs, - return_terms=True, - ) + out = super()._compute_loss(ctx) + loss = out.loss if self.dflash_lk_loss_type != "ce": - loss = self._lk_loss(terms) - if self.dflash_selector_loss_alpha <= 0 or draft_hidden is None: - return loss, accuracy + loss = self._lk_loss(ctx) + if self.dflash_selector_loss_alpha <= 0: + return out._replace(loss=loss) - bsz = input_ids.shape[0] + batch, targets = ctx.batch, ctx.targets + bsz, n_blocks = batch.anchors.shape block_size = self.dflash_block_size - n_blocks = anchor_positions.shape[1] - - # The backbone's supervision mask, without its position weighting (decay/D-PACE): - # that shapes *where* the backbone spends capacity, while the selector should learn - # every position's transition equally. - targets = _block_targets( - input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift - ) # Teacher-forced predecessor of block position k is the real token at anchor+k-1. # Only offsets 1..block_size-1 are supervised (the mask zeroes slot 0), so the @@ -251,18 +203,24 @@ def _compute_loss( # from the anchor token, so ITS position 0 corresponds to block offset 1. A caller # that hands greedy_path the full 0..block_size-1 candidate set is off by one. predecessor_ids = _block_tokens( - input_ids, anchor_positions, block_size, self._label_shift - 1 + batch.input_ids, batch.anchors, block_size, self._label_shift - 1 ) + # The backbone's supervision mask, without its position weighting (decay/D-PACE): + # that shapes *where* the backbone spends capacity, while the selector should learn + # every position's transition equally. selector_loss, selector_accuracy, selector_coverage = self._selector_loss( - logits.reshape(bsz, n_blocks, block_size, -1), + ctx.logits.backbone, targets.token_ids, - draft_hidden.reshape(bsz, n_blocks, block_size, -1), + batch.hidden.reshape(bsz, n_blocks, block_size, -1), predecessor_ids, targets.mask, ) - self._selector_metrics = { + metrics = { + **out.metrics, "selector_accuracy": selector_accuracy, "selector_coverage": selector_coverage, } - return loss + self.dflash_selector_loss_alpha * selector_loss, accuracy + return LossOutput( + loss + self.dflash_selector_loss_alpha * selector_loss, out.accuracy, metrics + ) diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py index 908881f8c20..10f98f30824 100644 --- a/modelopt/torch/speculative/plugins/hf_lilicorr.py +++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py @@ -77,7 +77,7 @@ from transformers import PreTrainedModel from ..dflash.conversion import LiLiCorrDMRegistry -from .hf_dflash import HFDFlashModel, _block_positions, _block_targets +from .hf_dflash import HFDFlashModel, LossOutput, _block_positions from .modeling_lilicorr import LiLiCorrModule logger = logging.getLogger(__name__) @@ -173,7 +173,6 @@ def modify(self, config): self.dflash_lilicorr_w_cal, self.dflash_lilicorr_margin, ) - self._lilicorr_metrics = None @property def _needs_base_logits(self) -> bool: @@ -190,77 +189,33 @@ def get_exporter(self): return LiLiCorrExporter(self) - def _compute_loss( - self, - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - *, - draft_hidden=None, - base_outputs=None, - ): + def _compute_loss(self, ctx): """Add the LiLiCorr lattice terms to the DFlash block loss. The returned accuracy is the reranker's greedy path-prefix fraction — the quantity that tracks acceptance length — rather than the backbone's per-token accuracy, which is reported as ``origin_accuracy`` in the metrics instead. - - The two things this needs beyond the shared signature -- the target-layer hidden - states it anchors on, and the target distribution its penalty and calibration - terms read -- both come from ``base_outputs``, the container - ``HFDFlashModel.forward`` already builds for them. """ - target_hidden = base_outputs.target_hidden if base_outputs is not None else None - if draft_hidden is None or target_hidden is None: - raise ValueError( - "LiLiCorr requires draft_hidden and target_hidden in _compute_loss: the " - "reranker reads the draft's per-slot hidden states and anchors on a " - "committed target row." - ) - - origin_loss, origin_accuracy = super()._compute_loss( - logits, - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - base_outputs=base_outputs, - ) + out = super()._compute_loss(ctx) + batch, targets = ctx.batch, ctx.targets # The backbone's own supervision, before its position weighting, so both losses # train the same positions. - targets = _block_targets( - input_ids, - anchor_positions, - block_keep_mask, - loss_mask, - self.dflash_block_size, - self._label_shift, - ) lilicorr_loss, accuracy, metrics = self._compute_lilicorr_loss( - logits=logits, + logits=ctx.logits.final, target_ids=targets.token_ids, # Block position 0 is the anchor -- the one row whose input embedding is a real # verified token -- so the lattice covers positions 1..bs-1. slot_mask=targets.mask[:, :, 1:] > 0.5, - anchor_positions=anchor_positions, - draft_hidden=draft_hidden, - target_hidden=target_hidden, - base_outputs=base_outputs, + anchor_positions=batch.anchors, + draft_hidden=batch.hidden, + target_hidden=batch.base_outputs.target_hidden, + base_outputs=batch.base_outputs, ) # The identity `loss == origin_loss + lilicorr_loss` is the cheap parity check # on this objective, so both halves are reported next to the total. - # Batched into one device-to-host sync rather than two per-step float() calls. - origin_scalars = torch.stack( - [ - origin_loss.detach().reshape(()).float(), - torch.as_tensor(origin_accuracy, dtype=torch.float32, device=origin_loss.device), - ] - ).tolist() - metrics["origin_loss"], metrics["origin_accuracy"] = origin_scalars - self._lilicorr_metrics = metrics - return origin_loss + lilicorr_loss, accuracy + metrics["origin_loss"] = out.loss.detach().float() + metrics["origin_accuracy"] = torch.full((), out.accuracy, device=out.loss.device) + return LossOutput(out.loss + lilicorr_loss, accuracy, metrics) def _compute_lilicorr_loss( self, @@ -277,7 +232,7 @@ def _compute_lilicorr_loss( Returns ``(loss, accuracy, metrics)``. ``loss`` already carries the absolute term weights, ``accuracy`` is the greedy path-prefix fraction, and ``metrics`` - holds Python floats read back in a single device sync. + holds detached 0-dim tensors. """ module = self.dflash_module head = module.lilicorr @@ -448,17 +403,18 @@ def _compute_lilicorr_loss( raw_metrics["lilicorr_penalty"] = penalty_loss raw_metrics["lilicorr_calibration"] = cal_loss - # One device sync for every scalar, rather than one per metric. - names = list(raw_metrics) - stacked = torch.stack([raw_metrics[name].detach().reshape(()).float() for name in names]) - metrics = dict(zip(names, stacked.tolist(), strict=True)) + metrics = {name: value.detach().reshape(()).float() for name, value in raw_metrics.items()} # Weights recorded alongside the terms so a run states its own composition - # without the reader expanding the config. - metrics["lilicorr_w_ce"] = self.dflash_lilicorr_w_ce - metrics["lilicorr_w_margin"] = self.dflash_lilicorr_w_margin - metrics["lilicorr_w_pen"] = self.dflash_lilicorr_w_pen - metrics["lilicorr_w_cal"] = self.dflash_lilicorr_w_cal - accuracy = metrics["lilicorr_selected_prefix"] / float(num_slots) + # without the reader expanding the config. float64 holds the configured values exactly. + weights = { + "lilicorr_w_ce": self.dflash_lilicorr_w_ce, + "lilicorr_w_margin": self.dflash_lilicorr_w_margin, + "lilicorr_w_pen": self.dflash_lilicorr_w_pen, + "lilicorr_w_cal": self.dflash_lilicorr_w_cal, + } + for name, weight in weights.items(): + metrics[name] = torch.full((), weight, dtype=torch.float64, device=loss.device) + accuracy = metrics["lilicorr_selected_prefix"].item() / float(num_slots) return loss, accuracy, metrics def _distractor_penalty( @@ -676,18 +632,3 @@ def block_mean(values): / supervised_count, } return metrics - - def forward(self, *args, **kwargs): - """Run the DFlash training forward and attach the LiLiCorr metrics. - - The whole variant is the objective, so the pipeline above it is inherited - verbatim rather than reimplemented; this only carries the metrics that - ``_compute_loss`` produced out to the caller, the way DSpark and Domino do - from their own forwards. - """ - self._lilicorr_metrics = None - outputs = super().forward(*args, **kwargs) - if self._lilicorr_metrics is not None: - outputs["lilicorr_metrics"] = self._lilicorr_metrics - self._lilicorr_metrics = None - return outputs From 3127fc7738cd7351165307ba89f085bf43455778 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:18:50 +0000 Subject: [PATCH 7/9] refactor(speculative): move the DFlash loss components into dflash_losses.py The block-supervision helpers (_block_positions, _block_gather, _block_tokens, _BlockTargets, _block_targets), the position weighting (_dpace_position_weights, _block_position_weights), _masked_top1, the loss context (DraftLogits, LossOutput, LossContext) and DSpark's per-token TVD (_tvd_chunk, _tvd_per_token) move unchanged from hf_dflash.py and hf_dspark.py into one module that names no variant. The next PR in this stack builds the shared loss terms on them, and hf_dspark.py would otherwise be imported by the generic loss code for its TVD. A pure move: every moved definition is byte-identical to its previous text, and only import lines change elsewhere. The new file carries the SpecForge and DeepSpec notices of the two files the code came from and is excluded from the license-insertion hook like them. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .pre-commit-config.yaml | 1 + .../speculative/plugins/dflash_losses.py | 279 ++++++++++++++++++ .../torch/speculative/plugins/hf_dflash.py | 197 +------------ .../torch/speculative/plugins/hf_dflash2.py | 3 +- .../torch/speculative/plugins/hf_domino.py | 9 +- .../torch/speculative/plugins/hf_dspark.py | 38 +-- .../torch/speculative/plugins/hf_lilicorr.py | 3 +- 7 files changed, 291 insertions(+), 239 deletions(-) create mode 100644 modelopt/torch/speculative/plugins/dflash_losses.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f60fad5cc15..88d7a89466d 100644 --- a/.pre-commit-config.yaml +++ b/.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| diff --git a/modelopt/torch/speculative/plugins/dflash_losses.py b/modelopt/torch/speculative/plugins/dflash_losses.py new file mode 100644 index 00000000000..b2e8279c3c3 --- /dev/null +++ b/modelopt/torch/speculative/plugins/dflash_losses.py @@ -0,0 +1,279 @@ +# Adapted from https://github.com/sgl-project/SpecForge/blob/8ea5ca6/specforge/core/dflash.py +# and https://github.com/deepseek-ai/DeepSpec/blob/add63ba/deepspec/modeling/dspark/loss.py +# Copyright (c) 2025 sgl-project +# Copyright (c) 2026 The DeepSpec Authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Loss components of the DFlash speculative-decoding family, independent of any variant. + +A training step's supervision: which token each block position is trained on and where +(:func:`_block_targets`), how the positions are weighted (:func:`_block_position_weights`), +and :class:`LossContext`, which carries them with the draft's logits to the loss. +""" + +import functools +from typing import NamedTuple + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint + +__all__ = ["DraftLogits", "LossContext", "LossOutput"] + + +def _block_positions(anchor_positions, block_size, offset, seq_len): + """Index ``anchor + offset + k`` of each block position ``k``, clamped onto the sequence. + + Returns ``(positions [B, N, bs], inside)``; ``inside`` is False past the sequence end. + """ + k = torch.arange(block_size, device=anchor_positions.device).view(1, 1, -1) + unclamped = anchor_positions.unsqueeze(-1) + k + offset + return unclamped.clamp(min=0, max=seq_len - 1), unclamped < seq_len + + +def _block_gather(x, positions): + """``x[b, positions[b, n, k]]`` for ``x`` [B, seq] and ``positions`` [B, N, bs].""" + return torch.gather(x.unsqueeze(1).expand(-1, positions.shape[1], -1), 2, positions) + + +def _block_tokens(input_ids, anchor_positions, block_size, offset): + """The token at ``anchor + offset + k`` for each block position ``k``: [B, N, bs].""" + positions, _ = _block_positions(anchor_positions, block_size, offset, input_ids.shape[1]) + return _block_gather(input_ids, positions) + + +class _BlockTargets(NamedTuple): + """What each block position is trained on: the token at ``anchor + k + label_shift``.""" + + token_ids: torch.Tensor # [B, N, bs] + positions: torch.Tensor # [B, N, bs] their sequence indices, clamped onto the sequence + mask: torch.Tensor # [B, N, bs] 1.0 where the position is supervised + + +def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, shift): + """Labels of every block position and where they are supervised. + + A position is supervised when its block is kept and its label lies inside the sequence + and the loss mask, and is not the anchor itself, which the draft is given as input. + """ + positions, inside = _block_positions(anchor_positions, block_size, shift, input_ids.shape[1]) + k = torch.arange(block_size, device=input_ids.device).view(1, 1, -1) + mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() + mask = mask * inside.float() + mask = mask * (k + shift > 0).float() + mask = mask * _block_gather(loss_mask, positions) + return _BlockTargets(_block_gather(input_ids, positions), positions, mask) + + +def _dpace_position_weights( + confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None +) -> torch.Tensor: + """Compute detached D-PACE per-position weights from draft confidences. + + Derived from D-PACE (arXiv:2605.18810). The paper factorizes the per-position + weight (Fig. 2 / Eq. 8) into a *cumulative confidence* times a *continuation + value*, which is equivalently the suffix sum of the cumulative confidences:: + + C_j = prod_{i<=j} q~_i # cumulative confidence (Eq. 8) + w_j = sum_{m>=j} C_m # = C_j * continuation value f~_j + + Each confidence is asymmetrically smoothed toward 1 (Eq. 7):: + + q~_i = (1 - alpha) * q_i + alpha, alpha in (0, 1], + + so the floor ``q~_i >= alpha`` keeps every cumulative product (hence every + weight) strictly positive. We evaluate the suffix sum from its definition as + ``total - exclusive_prefix_sum`` of ``C`` rather than reversing the tensor. + Positions with ``valid_mask == 0`` are multiplicative no-ops in ``C`` and + contribute nothing to the sum, matching the per-token loss mask. Weights are + detached (Eq. 9): they reweight the cross-entropy without adding gradient. + + Args: + confidences: ``[..., L]`` draft confidence ``q_i = exp(-CE_i)`` per position. + alpha: smoothing factor in (0, 1]; raises if outside that range. + valid_mask: optional ``[..., L]`` 0/1 mask; ``None`` treats all positions valid. + + Returns: + Detached weights with the same shape and dtype as ``confidences``. + """ + if not 0.0 < alpha <= 1.0: + raise ValueError(f"dflash_dpace_alpha must be in (0, 1], got {alpha}") + + with torch.no_grad(): + smoothed = alpha + (1.0 - alpha) * confidences.float() # Eq. 7 + if valid_mask is not None: + keep = valid_mask.to(torch.bool) + smoothed = torch.where(keep, smoothed, torch.ones_like(smoothed)) + cum_conf = torch.cumprod(smoothed, dim=-1) # Eq. 8 cumulative confidence C_j + if valid_mask is not None: + cum_conf = cum_conf * keep.to(cum_conf.dtype) + # Suffix sum w_j = sum_{m>=j} C_m, written as total minus the exclusive + # prefix sum so no axis reversal is needed (Eq. 8). + inclusive = torch.cumsum(cum_conf, dim=-1) + weights = inclusive[..., -1:] - inclusive + cum_conf + return weights.to(dtype=confidences.dtype) + + +def _block_position_weights(mask, scheme, first, decay_factor, dpace_alpha, target_ce=None): + """``mask`` weighted per block position, by D-PACE or by exponential decay. + + ``first`` is the first predicted block position. ``target_ce`` returns the draft's + cross-entropy on its labels ([B, N, bs]); only D-PACE calls it. + """ + block_size = mask.shape[-1] + if scheme == "dpace" and block_size > 1: + with torch.no_grad(): + # Draft confidence q_i = exp(-CE) on the label; detached, as in the paper (Eq. 9). + confidences = torch.exp(-target_ce()[..., first:].float()) + dpace = torch.ones_like(mask) + dpace[..., first:] = _dpace_position_weights( + confidences, dpace_alpha, valid_mask=mask[..., first:] + ) + return mask * dpace + if decay_factor > 0: + k = torch.arange(block_size, device=mask.device).view(1, 1, -1) + return mask * torch.exp(-(k - first).clamp(min=0).float() / decay_factor) + return mask + + +def _masked_top1(logits, target_ids, mask): + """Top-1 accuracy of ``logits`` against ``target_ids`` where ``mask`` is set (0-dim).""" + correct = (logits.argmax(dim=-1) == target_ids) & (mask > 0.5) + return correct.sum().float() / (mask.sum() + 1e-6) + + +def _tvd_chunk(a, b): + """Per-token TVD for one row chunk: ``(softmax(a) - softmax(b)).abs().sum(-1)``.""" + return (torch.softmax(a.float(), dim=-1) - torch.softmax(b.float(), dim=-1)).abs().sum(dim=-1) + + +def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024, chunk_fn=None): + """Total-variation distance ||softmax(a)-softmax(b)||_1 / ... per token, memory-lean. + + Materializing both [N, vocab] float32 softmax tensors at once OOMs at large + N (num_anchors*block_size) and a 150k vocab. We compute it in row chunks and + gradient-checkpoint each chunk so the wide softmaxes are recomputed in backward + rather than held — peak memory ~ chunk_size*vocab instead of N*vocab. The math + is identical to ``(softmax(final)-softmax(teacher)).abs().sum(-1)``. + + Chunks with ``Tensor.split``, not slicing (each slice's backward zero-fills a full + [N, vocab] tensor) nor ``torch.chunk`` (different chunk shapes, so recompiles). + """ + _chunk = chunk_fn or _tvd_chunk + outs = [] + for a, b in zip( + final_logits.split(chunk_size, dim=0), + teacher_logits.split(chunk_size, dim=0), + strict=True, + ): + if torch.is_grad_enabled() and a.requires_grad: + outs.append(torch.utils.checkpoint.checkpoint(_chunk, a, b, use_reentrant=False)) + else: + outs.append(_chunk(a, b)) + return torch.cat(outs, dim=0) + + +class DraftLogits(NamedTuple): + """The draft's logits per block position: [B, N, bs, V].""" + + final: torch.Tensor # the distribution the draft serves + backbone: torch.Tensor # before any head; ``final`` itself for a head-less variant + + +class LossOutput(NamedTuple): + """A training step's loss, the accuracy ``train_acc`` reports, and its metrics.""" + + loss: torch.Tensor + accuracy: float + metrics: dict # name -> detached 0-dim tensor + + +class LossContext: + """One training step's supervision; costly per-token quantities are computed lazily, once. + + ``weights`` are the per-position weights after decay or D-PACE, which reads the draft's + confidence off ``ce("final")``. ``ce`` keeps its gradient only when ``ce_grad`` says the + objective backpropagates through it; otherwise it runs under ``no_grad``. + """ + + def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad): + """Weigh the step's block positions. + + Args: + batch: The forward's ``DraftBatch``. + logits: The draft's ``DraftLogits``. + targets: The labels and their mask, from ``_block_targets``. + weighting: ``_block_position_weights``' scheme, first, decay_factor and dpace_alpha. + teacher: Returns the base model's logits for every label position. + ce_grad: Whether the objective backpropagates through ``ce``. + """ + self.batch = batch + self.logits = logits + self.targets = targets + self._teacher_fn = teacher + self._ce_grad = ce_grad and torch.is_grad_enabled() + self._ce = {} + self._teacher = None + self.weights = _block_position_weights( + targets.mask, + **weighting, + target_ce=lambda: self.ce("final").detach().view(targets.mask.shape), + ) + self.weight_sum = self.weights.sum() + 1e-6 + + @functools.cached_property + def trainable(self) -> bool: + """Whether the weights leave anything to train on (reads the device once).""" + return bool(self.weight_sum > 1.0) + + def flat_logits(self, source): + """``logits.`` as [B * N * bs, V].""" + logits = getattr(self.logits, source) + return logits.reshape(-1, logits.shape[-1]) + + def ce(self, source): + """Per-token cross-entropy of ``logits.`` on the labels: [B * N * bs].""" + if source not in self._ce: + with torch.set_grad_enabled(self._ce_grad): + self._ce[source] = F.cross_entropy( + self.flat_logits(source), self.targets.token_ids.view(-1), reduction="none" + ) + return self._ce[source] + + def teacher(self): + """The base model's logits for every label position: [B, N, bs, V].""" + if self._teacher is None: + self._teacher = self._teacher_fn() + return self._teacher diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index f3d915be237..93965b2b6e4 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -77,7 +77,6 @@ from typing import NamedTuple import torch -import torch.nn.functional as F from transformers import PreTrainedModel from transformers.models.qwen3.configuration_qwen3 import Qwen3Config as _Qwen3Config from transformers.trainer_pt_utils import LabelSmoother @@ -86,6 +85,7 @@ from ..dflash.conversion import DFlashDMRegistry from ..dflash.dflash_model import DFlashModel from .dflash_flex_attention import build_draft_block_mask +from .dflash_losses import DraftLogits, LossContext, LossOutput, _block_targets, _masked_top1 from .modeling_dflash import ( # noqa: F401 DFlashAttention, DFlashBaseModelOutput, @@ -134,78 +134,6 @@ def _gather_rows(x, positions): return torch.gather(x, 1, idx).reshape(*positions.shape, x.size(-1)) -def _block_positions(anchor_positions, block_size, offset, seq_len): - """Index ``anchor + offset + k`` of each block position ``k``, clamped onto the sequence. - - Returns ``(positions [B, N, bs], inside)``; ``inside`` is False past the sequence end. - """ - k = torch.arange(block_size, device=anchor_positions.device).view(1, 1, -1) - unclamped = anchor_positions.unsqueeze(-1) + k + offset - return unclamped.clamp(min=0, max=seq_len - 1), unclamped < seq_len - - -def _block_gather(x, positions): - """``x[b, positions[b, n, k]]`` for ``x`` [B, seq] and ``positions`` [B, N, bs].""" - return torch.gather(x.unsqueeze(1).expand(-1, positions.shape[1], -1), 2, positions) - - -def _block_tokens(input_ids, anchor_positions, block_size, offset): - """The token at ``anchor + offset + k`` for each block position ``k``: [B, N, bs].""" - positions, _ = _block_positions(anchor_positions, block_size, offset, input_ids.shape[1]) - return _block_gather(input_ids, positions) - - -class _BlockTargets(NamedTuple): - """What each block position is trained on: the token at ``anchor + k + label_shift``.""" - - token_ids: torch.Tensor # [B, N, bs] - positions: torch.Tensor # [B, N, bs] their sequence indices, clamped onto the sequence - mask: torch.Tensor # [B, N, bs] 1.0 where the position is supervised - - -def _block_targets(input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, shift): - """Labels of every block position and where they are supervised. - - A position is supervised when its block is kept and its label lies inside the sequence - and the loss mask, and is not the anchor itself, which the draft is given as input. - """ - positions, inside = _block_positions(anchor_positions, block_size, shift, input_ids.shape[1]) - k = torch.arange(block_size, device=input_ids.device).view(1, 1, -1) - mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float() - mask = mask * inside.float() - mask = mask * (k + shift > 0).float() - mask = mask * _block_gather(loss_mask, positions) - return _BlockTargets(_block_gather(input_ids, positions), positions, mask) - - -def _block_position_weights(mask, scheme, first, decay_factor, dpace_alpha, target_ce=None): - """``mask`` weighted per block position, by D-PACE or by exponential decay. - - ``first`` is the first predicted block position. ``target_ce`` returns the draft's - cross-entropy on its labels ([B, N, bs]); only D-PACE calls it. - """ - block_size = mask.shape[-1] - if scheme == "dpace" and block_size > 1: - with torch.no_grad(): - # Draft confidence q_i = exp(-CE) on the label; detached, as in the paper (Eq. 9). - confidences = torch.exp(-target_ce()[..., first:].float()) - dpace = torch.ones_like(mask) - dpace[..., first:] = _dpace_position_weights( - confidences, dpace_alpha, valid_mask=mask[..., first:] - ) - return mask * dpace - if decay_factor > 0: - k = torch.arange(block_size, device=mask.device).view(1, 1, -1) - return mask * torch.exp(-(k - first).clamp(min=0).float() / decay_factor) - return mask - - -def _masked_top1(logits, target_ids, mask): - """Top-1 accuracy of ``logits`` against ``target_ids`` where ``mask`` is set (0-dim).""" - correct = (logits.argmax(dim=-1) == target_ids) & (mask > 0.5) - return correct.sum().float() / (mask.sum() + 1e-6) - - def _draft_mask_mod(seq_len, anchor_positions, block_keep_mask, block_size, window, causal): """The draft's attention visibility rule, as a FlexAttention ``mask_mod``. @@ -248,129 +176,6 @@ def _dense_additive_mask(mask_mod, bsz, q_len, kv_len, dtype, device): return attn_mask -def _dpace_position_weights( - confidences: torch.Tensor, alpha: float, valid_mask: torch.Tensor | None = None -) -> torch.Tensor: - """Compute detached D-PACE per-position weights from draft confidences. - - Derived from D-PACE (arXiv:2605.18810). The paper factorizes the per-position - weight (Fig. 2 / Eq. 8) into a *cumulative confidence* times a *continuation - value*, which is equivalently the suffix sum of the cumulative confidences:: - - C_j = prod_{i<=j} q~_i # cumulative confidence (Eq. 8) - w_j = sum_{m>=j} C_m # = C_j * continuation value f~_j - - Each confidence is asymmetrically smoothed toward 1 (Eq. 7):: - - q~_i = (1 - alpha) * q_i + alpha, alpha in (0, 1], - - so the floor ``q~_i >= alpha`` keeps every cumulative product (hence every - weight) strictly positive. We evaluate the suffix sum from its definition as - ``total - exclusive_prefix_sum`` of ``C`` rather than reversing the tensor. - Positions with ``valid_mask == 0`` are multiplicative no-ops in ``C`` and - contribute nothing to the sum, matching the per-token loss mask. Weights are - detached (Eq. 9): they reweight the cross-entropy without adding gradient. - - Args: - confidences: ``[..., L]`` draft confidence ``q_i = exp(-CE_i)`` per position. - alpha: smoothing factor in (0, 1]; raises if outside that range. - valid_mask: optional ``[..., L]`` 0/1 mask; ``None`` treats all positions valid. - - Returns: - Detached weights with the same shape and dtype as ``confidences``. - """ - if not 0.0 < alpha <= 1.0: - raise ValueError(f"dflash_dpace_alpha must be in (0, 1], got {alpha}") - - with torch.no_grad(): - smoothed = alpha + (1.0 - alpha) * confidences.float() # Eq. 7 - if valid_mask is not None: - keep = valid_mask.to(torch.bool) - smoothed = torch.where(keep, smoothed, torch.ones_like(smoothed)) - cum_conf = torch.cumprod(smoothed, dim=-1) # Eq. 8 cumulative confidence C_j - if valid_mask is not None: - cum_conf = cum_conf * keep.to(cum_conf.dtype) - # Suffix sum w_j = sum_{m>=j} C_m, written as total minus the exclusive - # prefix sum so no axis reversal is needed (Eq. 8). - inclusive = torch.cumsum(cum_conf, dim=-1) - weights = inclusive[..., -1:] - inclusive + cum_conf - return weights.to(dtype=confidences.dtype) - - -class DraftLogits(NamedTuple): - """The draft's logits per block position: [B, N, bs, V].""" - - final: torch.Tensor # the distribution the draft serves - backbone: torch.Tensor # before any head; ``final`` itself for a head-less variant - - -class LossOutput(NamedTuple): - """A training step's loss, the accuracy ``train_acc`` reports, and its metrics.""" - - loss: torch.Tensor - accuracy: float - metrics: dict # name -> detached 0-dim tensor - - -class LossContext: - """One training step's supervision; costly per-token quantities are computed lazily, once. - - ``weights`` are the per-position weights after decay or D-PACE, which reads the draft's - confidence off ``ce("final")``. ``ce`` keeps its gradient only when ``ce_grad`` says the - objective backpropagates through it; otherwise it runs under ``no_grad``. - """ - - def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad): - """Weigh the step's block positions. - - Args: - batch: The forward's ``DraftBatch``. - logits: The draft's ``DraftLogits``. - targets: The labels and their mask, from ``_block_targets``. - weighting: ``_block_position_weights``' scheme, first, decay_factor and dpace_alpha. - teacher: Returns the base model's logits for every label position. - ce_grad: Whether the objective backpropagates through ``ce``. - """ - self.batch = batch - self.logits = logits - self.targets = targets - self._teacher_fn = teacher - self._ce_grad = ce_grad and torch.is_grad_enabled() - self._ce = {} - self._teacher = None - self.weights = _block_position_weights( - targets.mask, - **weighting, - target_ce=lambda: self.ce("final").detach().view(targets.mask.shape), - ) - self.weight_sum = self.weights.sum() + 1e-6 - - @functools.cached_property - def trainable(self) -> bool: - """Whether the weights leave anything to train on (reads the device once).""" - return bool(self.weight_sum > 1.0) - - def flat_logits(self, source): - """``logits.`` as [B * N * bs, V].""" - logits = getattr(self.logits, source) - return logits.reshape(-1, logits.shape[-1]) - - def ce(self, source): - """Per-token cross-entropy of ``logits.`` on the labels: [B * N * bs].""" - if source not in self._ce: - with torch.set_grad_enabled(self._ce_grad): - self._ce[source] = F.cross_entropy( - self.flat_logits(source), self.targets.token_ids.view(-1), reduction="none" - ) - return self._ce[source] - - def teacher(self): - """The base model's logits for every label position: [B, N, bs, V].""" - if self._teacher is None: - self._teacher = self._teacher_fn() - return self._teacher - - class DraftBatch(NamedTuple): """What a training step hands its loss: the batch, its blocks and the draft's output.""" diff --git a/modelopt/torch/speculative/plugins/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py index ebba52984a1..5d8760133e5 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash2.py +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -39,7 +39,8 @@ from transformers import PreTrainedModel from ..dflash.conversion import DFlash2DMRegistry -from .hf_dflash import HFDFlashModel, LossOutput, _block_tokens +from .dflash_losses import LossOutput, _block_tokens +from .hf_dflash import HFDFlashModel from .modeling_dflash2 import DFlash2Module __all__ = ["HFDFlash2Model"] diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 8d739f31253..2e924b121af 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -68,13 +68,8 @@ from transformers.utils import ModelOutput from ..dflash.conversion import DominoDMRegistry -from .hf_dflash import ( - HFDFlashModel, - _block_position_weights, - _block_targets, - _block_tokens, - _masked_top1, -) +from .dflash_losses import _block_position_weights, _block_targets, _block_tokens, _masked_top1 +from .hf_dflash import HFDFlashModel from .modeling_dflash import DFlashBaseModelOutput from .modeling_domino import DominoModule diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index bacabdf1790..0891f9a924b 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -63,19 +63,20 @@ import torch import torch.nn.functional as F -import torch.utils.checkpoint from transformers import PreTrainedModel from transformers.trainer_pt_utils import LabelSmoother from transformers.utils import ModelOutput from ..dflash.conversion import DSparkDMRegistry -from .hf_dflash import ( - HFDFlashModel, +from .dflash_losses import ( _block_position_weights, _block_targets, _block_tokens, _masked_top1, + _tvd_chunk, + _tvd_per_token, ) +from .hf_dflash import HFDFlashModel from .modeling_dflash import DFlashBaseModelOutput from .modeling_dspark import DSparkModule @@ -84,37 +85,6 @@ __all__ = ["HFDSparkModel"] -def _tvd_chunk(a, b): - """Per-token TVD for one row chunk: ``(softmax(a) - softmax(b)).abs().sum(-1)``.""" - return (torch.softmax(a.float(), dim=-1) - torch.softmax(b.float(), dim=-1)).abs().sum(dim=-1) - - -def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024, chunk_fn=None): - """Total-variation distance ||softmax(a)-softmax(b)||_1 / ... per token, memory-lean. - - Materializing both [N, vocab] float32 softmax tensors at once OOMs at large - N (num_anchors*block_size) and a 150k vocab. We compute it in row chunks and - gradient-checkpoint each chunk so the wide softmaxes are recomputed in backward - rather than held — peak memory ~ chunk_size*vocab instead of N*vocab. The math - is identical to ``(softmax(final)-softmax(teacher)).abs().sum(-1)``. - - Chunks with ``Tensor.split``, not slicing (each slice's backward zero-fills a full - [N, vocab] tensor) nor ``torch.chunk`` (different chunk shapes, so recompiles). - """ - _chunk = chunk_fn or _tvd_chunk - outs = [] - for a, b in zip( - final_logits.split(chunk_size, dim=0), - teacher_logits.split(chunk_size, dim=0), - strict=True, - ): - if torch.is_grad_enabled() and a.requires_grad: - outs.append(torch.utils.checkpoint.checkpoint(_chunk, a, b, use_reentrant=False)) - else: - outs.append(_chunk(a, b)) - return torch.cat(outs, dim=0) - - @DSparkDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) class HFDSparkModel(HFDFlashModel): """DFlash model with the DSpark sequential (Markov) + confidence head. diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py index 10f98f30824..97c8f6aea25 100644 --- a/modelopt/torch/speculative/plugins/hf_lilicorr.py +++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py @@ -77,7 +77,8 @@ from transformers import PreTrainedModel from ..dflash.conversion import LiLiCorrDMRegistry -from .hf_dflash import HFDFlashModel, LossOutput, _block_positions +from .dflash_losses import LossOutput, _block_positions +from .hf_dflash import HFDFlashModel from .modeling_lilicorr import LiLiCorrModule logger = logging.getLogger(__name__) From 16bd84e7dd0e00f8265fe4f46eb131a99a310cfb Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:15:03 +0000 Subject: [PATCH 8/9] test(speculative): follow the DFlash loss onto its context and dflash_losses.py Update the call sites of the internals the previous two commits changed or moved, keeping each test's intent: - _compute_loss tests (D-PACE/decay/KD branches, trailing padding blocks, KD from the base hidden vs. full logits) go through _loss_context and _compute_loss on given draft logits. - "KD without base_outputs raises" becomes "KD without a teacher raises": the context always carries the base outputs, so the remaining failure is a batch with neither the base logits nor its final hidden. - The selector and LiLiCorr metrics are read from draft_metrics, LiLiCorr's as tensors. - _block_targets, _dpace_position_weights and DSpark's TVD helpers are imported from dflash_losses, where the move put them. One new test pins which objectives keep the per-token cross-entropy's gradient. The context computes the CE once and sets its grad mode from ce_grad, overriding the no_grad under which D-PACE reads it first. The test inspects the context that a real DFlash2 training forward under D-PACE hands _compute_loss: ce, tv and lambda keep the CE's gradient, and KD does not. Mutation check. If _loss_context always asks for the gradient, LossContext ignores ce_grad, or the CE always keeps its graph, then each KD step under D-PACE (the default configuration) holds a vocab-wide graph that it never backpropagates. The loss and gradients stay bitwise equal (DFlash, DFlash2 and LiLiCorr KD, fp32 and bf16: 6/6), and the rest of the suite passes (512 passed); only [kd] fails. Dropping the gradient instead already failed three existing tests, and now also fails [ce], [tv] and [lambda]. tests/unit/torch/speculative and test_hf_spec_rope_export.py: 513 passed on transformers 5.5.4, 5.12.1 and 5.18.0. On 4 CPUs, the changed tests take at most 3.4 s each and the new cases under 0.1 s. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../speculative/plugins/test_hf_dflash.py | 42 ++++++++++++------ .../speculative/plugins/test_hf_dflash2.py | 36 ++++++++++++--- .../speculative/plugins/test_hf_dspark.py | 3 +- .../speculative/plugins/test_hf_lilicorr.py | 44 ++++++++++--------- 4 files changed, 84 insertions(+), 41 deletions(-) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index 306eb2679ea..57487aaa522 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -22,7 +22,7 @@ import logging import os from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -39,13 +39,17 @@ import modelopt.torch.speculative as mtsp import modelopt.torch.speculative.plugins.dflash_flex_attention as dflash_flex_attention import modelopt.torch.speculative.plugins.hf_dflash as hf_dflash +from modelopt.torch.speculative.plugins.dflash_losses import ( + DraftLogits, + _block_targets, + _dpace_position_weights, +) from modelopt.torch.speculative.plugins.hf_dflash import ( DFlashAttention, DFlashBaseModelOutput, DFlashModule, + DraftBatch, HFDFlashModel, - _block_targets, - _dpace_position_weights, build_target_layer_ids, ) from modelopt.torch.speculative.plugins.hf_dflash2 import HFDFlash2Model @@ -64,6 +68,15 @@ SEQ_LEN = 16 # must be multiple of BLOCK_SIZE +def _compute_loss(model, logits, input_ids, anchors, keep, loss_mask, base_outputs=None): + """``model._compute_loss`` on given draft logits [B, N * bs, V]: (loss, accuracy).""" + batch = DraftBatch(input_ids, loss_mask, anchors, keep, base_outputs, hidden=None) + block_logits = logits.view(*anchors.shape, BLOCK_SIZE, -1) + with patch.object(model, "_draft_logits", return_value=DraftLogits(block_logits, block_logits)): + out = model._compute_loss(model._loss_context(batch)) + return out.loss, out.accuracy + + class TestDFlashConvert: """Test DFlash model conversion.""" @@ -274,7 +287,7 @@ def _converted_model(self, objective, **overrides): def test_compute_loss_dpace_branch(self): """Default dpace objective produces a finite loss and valid accuracy.""" model = self._converted_model("dpace", dflash_self_logit_distillation=False) - loss, acc = model._compute_loss(*self._make_inputs()) + loss, acc = _compute_loss(model, *self._make_inputs()) assert torch.isfinite(loss).item() and loss.item() > 0 assert 0.0 <= acc <= 1.0 @@ -283,7 +296,7 @@ def test_compute_loss_decay_branch(self): model = self._converted_model( "decay", dflash_loss_decay_factor=4.0, dflash_self_logit_distillation=False ) - loss, acc = model._compute_loss(*self._make_inputs()) + loss, acc = _compute_loss(model, *self._make_inputs()) assert torch.isfinite(loss).item() and loss.item() > 0 assert 0.0 <= acc <= 1.0 @@ -293,15 +306,15 @@ def test_compute_loss_dpace_kd_branch(self): model = self._converted_model("dpace") inputs = self._make_inputs(vocab=vocab) teacher = DFlashBaseModelOutput(None, logits=torch.randn(1, SEQ_LEN, vocab)) - loss, acc = model._compute_loss(*inputs, base_outputs=teacher) + loss, acc = _compute_loss(model, *inputs, base_outputs=teacher) assert torch.isfinite(loss).item() assert 0.0 <= acc <= 1.0 - def test_kd_without_base_outputs_raises(self): + def test_kd_without_a_teacher_raises(self): """KD is decided by the config, so a missing teacher cannot fall back to CE.""" model = self._converted_model("dpace") - with pytest.raises(ValueError, match="base_outputs"): - model._compute_loss(*self._make_inputs()) + with pytest.raises(ValueError, match="base model's distribution"): + _compute_loss(model, *self._make_inputs(), base_outputs=DFlashBaseModelOutput(None)) class TestDFlashSaveRestore: @@ -1481,8 +1494,9 @@ def test_trailing_padding_blocks_do_not_change_the_loss(self): anchors = torch.tensor([[0, BLOCK_SIZE]])[:, :n_blocks] keep = torch.ones(bsz, n_blocks) - base_loss, base_acc = model._compute_loss(logits, input_ids, anchors, keep, loss_mask) - padded_loss, padded_acc = model._compute_loss( + base_loss, base_acc = _compute_loss(model, logits, input_ids, anchors, keep, loss_mask) + padded_loss, padded_acc = _compute_loss( + model, torch.cat([logits, torch.randn(bsz, pad * BLOCK_SIZE, vocab)], dim=1), input_ids, torch.cat([anchors, torch.zeros(bsz, pad, dtype=anchors.dtype)], dim=1), @@ -1652,9 +1666,9 @@ def test_kd_loss_matches_full_logits(self): hidden = self._hidden(model, bsz=1) inputs = TestDPaceLossIntegration._make_inputs(vocab=model.config.vocab_size) full = DFlashBaseModelOutput(None, logits=model._base_model_lm_head(hidden)) - loss_full, _ = model._compute_loss(*inputs, base_outputs=full) - loss, _ = model._compute_loss( - *inputs, base_outputs=DFlashBaseModelOutput(None, base_hidden=hidden) + loss_full, _ = _compute_loss(model, *inputs, base_outputs=full) + loss, _ = _compute_loss( + model, *inputs, base_outputs=DFlashBaseModelOutput(None, base_hidden=hidden) ) torch.testing.assert_close(loss, loss_full) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py index c43677818c0..21ca9ca59e5 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py @@ -26,6 +26,7 @@ import json from copy import deepcopy +from unittest.mock import patch import pytest import torch @@ -312,7 +313,7 @@ def test_selector_metrics_reported(self): out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) # Carried out on the forward output, not left on a private attribute, so the # trainer can log them; and kept as tensors to avoid a per-step CPU-GPU sync. - metrics = out["selector_metrics"] + metrics = out["draft_metrics"] for key in ("selector_accuracy", "selector_coverage"): assert torch.is_tensor(metrics[key]) assert 0.0 <= metrics[key].item() <= 1.0 @@ -373,16 +374,16 @@ def test_overfits_a_single_batch(self): optimizer.step() assert out.train_acc[0][0] > 0.9 - assert out["selector_metrics"]["selector_accuracy"].item() > 0.9 - assert out["selector_metrics"]["selector_coverage"].item() > 0.9 + assert out["draft_metrics"]["selector_accuracy"].item() > 0.9 + assert out["draft_metrics"]["selector_coverage"].item() > 0.9 # Candidate selection is a topk over draft logits; check it in bf16, the training dtype. model.to(torch.bfloat16) torch.manual_seed(0) with torch.no_grad(): out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) - assert out["selector_metrics"]["selector_accuracy"].item() > 0.9 - assert out["selector_metrics"]["selector_coverage"].item() > 0.9 + assert out["draft_metrics"]["selector_accuracy"].item() > 0.9 + assert out["draft_metrics"]["selector_coverage"].item() > 0.9 class TestCandidateSelectorAlignment: @@ -540,6 +541,31 @@ def test_acceptance_term_is_a_probability(self): loss = self._loss(dflash_lk_loss_type="tv", dflash_selector_loss_alpha=0.0) assert 0.0 <= loss <= 1.0 + @pytest.mark.parametrize( + ("loss_type", "distill", "ce_grad"), + [("ce", False, True), ("tv", False, True), ("lambda", False, True), ("ce", True, False)], + ids=["ce", "tv", "lambda", "kd"], + ) + def test_cross_entropy_gradient_follows_the_objective(self, loss_type, distill, ce_grad): + """The per-token CE keeps its graph exactly when the objective backpropagates through it. + + D-PACE reads it first, under no_grad. A tv or lambda objective computed on that would + train the backbone through the selector alone; KD would hold a vocab-wide graph unused. + """ + config = _get_dflash2_config() + config.update( + dflash_lk_loss_type=loss_type, + dflash_self_logit_distillation=distill, + dflash_loss_objective="dpace", + ) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + model.float().train() + input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) + with patch.object(model, "_compute_loss", wraps=model._compute_loss) as compute_loss: + model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + assert compute_loss.call_args.args[0].ce("final").requires_grad is ce_grad + @pytest.mark.parametrize("loss_type", ["tv", "lambda"]) def test_distillation_conflict_is_rejected(self, loss_type): """Both terms read q(gold), which the KD path never forms.""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index ad8db23babd..66fc64d37db 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -36,8 +36,9 @@ import modelopt.torch.quantization as mtq import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG +from modelopt.torch.speculative.plugins.dflash_losses import _tvd_chunk, _tvd_per_token from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel -from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel, _tvd_chunk, _tvd_per_token +from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel from modelopt.torch.speculative.plugins.modeling_dflash import ( DFlashBaseModelOutput, DFlashModule, diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py index 0d132f9ce49..301ff60946e 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py +++ b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py @@ -269,7 +269,7 @@ def test_forward_loss_and_metrics(self, variant): assert out.loss.dim() == 0 assert torch.isfinite(out.loss) - metrics = out.lilicorr_metrics + metrics = out.draft_metrics for key in ( "origin_loss", "lilicorr_loss", @@ -291,9 +291,9 @@ def test_loss_is_exactly_origin_plus_lilicorr(self, variant): model = _converted(variant=variant) model.train() out = model(**_make_batch(model.dflash_config.vocab_size)) - metrics = out.lilicorr_metrics + metrics = out.draft_metrics assert out.loss.item() == pytest.approx( - metrics["origin_loss"] + metrics["lilicorr_loss"], abs=1e-5 + (metrics["origin_loss"] + metrics["lilicorr_loss"]).item(), abs=1e-5 ) def test_hinge_only_active_in_the_margin_variant(self): @@ -301,20 +301,20 @@ def test_hinge_only_active_in_the_margin_variant(self): base = _converted(variant="base") base.train() assert ( - base(**_make_batch(base.dflash_config.vocab_size)).lilicorr_metrics["lilicorr_margin"] + base(**_make_batch(base.dflash_config.vocab_size)).draft_metrics["lilicorr_margin"] == 0.0 ) margin = _converted(variant="margin") margin.train() - metrics = margin(**_make_batch(margin.dflash_config.vocab_size)).lilicorr_metrics + metrics = margin(**_make_batch(margin.dflash_config.vocab_size)).draft_metrics # The hinge is relu(margin - gap) on an untrained head, so it is active. assert metrics["lilicorr_margin"] > 0.0 def test_penalty_term_is_active(self): model = _converted() model.train() - metrics = model(**_make_batch(model.dflash_config.vocab_size)).lilicorr_metrics + metrics = model(**_make_batch(model.dflash_config.vocab_size)).draft_metrics assert metrics["lilicorr_penalty"] > 0.0 @pytest.mark.parametrize("w_cal", [0.0, 0.5]) @@ -324,14 +324,16 @@ def test_calibration_enters_the_loss_with_its_weight(self, w_cal): # term has supervised slots to be positive on. model = _converted(topk=VOCAB_SIZE, dflash_lilicorr_w_cal=w_cal) model.train() - metrics = model(**_make_batch(model.dflash_config.vocab_size)).lilicorr_metrics + metrics = model(**_make_batch(model.dflash_config.vocab_size)).draft_metrics assert metrics["lilicorr_w_cal"] == w_cal - assert (metrics["lilicorr_calibration"] > 0.0) is (w_cal > 0.0) + assert bool(metrics["lilicorr_calibration"] > 0.0) is (w_cal > 0.0) weights = VARIANTS["base"] - assert metrics["lilicorr_loss"] == pytest.approx( - weights["w_ce"] * metrics["lilicorr_ce"] - + weights["w_pen"] * metrics["lilicorr_penalty"] - + w_cal * metrics["lilicorr_calibration"], + assert metrics["lilicorr_loss"].item() == pytest.approx( + ( + weights["w_ce"] * metrics["lilicorr_ce"] + + weights["w_pen"] * metrics["lilicorr_penalty"] + + w_cal * metrics["lilicorr_calibration"] + ).item(), abs=1e-5, ) @@ -423,9 +425,9 @@ def test_offline_forward_matches_online(self, prenorm): assert offline_out.loss.item() == pytest.approx(online_out.loss.item(), abs=1e-5) for term in ("lilicorr_penalty", "lilicorr_calibration"): - assert online_out.lilicorr_metrics[term] > 0.0 - assert offline_out.lilicorr_metrics[term] == pytest.approx( - online_out.lilicorr_metrics[term], abs=1e-5 + assert online_out.draft_metrics[term] > 0.0 + assert offline_out.draft_metrics[term].item() == pytest.approx( + online_out.draft_metrics[term].item(), abs=1e-5 ) @pytest.mark.parametrize( @@ -449,8 +451,8 @@ def test_degenerate_lattice_reports_finite_metrics(self, kwargs, degenerate): out = model(**_make_batch(model.dflash_config.vocab_size)) assert torch.isfinite(out.loss), degenerate - for name, value in out.lilicorr_metrics.items(): - assert math.isfinite(value), f"{name} is {value} ({degenerate})" + for name, value in out.draft_metrics.items(): + assert torch.isfinite(value), f"{name} is {value} ({degenerate})" def test_gradient_reaches_head_and_drafter_body(self): """Co-training: the lattice terms must reach the backbone, not only the head. @@ -492,14 +494,14 @@ def test_supervision_follows_lattice_coverage(self): """ wide = _converted(topk=VOCAB_SIZE) wide.train() - metrics = wide(**_make_batch(wide.dflash_config.vocab_size)).lilicorr_metrics + metrics = wide(**_make_batch(wide.dflash_config.vocab_size)).draft_metrics assert metrics["lilicorr_slot_gt_in_lattice"] == 1.0 assert metrics["lilicorr_oracle_prefix"] == float(BLOCK_SIZE - 1) assert metrics["lilicorr_zero_prefix_rate"] == 0.0 narrow = _converted(topk=CANDIDATE_TOPK) narrow.train() - metrics = narrow(**_make_batch(narrow.dflash_config.vocab_size)).lilicorr_metrics + metrics = narrow(**_make_batch(narrow.dflash_config.vocab_size)).draft_metrics assert metrics["lilicorr_slot_gt_in_lattice"] < 1.0 assert metrics["lilicorr_zero_prefix_rate"] > 0.0 @@ -531,7 +533,7 @@ def test_eval_forward_delegates_to_the_base_model(self): model.eval() batch = _make_batch(model.dflash_config.vocab_size) out = model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"]) - assert not hasattr(out, "lilicorr_metrics") or out.get("lilicorr_metrics") is None + assert not hasattr(out, "draft_metrics") or out.get("draft_metrics") is None def test_pseudo_speculative_generate_still_runs(self): """AR validation falls back to the backbone; it must not raise.""" @@ -658,7 +660,7 @@ def test_single_batch_overfit(self): out = model(**batch) out.loss.backward() optimizer.step() - losses.append(out.lilicorr_metrics["lilicorr_loss"]) + losses.append(out.draft_metrics["lilicorr_loss"].item()) assert all(math.isfinite(loss) for loss in losses), losses # A trend, not two endpoints. The DFlash forward draws its own masks, so successive From f11b7d820a87bc9998693e16a44cb497304c11a8 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:15:35 +0000 Subject: [PATCH 9/9] docs(speculative): list the draft_metrics change under backward breaking changes DFlash2's selector_metrics and LiLiCorr's lilicorr_metrics arrive in 0.48.0, which this change follows. The first commit of this series replaces both keys with draft_metrics, keeps their inner keys, and returns tensors where LiLiCorr returned Python floats. Nothing in the repo reads these keys, but user code that reads them breaks. The entry goes under the 0.49.0 heading: release/0.48.0 is branched and ships both keys, so this change first reaches users in 0.49.0. 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 815cae82aac..294bcc88e82 100755 --- a/CHANGELOG.rst +++ b/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**