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. diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 92ce9121b53..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": getattr(draft_config, "shift_label", True), - "pure_draft_prefix_len": getattr(draft_config, "pure_draft_prefix_len", 1), - "gru_hidden_dim": gru_hidden_dim, - "emb_dim": emb_dim, + "shift_label": bool(self.model._label_shift), + "pure_draft_prefix_len": head.pure_draft_prefix_len, + "gru_hidden_dim": head.gru_hidden_dim, + "emb_dim": head.emb_dim, } ) return config @@ -615,14 +615,17 @@ def _export_config(self): """Extend the DFlash config with the DSpark head fields.""" config = super()._export_config() draft_config = self.model.dflash_config + # Read off the built head, not the config dict: DSparkModule casts markov_rank to int + # and lower-cases markov_head_type, and the export must describe what was trained. + head = self.model.dflash_module config["dflash_config"].update( { "projector_type": getattr(draft_config, "projector_type", "dspark"), - "shift_label": getattr(draft_config, "shift_label", True), - "markov_rank": draft_config.markov_rank, - "markov_head_type": self.model.dflash_module.markov_head_type, - "use_confidence_head": bool(getattr(draft_config, "use_confidence_head", False)), + "shift_label": bool(self.model._label_shift), + "markov_rank": head.markov_rank, + "markov_head_type": head.markov_head_type, + "use_confidence_head": head.use_confidence_head, } ) if config["dflash_config"]["markov_head_type"] == "vanilla": @@ -634,7 +637,7 @@ def _export_config(self): "mask_token_id": config["dflash_config"]["mask_token_id"], "target_layer_ids": config["dflash_config"]["target_layer_ids"], "markov_rank": config["dflash_config"]["markov_rank"], - "sample_from_anchor": True, + "sample_from_anchor": config["dflash_config"]["shift_label"], "enable_confidence_head": config["dflash_config"]["use_confidence_head"], "confidence_head_with_markov": config["dflash_config"]["use_confidence_head"], } 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 d4d76d7e550..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 @@ -1219,36 +1229,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..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.""" @@ -266,14 +262,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 +353,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/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/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/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..8494f6f4b76 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 @@ -107,9 +98,8 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i | Error pattern | Root cause | Fix | | --- | --- | --- | | `Domino (projector_type='domino') requires ['emb_dim', 'gru_hidden_dim'] in dflash_architecture_config` | GRU head dims missing | Set both in `dflash_architecture_config` | -| `Domino currently supports shift_label=True (next-token alignment) only` | `shift_label=false` | Leave it at `true` | +| `shift_label=False is not supported: projector_type 'domino' trains only next-token alignment` | `shift_label: false` | Leave it at `true` | | `pure_draft_prefix_len must be in [0, N] (block_size=B), got X` | Prefix length >= block size | Lower it below `block_size` | | `DominoLambdaCallback: state.max_steps unset (<=0); lambda_base curriculum disabled` (warning) | `training.max_steps` not set | Set `training.max_steps`, else the curriculum never runs | -| `Domino eval uses the DFlash backbone only ...` (warning) | Correction head not applied at eval | Expected — do not chase it; evaluate after export once the inference path lands | | 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..bb2c6992ac5 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 @@ -151,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/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 5f2e8adc045..a0fc7301f39 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,48 @@ 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}, + {"projector_type": "domino", "emb_dim": 8, "gru_hidden_dim": 8}, + ], + ids=["dflash", "dspark", "domino"], + ) + 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. 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. 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/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 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