From 54dd95f25f0f0ca1c81e5c9d40c1ad78536050a1 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 1 Oct 2026 09:32:37 +0000 Subject: [PATCH 1/5] [3/8] Add the external-draft recipe and training entrypoint Registers the speculative_external recipe type, its schema class, and the external_draft.yaml recipe, and teaches main.py to train the mode added in [2/8]. The recipe inverts the usual relationship: model.model_name_or_path is the base, used only for its lm_head and vocabulary, while draft_model_name_or_path names the pretrained draft that is trained. The vocabulary check runs before training because a mismatch otherwise yields a plausible loss curve and a draft with no acceptance. Documents the mode in the example README and adds the changelog entry for the series. Signed-off-by: Peter Hicks Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- CHANGELOG.rst | 4 + examples/speculative_decoding/README.md | 39 ++++ examples/speculative_decoding/main.py | 120 +++++++++++- modelopt/recipe/config.py | 34 +++- modelopt/recipe/loader.py | 1 + .../torch/speculative/plugins/hf_dflash.py | 11 +- .../torch/speculative/plugins/hf_eagle.py | 11 +- .../speculative/plugins/modeling_fakebase.py | 26 +++ .../speculative_decoding/external_draft.yaml | 58 ++++++ .../test_speculative_decoding_main.py | 172 ++++++++++++++++++ tests/unit/recipe/test_loader.py | 1 + .../test_speculative_external_recipe.py | 34 ++++ .../speculative/plugins/test_fakebase.py | 39 +++- 13 files changed, 529 insertions(+), 21 deletions(-) create mode 100644 modelopt_recipes/general/speculative_decoding/external_draft.yaml create mode 100644 tests/unit/examples/test_speculative_decoding_main.py create mode 100644 tests/unit/recipe/test_speculative_external_recipe.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 84fd08d1233..6093d8ac132 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -11,6 +11,10 @@ Changelog - Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ on the dense linears and imatrix on the experts GPTQ skips; the built-in IQ1 and IQ2 recipes now do this, and new ``_experts_only`` recipes apply imatrix to the routed MoE experts alone. - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. +*Speculative Decoding* + +- Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint. + **Backward Breaking Changes** **Deprecations** diff --git a/examples/speculative_decoding/README.md b/examples/speculative_decoding/README.md index 23ba03f6e39..943cc6e005f 100644 --- a/examples/speculative_decoding/README.md +++ b/examples/speculative_decoding/README.md @@ -19,6 +19,7 @@ This example focuses on training with Hugging Face. To train with Megatron‑LM, | Online Training | Train draft model alongside base model in GPU memory | \[[Link](#training-draft-model-with-online-base-model)\] | | Offline Training | Train draft model using pre-computed hidden states | \[[Link](#training-draft-model-with-offline-base-model)\] | | Streaming Training | Train draft on hidden states streamed from a live vLLM serve (no disk dump) | \[[Link](#training-draft-model-with-streaming-base-model)\] | +| External Draft Training | Train a standalone pretrained LM as the draft, not a grafted head | \[[Link](#training-a-standalone-external-draft-model)\] | | After Training | Evaluation, export and deployment | \[[Link](#model-validation)\] | | Advanced Usage | Data synthesis, vocab compression, and configuration | \[[Link](#advanced-usage)\] | | Support Matrix | Supported models for speculative decoding training | \[[Link](#support-matrix)\] | @@ -132,6 +133,44 @@ Once we finish dumping hidden states, launch offline training pointing to the hi For large base models, you can stream hidden states from a live `vllm serve` instead of dumping them to disk: a co-located server produces the base-model hidden states on the fly and sends them to the trainer over NIXL RDMA, scaling to multiple nodes (dedicated serve replicas + DDP trainers). See the launcher examples, e.g. [Kimi-K2.5 streaming EAGLE3](../../tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_eagle3_multi_node.yaml) and [streaming DFlash](../../tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_dflash_multi_node.yaml). +## Training a Standalone (External) Draft Model + +The modes above graft a draft head onto the target. The `external` mode instead +trains an ordinary pretrained causal LM as the draft: it keeps its own embeddings +and lm_head, the target is never modified, and the result exports as a plain HF +checkpoint that TRT-LLM, vLLM and SGLang load with no mode-specific support. The +draft must index the same vocabulary as the target. + +Training consumes dumped base hidden states, as in the offline flow above. Pass +`--no-aux-hidden-states` when dumping: this mode never reads them and they +dominate the dump size. Streaming feeds the same hidden states from a live +`vllm serve` and is accepted here, though it has not been exercised for this +mode. Online training is rejected: it produces no cached base outputs for the +draft to be scored against. + +```bash +torchrun --nproc_per_node 8 main.py \ + --config ../../modelopt_recipes/general/speculative_decoding/external_draft.yaml \ + model.model_name_or_path=$BASE_MODEL \ + draft_model_name_or_path=$DRAFT_MODEL \ + data.offline_data_path=$HIDDEN_STATES_DIR \ + external.external_loss=tvd \ + training.output_dir=$OUTPUT_DIR +``` + +`external.external_loss` selects the objective: + +| value | what it scores | +|---|---| +| `soft_ce` | cross-entropy against the base's soft distribution (the EAGLE offline objective) | +| `tvd` | total-variation distance between the two full distributions | +| `tvd_deploy` | the same, with both sides put through the serving filter (`external_top_k`, then `external_top_p`) first, so the loss is the acceptance the deployed pair sees | + +Under rejection sampling the acceptance probability is exactly `sum(min(p, q))`, +so `tvd_deploy` is `1 - acceptance` — the deployment metric itself rather than a +proxy for it. Report which truncation you trained under, since `top_k` and +`top_p` settings are not interchangeable. + ## Model Validation For online training checkpoints, we can run in-framework evaluation on MT-bench: diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index bd06f24e8ac..5ed811f71fb 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -53,15 +53,18 @@ from modelopt.recipe.config import ( ModelOptDFlashRecipe, ModelOptEagleRecipe, + ModelOptExternalDraftRecipe, ModelOptMedusaRecipe, ModelOptSpeculativeRecipeBase, ) +from modelopt.torch.opt.conversion import ModeloptStateManager from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel from modelopt.torch.speculative.plugins.hf_domino import DominoLambdaCallback from modelopt.torch.speculative.plugins.hf_training_args import ( TrainingArguments as SpecTrainingArgs, ) from modelopt.torch.speculative.plugins.master_weight_adamw import VerifyMasterWeightsCallback +from modelopt.torch.speculative.plugins.modeling_fakebase import find_final_norm_path from modelopt.torch.speculative.utils import load_vlm_or_llm, patch_transformers5_params_loading from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.distributed import is_master, local_rank @@ -170,12 +173,91 @@ def _is_hf_format_checkpoint(checkpoint: str | None) -> bool: return any(os.path.isfile(os.path.join(checkpoint, f)) for f in hf_files) +def setup_external_draft(recipe, training_args, base_model, draft=None, dry_run=False): + """Convert an external draft and wire it to the base model. + + Runs on both the fresh-load and the HF-resume path. The base lm_head is never part + of the draft's own state, so a resumed draft must be re-attached to it either way; + doing this only on the fresh path leaves a resumed run to fail at the first step + with no lm_head attached. + + A checkpoint written by this mode does carry ``modelopt_state.pth``, so the patched + ``from_pretrained`` already restored the ``external`` mode and handed back a + converted draft. Converting it a second time raises ``TypeError: Cannot create a + consistent method resolution order``, so the conversion is skipped when the draft is + already converted -- the validations and the head attachment below still run. + + Args: + base_model: the target, kept only for its lm_head and vocabulary. + draft: the already-loaded draft on the resume path; loaded here when None. + dry_run: recipe validation only; skips the data-source requirement, as the + ``data.data_path`` check in :func:`train` does. + """ + if recipe.draft_model_name_or_path is None: + raise ValueError( + "draft_model_name_or_path must be set in the recipe YAML or via a dotlist override." + ) + # Checked here, not on the recipe schema: the loader validates the recipe file once + # before splicing CLI overrides, so a schema-level raise would make it impossible to + # pass the data path on the command line -- the shipped recipe carries none. + # Only "online" is rejected: it is the one mode that produces no cached base + # outputs at all. "streaming" feeds the same pre-computed hidden states offline + # does, so it is permitted but untested for this mode -- if its batches ever lack + # base_model_outputs the forward raises rather than training on nothing. + if not dry_run and recipe.data.mode == "online": + raise ValueError( + "external draft training has no online path: the draft is scored against " + "cached base outputs, which an online run does not produce. Set " + "data.offline_data_path." + ) + base_vocab_size = base_model.config.vocab_size + if draft is None: + draft = transformers.AutoModelForCausalLM.from_pretrained( + recipe.draft_model_name_or_path, + dtype="auto", + device_map="cpu", + trust_remote_code=recipe.model.trust_remote_code, + ) + external_cfg: dict = recipe.external.model_dump() + # The draft's own tokenizer; the base's is used only for its lm_head. + tokenizer = transformers.AutoTokenizer.from_pretrained( + recipe.draft_model_name_or_path, + model_max_length=training_args.training_seq_len, + trust_remote_code=recipe.model.trust_remote_code, + ) + if not ModeloptStateManager.is_converted(draft): + mtsp.convert(draft, [("external", external_cfg)]) + # Fail before training: a vocab mismatch yields a plausible loss curve and a draft + # with no acceptance. Equal sizes are not enough -- different tokenizers are + # routinely padded to the same width -- so also check that the two agree on the ids. + draft.validate_against_base(base_vocab_size) + base_tokenizer = transformers.AutoTokenizer.from_pretrained( + recipe.model.model_name_or_path, + trust_remote_code=recipe.model.trust_remote_code, + ) + draft.validate_tokenizer_against_base(tokenizer, base_tokenizer) + # The norm is not always at base_model.model.norm -- FakeBaseModel keeps it at .norm + # -- so search the same paths the other offline modes use. It stays optional; + # _teacher_logits raises only if the dump declares a pre-norm hidden. + norm_path = find_final_norm_path(base_model) + base_final_norm = base_model.get_submodule(norm_path) if norm_path else None + # Attach on the device training will run on: the draft is loaded with + # device_map="cpu", so otherwise the first forward moves the copies mid-step on + # every rank. training_args.device is the device accelerate assigned this process, + # which is correct both when every rank sees all GPUs and when CUDA_VISIBLE_DEVICES + # pins one GPU per rank -- local_rank arithmetic gets the second case wrong. + draft.attach_base_lm_head( + base_model.get_output_embeddings(), base_final_norm, device=training_args.device + ) + return draft, tokenizer + + def train(): config_path, dry_run, overrides = _parse_cli() recipe = load_recipe(config_path, overrides=overrides) if not isinstance(recipe, ModelOptSpeculativeRecipeBase): raise ValueError( - f"main.py expects a speculative-decoding recipe (eagle / dflash / medusa); " + f"main.py expects a speculative-decoding recipe (eagle / dflash / medusa / external); " f"got {type(recipe).__name__} from {config_path!r}." ) @@ -227,6 +309,32 @@ def train(): tokenizer = transformers.AutoTokenizer.from_pretrained( checkpoint, trust_remote_code=recipe.model.trust_remote_code ) + if isinstance(recipe, ModelOptExternalDraftRecipe): + # The restored draft is already converted (the checkpoint carries + # modelopt_state.pth), but the base lm_head it is scored against is not part + # of its state, so it still has to be re-attached. The base is not on this + # path at all, so load it here; use_fake_base_for_offline keeps that cheap. + if recipe.model.model_name_or_path is None: + raise ValueError( + "model.model_name_or_path must be set to resume an external-draft run: " + "the base model supplies the lm_head the draft is scored against." + ) + # Same arguments as the fresh path below. use_offline_training is True + # unconditionally because this mode has no online path -- and without it + # load_vlm_or_llm ignores use_fake_base and materialises every decoder + # layer, on every rank, to borrow an lm_head and a final norm. + base_model = load_vlm_or_llm( + recipe.model.model_name_or_path, + use_fake_base=recipe.model.use_fake_base_for_offline, + use_offline_training=True, + dtype="auto", + device_map="cpu", + trust_remote_code=recipe.model.trust_remote_code, + config_overrides=recipe.model.config_overrides, + ) + model, tokenizer = setup_external_draft( + recipe, training_args, base_model=base_model, draft=model, dry_run=dry_run + ) else: if checkpoint: print_rank_0( @@ -273,6 +381,12 @@ def train(): ) dflash_cfg: dict = recipe.dflash.model_dump() mtsp.convert(model, [("dflash", dflash_cfg)]) + elif isinstance(recipe, ModelOptExternalDraftRecipe): + # This mode inverts the usual relationship: we train a pretrained draft, and + # `model` loaded above is the base, kept only for its lm_head and vocabulary. + model, tokenizer = setup_external_draft( + recipe, training_args, base_model=model, dry_run=dry_run + ) else: raise ValueError(f"Unsupported speculative recipe type: {type(recipe).__name__}") @@ -306,6 +420,10 @@ def train(): answer_only_loss=training_args.answer_only_loss, shift_labels=not is_dflash, final_aux_is_base_hidden=recipe.data.final_aux_is_base_hidden, + # The external draft never reads the aux planes, so it can consume dumps + # written with --no-aux-hidden-states. Every other mode must still fail + # loudly on a dump that omits them. + aux_hidden_states_optional=isinstance(recipe, ModelOptExternalDraftRecipe), ) callbacks = [EagleTrainingPlot(training_args.ar_validate_steps, training_args.estimate_ar)] diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index 9639ffec56c..a1e71efc23b 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -25,7 +25,12 @@ from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField from modelopt.torch.quantization.config import QuantizeConfig # noqa: TC001 -from modelopt.torch.speculative.config import DFlashConfig, EagleConfig, MedusaConfig +from modelopt.torch.speculative.config import ( + DFlashConfig, + EagleConfig, + ExternalDraftConfig, + MedusaConfig, +) from modelopt.torch.speculative.plugins.hf_training_args import DataArguments as SpecDataArgs from modelopt.torch.speculative.plugins.hf_training_args import ModelArguments as SpecModelArgs from modelopt.torch.speculative.plugins.hf_training_args import ( @@ -58,6 +63,7 @@ class RecipeType(str, Enum): SPECULATIVE_EAGLE = "speculative_eagle" SPECULATIVE_DFLASH = "speculative_dflash" SPECULATIVE_MEDUSA = "speculative_medusa" + SPECULATIVE_EXTERNAL = "speculative_external" # QAT = "qat" # Not implemented yet, will be added in the future. @@ -524,6 +530,31 @@ def _derive_dflash_offline(self) -> ModelOptDFlashRecipe: return self +class ModelOptExternalDraftRecipe(ModelOptSpeculativeRecipeBase): + """Our config class for external (standalone) draft speculative decoding recipes.""" + + RECIPE_TYPE: ClassVar[RecipeType] = RecipeType.SPECULATIVE_EXTERNAL + + metadata: RecipeMetadataConfig = _metadata_field() + + external: ExternalDraftConfig = ModeloptField( + default=ExternalDraftConfig(), + title="External draft config", + description="External (standalone) draft model configuration.", + validate_default=True, + ) + + draft_model_name_or_path: str | None = ModeloptField( + default=None, + title="Draft model", + description=( + "Path or HF id of the pretrained draft model to train. Unlike the head-based " + "recipes, ``model.model_name_or_path`` refers to the *base* model, which is " + "used only for its lm_head and vocabulary." + ), + ) + + class ModelOptMedusaRecipe(ModelOptSpeculativeRecipeBase): """Our config class for Medusa speculative decoding recipes.""" @@ -547,4 +578,5 @@ class ModelOptMedusaRecipe(ModelOptSpeculativeRecipeBase): RecipeType.SPECULATIVE_EAGLE: ModelOptEagleRecipe, RecipeType.SPECULATIVE_DFLASH: ModelOptDFlashRecipe, RecipeType.SPECULATIVE_MEDUSA: ModelOptMedusaRecipe, + RecipeType.SPECULATIVE_EXTERNAL: ModelOptExternalDraftRecipe, } diff --git a/modelopt/recipe/loader.py b/modelopt/recipe/loader.py index d0a01718e0b..87af5f6ff48 100644 --- a/modelopt/recipe/loader.py +++ b/modelopt/recipe/loader.py @@ -52,6 +52,7 @@ RecipeType.SPECULATIVE_EAGLE: "eagle", RecipeType.SPECULATIVE_DFLASH: "dflash", RecipeType.SPECULATIVE_MEDUSA: "medusa", + RecipeType.SPECULATIVE_EXTERNAL: "external", } diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index cd596f773d9..9833f285521 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -93,8 +93,8 @@ from .modeling_fakebase import ( _BASE_MODEL_PATHS, _EMBED_TOKENS_PATHS, - _FINAL_NORM_PATHS, _LM_HEAD_PATHS, + find_final_norm_path, ) logger = logging.getLogger(__name__) @@ -250,14 +250,7 @@ def _find_base_model_parts(self): raise ValueError(f"Part {name} not found in model") # Final pre-lm_head norm is OPTIONAL (set None if absent): used to re-normalize the # un-normed final hidden collect by vllm. - self.base_model_norm_path = None - for path in _FINAL_NORM_PATHS: - try: - assert isinstance(self.get_submodule(path), torch.nn.Module) - self.base_model_norm_path = path - break - except Exception: - continue + self.base_model_norm_path = find_final_norm_path(self) def modify(self, config): """Initialize DFlash draft module.""" diff --git a/modelopt/torch/speculative/plugins/hf_eagle.py b/modelopt/torch/speculative/plugins/hf_eagle.py index 94399912983..e7c931c1c86 100644 --- a/modelopt/torch/speculative/plugins/hf_eagle.py +++ b/modelopt/torch/speculative/plugins/hf_eagle.py @@ -43,8 +43,8 @@ from .modeling_fakebase import ( _BASE_MODEL_PATHS, _EMBED_TOKENS_PATHS, - _FINAL_NORM_PATHS, _LM_HEAD_PATHS, + find_final_norm_path, ) from .modeling_final_norm import _maybe_apply_base_final_norm @@ -138,14 +138,7 @@ def _find_base_model_parts(self): # Final pre-lm_head norm is OPTIONAL (set None if absent): used to re-normalize the # un-normed final hidden captured by vLLM in the offline/streaming path. - self.base_model_norm_path = None - for path in _FINAL_NORM_PATHS: - try: - assert isinstance(self.get_submodule(path), torch.nn.Module) - self.base_model_norm_path = path - break - except Exception: - continue + self.base_model_norm_path = find_final_norm_path(self) def _activate_torch_compile(self): import torch._dynamo diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index 9ef6a1d2671..85a3fe3a9f4 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py @@ -58,6 +58,32 @@ "language_model.backbone.norm", "model.language_model.norm", ] + + +def find_final_norm_path(model) -> str | None: + """Return the attribute path of ``model``'s final norm, or ``None`` if it has none. + + The norm is not always at ``model.norm``: a real causal LM keeps it there while + ``FakeBaseModel`` holds it at ``.norm``. Returns the path rather than the module + because the speculative plugins store the path and resolve it later, while callers + that want the module can ``get_submodule`` it. + + Resolves each path with ``get_submodule`` rather than a ``named_modules()`` dict: + ``named_modules`` de-duplicates by default, so a norm instance also registered + under an earlier path would be missing from the dict and the lookup would return + ``None`` for a model that has one. ``get_submodule`` walks the attribute path and + is indifferent to aliasing. + """ + for path in _FINAL_NORM_PATHS: + try: + module = model.get_submodule(path) + except AttributeError: + continue + if isinstance(module, torch.nn.Module): + return path + return None + + _BASE_MODEL_PATHS = [ "language_model.model", "model.language_model", diff --git a/modelopt_recipes/general/speculative_decoding/external_draft.yaml b/modelopt_recipes/general/speculative_decoding/external_draft.yaml new file mode 100644 index 00000000000..106d9072b79 --- /dev/null +++ b/modelopt_recipes/general/speculative_decoding/external_draft.yaml @@ -0,0 +1,58 @@ +# External (standalone) draft speculative-decoding training recipe. +# Unlike the head-based recipes, model.model_name_or_path is the *base* model -- +# used only for its lm_head and vocabulary -- and draft_model_name_or_path names +# the pretrained draft that is actually trained. +# Override fields via OmegaConf dotlist on the CLI. + +# modelopt-schema: modelopt.recipe.config.ModelOptExternalDraftRecipe +metadata: + description: External standalone draft training recipe (model/data/training/external bundled). + +# maps to ModelArguments (main.py); this is the BASE model +model: + model_name_or_path: + # Matches every other recipe and the ModelArguments default. Pass + # model.trust_remote_code=true on the CLI for a base that genuinely needs it. + trust_remote_code: false + use_fake_base_for_offline: true + +# the pretrained draft to train +draft_model_name_or_path: + +# maps to DataArguments (main.py) +data: + # Auto-derived from the data-source fields below; leave unset. + mode: + data_path: + offline_data_path: + draft_vocab_cache: + vlm_img_dir: + vlm_processor: + +# maps to TrainingArguments (main.py) +training: + output_dir: + num_train_epochs: 1 + per_device_train_batch_size: 1 + learning_rate: 1.0e-4 + warmup_steps: 20 + training_seq_len: 2048 + logging_steps: 1 + save_steps: 8192 + cp_size: 1 + disable_tqdm: false + estimate_ar: false + ar_validate_steps: -1 + answer_only_loss: true + bf16: true + # The offline collator needs base_model_hidden_states/aux_hidden_states, which HF + # Trainer's RemoveColumnsCollator would strip: they are not named in the model's + # forward signature. Every speculative recipe sets this. + remove_unused_columns: false + +external: + # soft_ce matches the EAGLE offline objective. tvd is plain total-variation + # distance; tvd_deploy filters both sides to the serving policy first and is the + # acceptance the deployed pair sees. + external_loss: tvd + external_report_acc: true diff --git a/tests/unit/examples/test_speculative_decoding_main.py b/tests/unit/examples/test_speculative_decoding_main.py new file mode 100644 index 00000000000..e6e86e75a97 --- /dev/null +++ b/tests/unit/examples/test_speculative_decoding_main.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""``setup_external_draft`` wiring in the speculative-decoding entrypoint. + +Lives under ``tests/unit`` rather than ``tests/examples/`` so CI runs it on every +change: the example lanes cover a fixed allowlist. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch + +transformers = pytest.importorskip("transformers") + +import modelopt.torch.speculative as mtsp +from modelopt.recipe.config import ModelOptExternalDraftRecipe + +_EXAMPLE_DIR = Path(__file__).resolve().parents[3] / "examples" / "speculative_decoding" + +VOCAB = 64 +HIDDEN = 32 + + +@pytest.fixture(scope="module") +def main_mod(): + """Import the entrypoint by path; its sibling modules need the example dir on sys.path.""" + sys.path.insert(0, str(_EXAMPLE_DIR)) + try: + spec = importlib.util.spec_from_file_location("sd_main", _EXAMPLE_DIR / "main.py") + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + except ImportError as e: # optional example-only deps + pytest.skip(f"speculative_decoding/main.py not importable: {e}") + finally: + sys.path.pop(0) + return module + + +def _tiny_causal_lm(vocab_size=VOCAB, hidden=HIDDEN): + cfg = transformers.LlamaConfig( + vocab_size=vocab_size, + hidden_size=hidden, + intermediate_size=2 * hidden, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=128, + ) + return transformers.LlamaForCausalLM(cfg) + + +class _FakeTok: + """Enough of a tokenizer for the vocabulary checks; both sides share one vocab.""" + + def __init__(self): + self._vocab = {chr(i): i for i in range(VOCAB)} + + def get_vocab(self): + return dict(self._vocab) + + def __len__(self): + return len(self._vocab) + + +def _recipe(**data): + return ModelOptExternalDraftRecipe.model_validate( + { + "metadata": {"description": "d"}, + "model": {"model_name_or_path": "base"}, + "draft_model_name_or_path": "draft", + "data": data, + } + ) + + +@pytest.fixture +def stub_tokenizers(monkeypatch, main_mod): + monkeypatch.setattr( + main_mod.transformers.AutoTokenizer, + "from_pretrained", + classmethod(lambda cls, *a, **k: _FakeTok()), + ) + + +def _training_args(): + return types.SimpleNamespace(training_seq_len=128, device=torch.device("cpu")) + + +def test_resume_of_a_converted_checkpoint_does_not_reconvert(main_mod, stub_tokenizers, tmp_path): + """A checkpoint this mode writes carries modelopt_state, so it restores already converted. + + Converting it a second time raises ``TypeError: Cannot create a consistent method + resolution order``, so any HF-format resume died at startup. + """ + draft = _tiny_causal_lm() + mtsp.convert(draft, [("external", {})]) + draft.save_pretrained(tmp_path / "ckpt") + + restored = transformers.AutoModelForCausalLM.from_pretrained(tmp_path / "ckpt") + assert main_mod.ModeloptStateManager.is_converted(restored), ( + "checkpoint did not restore as converted; this test no longer covers the bug" + ) + + base = _tiny_causal_lm(hidden=2 * HIDDEN) + out, _ = main_mod.setup_external_draft( + _recipe(offline_data_path="/tmp/hidden"), + _training_args(), + base_model=base, + draft=restored, + ) + assert out is restored + # The head is not part of the draft's own state, so a resume must re-attach it. + assert out._base_lm_head is not None + + +def test_exported_architecture_survives_a_resume(main_mod, stub_tokenizers, tmp_path): + """A second convert would nest the dynamic class and export it as ``architectures``.""" + draft = _tiny_causal_lm() + mtsp.convert(draft, [("external", {})]) + draft.save_pretrained(tmp_path / "ckpt") + restored = transformers.AutoModelForCausalLM.from_pretrained(tmp_path / "ckpt") + + main_mod.setup_external_draft( + _recipe(offline_data_path="/tmp/hidden"), + _training_args(), + base_model=_tiny_causal_lm(hidden=2 * HIDDEN), + draft=restored, + ) + restored.save_pretrained(tmp_path / "out") + config = transformers.AutoConfig.from_pretrained(tmp_path / "out") + assert config.architectures == ["LlamaForCausalLM"], ( + f"draft exported as {config.architectures}; serving frameworks key off this name" + ) + + +def test_dry_run_is_exempt_from_the_data_source_requirement(main_mod, stub_tokenizers): + """The shipped recipe carries no data path, so it derives mode="online". + + --dry_run only validates the recipe and writes an untrained checkpoint, so it must + not require a dataset -- matching the data_path check in train(). + """ + recipe = _recipe() + assert recipe.data.mode == "online" + base = _tiny_causal_lm(hidden=2 * HIDDEN) + + with pytest.raises(ValueError, match="no online path"): + main_mod.setup_external_draft( + recipe, _training_args(), base_model=base, draft=_tiny_causal_lm() + ) + + draft, _ = main_mod.setup_external_draft( + recipe, _training_args(), base_model=base, draft=_tiny_causal_lm(), dry_run=True + ) + assert draft is not None diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index e1add8623bf..a02128e5e39 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -860,6 +860,7 @@ def test_metadata_config_recipe_type_is_optional(tmp_path): RecipeType.SPECULATIVE_EAGLE: {}, # body sections have field defaults RecipeType.SPECULATIVE_DFLASH: {}, RecipeType.SPECULATIVE_MEDUSA: {}, + RecipeType.SPECULATIVE_EXTERNAL: {}, } diff --git a/tests/unit/recipe/test_speculative_external_recipe.py b/tests/unit/recipe/test_speculative_external_recipe.py new file mode 100644 index 00000000000..daf5674edf8 --- /dev/null +++ b/tests/unit/recipe/test_speculative_external_recipe.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Schema tests for the ``speculative_external`` recipe.""" + +from modelopt.recipe.config import ModelOptExternalDraftRecipe + + +def test_recipe_validates_without_a_data_path_so_overrides_stay_reachable(): + """The loader validates the recipe file *before* splicing CLI overrides. + + A schema-level raise on the missing data path would therefore fire on the shipped + recipe, which carries none, and make `data.offline_data_path=...` on the command line + impossible to pass. The gate lives in main.py, after the overrides are in. + """ + bare = ModelOptExternalDraftRecipe.model_validate({"metadata": {"description": "d"}}) + assert bare.data.mode == "online" + + ok = ModelOptExternalDraftRecipe.model_validate( + {"metadata": {"description": "d"}, "data": {"offline_data_path": "/tmp/hidden"}} + ) + assert ok.data.mode == "offline" diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index 99dace9ea9c..44cc19ed90c 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/tests/unit/torch/speculative/plugins/test_fakebase.py @@ -24,7 +24,11 @@ pytest.importorskip("transformers") import transformers -from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseConfig, FakeBaseModel +from modelopt.torch.speculative.plugins.modeling_fakebase import ( + FakeBaseConfig, + FakeBaseModel, + find_final_norm_path, +) from modelopt.torch.speculative.utils import load_vlm_or_llm _HIDDEN_SIZE = 16 @@ -241,3 +245,36 @@ def test_unknown_theta_publishes_no_dict(self): config = FakeBaseConfig(num_hidden_layers=2, hidden_size=32, rope_theta=None) assert config.rope_theta is None assert not getattr(config, "rope_parameters", None) + + +def test_find_final_norm_path_resolves_a_shared_norm_instance(): + """A norm instance also registered under an earlier path must still be found. + + ``named_modules()`` de-duplicates by default and yields a shared module only under + the first name it reaches, so a dict built from it can be missing ``model.norm`` + entirely. The failure is silent: callers read ``None``, skip the norm, and feed an + un-normed residual to lm_head -- a plausible loss curve against a garbage teacher. + """ + shared = torch.nn.LayerNorm(8) + + class _Aliased(torch.nn.Module): + def __init__(self): + super().__init__() + # "aaa" sorts before "model", so traversal reaches the alias first. + self.aaa = torch.nn.Module() + self.aaa.norm = shared + self.model = torch.nn.Module() + self.model.norm = shared + + model = _Aliased() + assert "model.norm" not in dict(model.named_modules()), "alias no longer shadows" + assert find_final_norm_path(model) == "model.norm" + + +def test_find_final_norm_path_returns_none_without_a_norm(): + class _NoNorm(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Linear(4, 4) + + assert find_final_norm_path(_NoNorm()) is None From 8c2cb2ae59e9f1ed657239ec52e0efd5fcf3079e Mon Sep 17 00:00:00 2001 From: Peter Hicks <60696413+phicksy@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:09:44 +0100 Subject: [PATCH 2/5] Update examples/speculative_decoding/main.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Peter Hicks <60696413+phicksy@users.noreply.github.com> --- examples/speculative_decoding/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 5ed811f71fb..8b741b75b35 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -225,7 +225,7 @@ def setup_external_draft(recipe, training_args, base_model, draft=None, dry_run= model_max_length=training_args.training_seq_len, trust_remote_code=recipe.model.trust_remote_code, ) - if not ModeloptStateManager.is_converted(draft): + if not ModeloptStateManager.has_state_for_mode_type("speculative", model=draft): mtsp.convert(draft, [("external", external_cfg)]) # Fail before training: a vocab mismatch yields a plausible loss curve and a draft # with no acceptance. Equal sizes are not enough -- different tokenizers are From 930e19fe74308e07a9b5de7c5ccad640bc05b3af Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 15 Oct 2026 15:25:49 +0000 Subject: [PATCH 3/5] Re-apply the recipe's external config when resuming a converted draft The HF-resume path skips mtsp.convert when the checkpoint already carries speculative state, which also skips modify(). The restored draft therefore keeps the objective saved in the checkpoint while the printed recipe reports the edited one, so resuming with a changed external_loss trains the old objective silently. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- examples/speculative_decoding/main.py | 7 +++++ .../test_speculative_decoding_main.py | 26 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 8b741b75b35..05292309e72 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -58,6 +58,7 @@ ModelOptSpeculativeRecipeBase, ) from modelopt.torch.opt.conversion import ModeloptStateManager +from modelopt.torch.speculative.config import ExternalDraftConfig from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel from modelopt.torch.speculative.plugins.hf_domino import DominoLambdaCallback from modelopt.torch.speculative.plugins.hf_training_args import ( @@ -227,6 +228,12 @@ def setup_external_draft(recipe, training_args, base_model, draft=None, dry_run= ) if not ModeloptStateManager.has_state_for_mode_type("speculative", model=draft): mtsp.convert(draft, [("external", external_cfg)]) + else: + # The restore replayed the config stored in the checkpoint, so modify() has + # already run with the values this draft was last trained under. Re-apply the + # recipe's: otherwise an edited external_loss / top_k / top_p is dropped while + # the banner and the recipe saved beside the checkpoint report the new one. + draft.modify(ExternalDraftConfig(**external_cfg)) # Fail before training: a vocab mismatch yields a plausible loss curve and a draft # with no acceptance. Equal sizes are not enough -- different tokenizers are # routinely padded to the same width -- so also check that the two agree on the ids. diff --git a/tests/unit/examples/test_speculative_decoding_main.py b/tests/unit/examples/test_speculative_decoding_main.py index e6e86e75a97..5e2e8825b33 100644 --- a/tests/unit/examples/test_speculative_decoding_main.py +++ b/tests/unit/examples/test_speculative_decoding_main.py @@ -170,3 +170,29 @@ def test_dry_run_is_exempt_from_the_data_source_requirement(main_mod, stub_token recipe, _training_args(), base_model=base, draft=_tiny_causal_lm(), dry_run=True ) assert draft is not None + + +def test_resume_honours_an_edited_external_section(main_mod, stub_tokenizers, tmp_path): + """An edited `external:` section must take effect on resume, not be silently dropped. + + The restore replays the config saved in the checkpoint, so skipping convert also + skips modify(). Resuming with a changed objective would then train the old one while + the printed recipe reports the new -- a plausible loss curve and the wrong draft. + """ + draft = _tiny_causal_lm() + mtsp.convert(draft, [("external", {"external_loss": "soft_ce"})]) + draft.save_pretrained(tmp_path / "ckpt") + restored = transformers.AutoModelForCausalLM.from_pretrained(tmp_path / "ckpt") + assert restored.external_loss == "soft_ce", "checkpoint did not carry the objective" + + recipe = _recipe(offline_data_path="/tmp/hidden") + recipe.external.external_loss = "tvd_deploy" + out, _ = main_mod.setup_external_draft( + recipe, + _training_args(), + base_model=_tiny_causal_lm(hidden=2 * HIDDEN), + draft=restored, + ) + assert out.external_loss == "tvd_deploy", ( + f"resume trained {out.external_loss!r} while the recipe asked for 'tvd_deploy'" + ) From 6829ae71833153d6483eb1044734d33a42229413 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 15 Oct 2026 16:27:55 +0000 Subject: [PATCH 4/5] Pin offline loading for the external draft, reject cp_size, reuse the base tokenizer The fresh path still derived use_offline_training from data.mode. The shipped recipe carries no data-source field, so under --dry_run it resolves to "online", load_vlm_or_llm ignores use_fake_base_for_offline, and every base decoder layer materialises on every rank -- the same inversion already fixed on the resume path. Also reject cp_size > 1, which this mode has no implementation for but was silently accepting, and let the fresh path hand over the base tokenizer it has already loaded instead of fetching it twice. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- examples/speculative_decoding/main.py | 43 ++++++++++++++----- .../test_speculative_decoding_main.py | 40 ++++++++++++++++- 2 files changed, 71 insertions(+), 12 deletions(-) diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 05292309e72..86ff02e4983 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -174,7 +174,9 @@ def _is_hf_format_checkpoint(checkpoint: str | None) -> bool: return any(os.path.isfile(os.path.join(checkpoint, f)) for f in hf_files) -def setup_external_draft(recipe, training_args, base_model, draft=None, dry_run=False): +def setup_external_draft( + recipe, training_args, base_model, draft=None, dry_run=False, base_tokenizer=None +): """Convert an external draft and wire it to the base model. Runs on both the fresh-load and the HF-resume path. The base lm_head is never part @@ -193,6 +195,8 @@ def setup_external_draft(recipe, training_args, base_model, draft=None, dry_run= draft: the already-loaded draft on the resume path; loaded here when None. dry_run: recipe validation only; skips the data-source requirement, as the ``data.data_path`` check in :func:`train` does. + base_tokenizer: the base's tokenizer when the caller already loaded it; the fresh + path has it, the resume path does not and leaves this None. """ if recipe.draft_model_name_or_path is None: raise ValueError( @@ -211,6 +215,12 @@ def setup_external_draft(recipe, training_args, base_model, draft=None, dry_run= "cached base outputs, which an online run does not produce. Set " "data.offline_data_path." ) + if training_args.cp_size > 1: + raise ValueError( + "external draft training does not support context parallelism: the draft is a " + "standalone model scored against cached base outputs, and nothing in this mode " + f"shards the sequence. Got cp_size={training_args.cp_size}; set it to 1." + ) base_vocab_size = base_model.config.vocab_size if draft is None: draft = transformers.AutoModelForCausalLM.from_pretrained( @@ -238,10 +248,11 @@ def setup_external_draft(recipe, training_args, base_model, draft=None, dry_run= # with no acceptance. Equal sizes are not enough -- different tokenizers are # routinely padded to the same width -- so also check that the two agree on the ids. draft.validate_against_base(base_vocab_size) - base_tokenizer = transformers.AutoTokenizer.from_pretrained( - recipe.model.model_name_or_path, - trust_remote_code=recipe.model.trust_remote_code, - ) + if base_tokenizer is None: + base_tokenizer = transformers.AutoTokenizer.from_pretrained( + recipe.model.model_name_or_path, + trust_remote_code=recipe.model.trust_remote_code, + ) draft.validate_tokenizer_against_base(tokenizer, base_tokenizer) # The norm is not always at base_model.model.norm -- FakeBaseModel keeps it at .norm # -- so search the same paths the other offline modes use. It stays optional; @@ -298,6 +309,14 @@ def train(): checkpoint = training_args.resume_from_checkpoint or last_checkpoint use_offline_training = recipe.data.mode != "online" + if isinstance(recipe, ModelOptExternalDraftRecipe): + # This mode has no online path (see setup_external_draft), and the shipped recipe + # carries no data-source field, so mode resolves to "online" under --dry_run. + # Deriving the flag from it there would make load_vlm_or_llm ignore + # use_fake_base_for_offline and materialise every base decoder layer on every + # rank, and would route a VLM base to the real multimodal model, whose parent + # config has no vocab_size for validate_against_base to read. + use_offline_training = True # Resume path depends on the existing checkpoint's on-disk format: consolidated HF # weights load via from_pretrained; FSDP sharded checkpoints load the base model and @@ -326,14 +345,12 @@ def train(): "model.model_name_or_path must be set to resume an external-draft run: " "the base model supplies the lm_head the draft is scored against." ) - # Same arguments as the fresh path below. use_offline_training is True - # unconditionally because this mode has no online path -- and without it - # load_vlm_or_llm ignores use_fake_base and materialises every decoder - # layer, on every rank, to borrow an lm_head and a final norm. + # Same arguments as the fresh path below; use_offline_training is pinned + # True for this recipe above. base_model = load_vlm_or_llm( recipe.model.model_name_or_path, use_fake_base=recipe.model.use_fake_base_for_offline, - use_offline_training=True, + use_offline_training=use_offline_training, dtype="auto", device_map="cpu", trust_remote_code=recipe.model.trust_remote_code, @@ -392,7 +409,11 @@ def train(): # This mode inverts the usual relationship: we train a pretrained draft, and # `model` loaded above is the base, kept only for its lm_head and vocabulary. model, tokenizer = setup_external_draft( - recipe, training_args, base_model=model, dry_run=dry_run + recipe, + training_args, + base_model=model, + dry_run=dry_run, + base_tokenizer=tokenizer, ) else: raise ValueError(f"Unsupported speculative recipe type: {type(recipe).__name__}") diff --git a/tests/unit/examples/test_speculative_decoding_main.py b/tests/unit/examples/test_speculative_decoding_main.py index 5e2e8825b33..5d633293296 100644 --- a/tests/unit/examples/test_speculative_decoding_main.py +++ b/tests/unit/examples/test_speculative_decoding_main.py @@ -101,7 +101,7 @@ def stub_tokenizers(monkeypatch, main_mod): def _training_args(): - return types.SimpleNamespace(training_seq_len=128, device=torch.device("cpu")) + return types.SimpleNamespace(training_seq_len=128, device=torch.device("cpu"), cp_size=1) def test_resume_of_a_converted_checkpoint_does_not_reconvert(main_mod, stub_tokenizers, tmp_path): @@ -196,3 +196,41 @@ def test_resume_honours_an_edited_external_section(main_mod, stub_tokenizers, tm assert out.external_loss == "tvd_deploy", ( f"resume trained {out.external_loss!r} while the recipe asked for 'tvd_deploy'" ) + + +def test_context_parallelism_is_rejected(main_mod, stub_tokenizers): + """Nothing in this mode shards the sequence, so cp_size > 1 must refuse, not run. + + main.py sets model.eagle_cp_size unconditionally further down, so without this the + draft trains with CP configured and silently ignored. + """ + args = _training_args() + args.cp_size = 2 + with pytest.raises(ValueError, match="context parallelism"): + main_mod.setup_external_draft( + _recipe(offline_data_path="/tmp/hidden"), + args, + base_model=_tiny_causal_lm(hidden=2 * HIDDEN), + draft=_tiny_causal_lm(), + ) + + +def test_caller_supplied_base_tokenizer_is_used(main_mod, monkeypatch): + """The fresh path already holds the base tokenizer; it must not be re-fetched.""" + calls = [] + + def _from_pretrained(cls, *a, **k): + calls.append(a[0] if a else None) + return _FakeTok() + + monkeypatch.setattr( + main_mod.transformers.AutoTokenizer, "from_pretrained", classmethod(_from_pretrained) + ) + main_mod.setup_external_draft( + _recipe(offline_data_path="/tmp/hidden"), + _training_args(), + base_model=_tiny_causal_lm(hidden=2 * HIDDEN), + draft=_tiny_causal_lm(), + base_tokenizer=_FakeTok(), + ) + assert calls == ["draft"], f"base tokenizer re-loaded: {calls}" From e4fa8a2719c705a564397e5844a2cac2d1f02678 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 15 Oct 2026 18:19:19 +0000 Subject: [PATCH 5/5] Keep the stored modelopt config in step with a re-applied recipe modify() writes only the live attributes, so a resume with an edited external: section trained the new objective while the next checkpoint still recorded the old one -- and a restore outside main.py handed back that stale objective. Update the state manager's config alongside modify(). Also refuse a draft restored from a different speculative mode: the guard asked only whether some speculative mode was present, so an eagle or dflash checkpoint skipped the convert and failed later on an unrelated attribute. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- examples/speculative_decoding/main.py | 28 ++++++++++++----- .../test_speculative_decoding_main.py | 30 +++++++++++++++++++ 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 86ff02e4983..0d01cd81c48 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -57,7 +57,7 @@ ModelOptMedusaRecipe, ModelOptSpeculativeRecipeBase, ) -from modelopt.torch.opt.conversion import ModeloptStateManager +from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state from modelopt.torch.speculative.config import ExternalDraftConfig from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel from modelopt.torch.speculative.plugins.hf_domino import DominoLambdaCallback @@ -236,14 +236,26 @@ def setup_external_draft( model_max_length=training_args.training_seq_len, trust_remote_code=recipe.model.trust_remote_code, ) - if not ModeloptStateManager.has_state_for_mode_type("speculative", model=draft): - mtsp.convert(draft, [("external", external_cfg)]) + restored_modes = {m for m, _ in modelopt_state(draft)["modelopt_state_dict"]} + if "external" in restored_modes: + # The restore replayed the config stored in the checkpoint, so modify() has already + # run with the values this draft was last trained under. Re-apply the recipe's: + # otherwise an edited external_loss / top_k / top_p is dropped while the banner and + # the recipe saved beside the checkpoint report the new one. The manager's config is + # updated alongside it because modify() writes only the live attributes -- without + # this the next save_pretrained records the objective the draft was never trained + # under, and a restore outside this script hands back that stale objective. + external_config = ExternalDraftConfig(**external_cfg) + draft.modify(external_config) + ModeloptStateManager(draft)._last_config = external_config + elif ModeloptStateManager.has_state_for_mode_type("speculative", model=draft): + raise ValueError( + f"the draft checkpoint carries speculative mode(s) {sorted(restored_modes)}, not " + "'external'. Skipping the convert would leave it as that mode's model type and " + "fail later with an unrelated error, so refuse here." + ) else: - # The restore replayed the config stored in the checkpoint, so modify() has - # already run with the values this draft was last trained under. Re-apply the - # recipe's: otherwise an edited external_loss / top_k / top_p is dropped while - # the banner and the recipe saved beside the checkpoint report the new one. - draft.modify(ExternalDraftConfig(**external_cfg)) + mtsp.convert(draft, [("external", external_cfg)]) # Fail before training: a vocab mismatch yields a plausible loss curve and a draft # with no acceptance. Equal sizes are not enough -- different tokenizers are # routinely padded to the same width -- so also check that the two agree on the ids. diff --git a/tests/unit/examples/test_speculative_decoding_main.py b/tests/unit/examples/test_speculative_decoding_main.py index 5d633293296..13cd3fc3597 100644 --- a/tests/unit/examples/test_speculative_decoding_main.py +++ b/tests/unit/examples/test_speculative_decoding_main.py @@ -197,6 +197,15 @@ def test_resume_honours_an_edited_external_section(main_mod, stub_tokenizers, tm f"resume trained {out.external_loss!r} while the recipe asked for 'tvd_deploy'" ) + # modify() writes only the live attributes, so the stored state has to be updated too: + # otherwise the next checkpoint records the objective the draft was never trained under + # and a restore outside main.py hands that back. + out.save_pretrained(tmp_path / "out") + again = transformers.AutoModelForCausalLM.from_pretrained(tmp_path / "out") + assert again.external_loss == "tvd_deploy", ( + f"checkpoint recorded {again.external_loss!r}; the draft was trained with 'tvd_deploy'" + ) + def test_context_parallelism_is_rejected(main_mod, stub_tokenizers): """Nothing in this mode shards the sequence, so cp_size > 1 must refuse, not run. @@ -234,3 +243,24 @@ def _from_pretrained(cls, *a, **k): base_tokenizer=_FakeTok(), ) assert calls == ["draft"], f"base tokenizer re-loaded: {calls}" + + +def test_resuming_a_different_speculative_mode_is_refused(main_mod, stub_tokenizers, tmp_path): + """A checkpoint from another speculative mode must fail with a clear message. + + The skip-convert guard only asks whether *some* speculative mode is present, so without + this an eagle/dflash checkpoint would skip the convert, never become an external draft, + and die later on an unrelated attribute. + """ + draft = _tiny_causal_lm() + mtsp.convert(draft, [("eagle", {})]) + draft.save_pretrained(tmp_path / "ckpt") + restored = transformers.AutoModelForCausalLM.from_pretrained(tmp_path / "ckpt") + + with pytest.raises(ValueError, match="not 'external'"): + main_mod.setup_external_draft( + _recipe(offline_data_path="/tmp/hidden"), + _training_args(), + base_model=_tiny_causal_lm(hidden=2 * HIDDEN), + draft=restored, + )