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..0d01cd81c48 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -53,15 +53,19 @@ from modelopt.recipe.config import ( ModelOptDFlashRecipe, ModelOptEagleRecipe, + ModelOptExternalDraftRecipe, ModelOptMedusaRecipe, ModelOptSpeculativeRecipeBase, ) +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 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 +174,120 @@ 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, 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 + 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. + 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( + "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." + ) + 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( + 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, + ) + 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: + 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) + 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; + # _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}." ) @@ -209,6 +321,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 @@ -227,6 +347,30 @@ 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 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=use_offline_training, + 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 +417,16 @@ 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, + base_tokenizer=tokenizer, + ) else: raise ValueError(f"Unsupported speculative recipe type: {type(recipe).__name__}") @@ -306,6 +460,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..13cd3fc3597 --- /dev/null +++ b/tests/unit/examples/test_speculative_decoding_main.py @@ -0,0 +1,266 @@ +# 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"), cp_size=1) + + +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 + + +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'" + ) + + # 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. + + 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}" + + +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, + ) 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