From 35247e34b57cece91ebfa0047dadefa03afc81d2 Mon Sep 17 00:00:00 2001 From: Peter Hicks <60696413+phicksy@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:37:48 +0100 Subject: [PATCH 1/4] Update examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Peter Hicks <60696413+phicksy@users.noreply.github.com> --- .../collect_hidden_states/compute_hidden_states_hf.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 7d23ee4bcc1..047dfb28863 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -152,7 +152,11 @@ def keep_conversation(entry): ) # Multimodal checkpoints (e.g. *ForConditionalGeneration) keep the language # model's depth under config.text_config. - text_config = getattr(model.config, "text_config", None) or model.config + text_config = ( + getattr(model.config, "text_config", None) + or getattr(model.config, "llm_config", None) + or model.config + ) num_hidden_layers = getattr(text_config, "num_hidden_layers", None) if num_hidden_layers is None: raise ValueError(f"model.config has no 'num_hidden_layers' attribute: {model.config}") From b2faa5274f8cc9eb3a4dd16c7f72af0c4d7a4ebb Mon Sep 17 00:00:00 2001 From: Peter Hicks <60696413+phicksy@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:38:03 +0100 Subject: [PATCH 2/4] Update examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Peter Hicks <60696413+phicksy@users.noreply.github.com> --- .../collect_hidden_states/compute_hidden_states_hf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 047dfb28863..3b51994b52d 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -159,7 +159,7 @@ def keep_conversation(entry): ) num_hidden_layers = getattr(text_config, "num_hidden_layers", None) if num_hidden_layers is None: - raise ValueError(f"model.config has no 'num_hidden_layers' attribute: {model.config}") + raise ValueError(f"config has no 'num_hidden_layers' attribute: {text_config}") selected_layer_ids = resolve_aux_layers(args, num_hidden_layers) tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) From c00021c53cbbe54966d8a868bb3acdf5e7cfb087 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Tue, 13 Oct 2026 12:40:17 +0000 Subject: [PATCH 3/4] Address review suggestions on the aux-hidden-states dump path - Skip building the aux planes under --no-aux-hidden-states: the concatenation and device-to-host copy ran per conversation on a dump about to drop them. - Document aux_hidden_states_optional beside its sibling arguments, and record in the dump-format key list that the key is absent under --no-aux-hidden-states. That list is the contract a new producer or consumer is written against. - Let sample_hidden_states inspect dumps that omit an optional key. It compared the whole key set, so it reported every --no-aux-hidden-states dump as damaged, and every --answer-only-loss dump along with it. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- .../compute_hidden_states_hf.py | 12 +++++++++-- .../sample_hidden_states.py | 20 +++++++++++-------- modelopt/torch/speculative/eagle/utils.py | 16 +++++++-------- 3 files changed, 30 insertions(+), 18 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 3b51994b52d..21e1b652527 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -193,8 +193,16 @@ async def dump_hidden_states( outputs = model(input_ids=input_ids.to(model.device), output_hidden_states=True) # outputs.hidden_states[0] is the embedding output; layer k output is at index k+1. hidden_states = outputs.hidden_states - aux_hidden_states = torch.cat( - [hidden_states[lid + 1].squeeze(0).cpu() for lid in selected_layer_ids], dim=-1 + # Only build the aux planes when they will be stored: assembling them costs a + # concatenation and a device-to-host copy per conversation, which is pure waste + # on a dump that is about to drop them. + aux_hidden_states = ( + None + if args.no_aux_hidden_states + else torch.cat( + [hidden_states[lid + 1].squeeze(0).cpu() for lid in selected_layer_ids], + dim=-1, + ) ) output_hidden_states = hidden_states[-1].squeeze(0).cpu() # Write-then-rename: keep_conversation() skips any existing .pt as finished work, diff --git a/examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py b/examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py index 75a88969d38..618ef1f6597 100644 --- a/examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py +++ b/examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py @@ -61,18 +61,22 @@ def main(args: argparse.Namespace) -> None: for i, file in enumerate(sampled_files): data = torch.load(file) - expected_keys = [ - "input_ids", - "hidden_states", - "aux_hidden_states", - "conversation_id", - ] - if set(expected_keys) != set(data.keys()): - print(f"File {file} does not contain all expected keys: {expected_keys}") + # Require only the keys every dump carries, and report what is missing rather + # than comparing the whole set: the optional keys are legitimately absent on + # some dumps -- ``aux_hidden_states`` under ``--no-aux-hidden-states``, and + # ``loss_mask`` unless ``--answer-only-loss`` was used -- and an exact-set + # check rejects those as malformed. + required_keys = ["input_ids", "hidden_states", "conversation_id"] + missing = [k for k in required_keys if k not in data] + if missing: + print(f"File {file} is missing required keys: {missing}") print(f" Found keys: {list(data.keys())}") continue print(f"Sample {i + 1}: {file.name}") for key in ["input_ids", "hidden_states", "aux_hidden_states"]: + if key not in data: + print(f"{key}: absent") + continue print(f"{key}: {data[key].shape} {data[key].dtype} {data[key].device}") print(f"conversation_id: {data['conversation_id']}") input_ids_list = data["input_ids"].tolist() diff --git a/modelopt/torch/speculative/eagle/utils.py b/modelopt/torch/speculative/eagle/utils.py index c00b691c083..4528bbfe5b5 100644 --- a/modelopt/torch/speculative/eagle/utils.py +++ b/modelopt/torch/speculative/eagle/utils.py @@ -90,7 +90,9 @@ class OfflineSupervisedDataset(Dataset): - ``input_ids``: token IDs of shape ``(seq_len,)`` - ``hidden_states``: base model last hidden states of shape ``(seq_len, hidden_size)`` - - ``aux_hidden_states``: auxiliary hidden states of shape ``(seq_len, hidden_size)`` + - ``aux_hidden_states``: auxiliary hidden states of shape ``(seq_len, hidden_size)``. + Absent from dumps written with ``--no-aux-hidden-states``; reading one of those + requires ``aux_hidden_states_optional=True``. - ``loss_mask`` (optional): per-token assistant mask of shape ``(seq_len,)``, present when the dump was produced with ``--answer-only-loss``. @@ -106,6 +108,10 @@ class OfflineSupervisedDataset(Dataset): is stored in the file (backward compatible). tokenizer: Optional tokenizer used to recover the assistant mask for dumps that lack a stored ``loss_mask``. + aux_hidden_states_optional (bool): Tolerate dumps written with + ``--no-aux-hidden-states``, yielding an empty ``(seq_len, 0)`` plane in + their place. Only modes that never read the aux planes may set this; for + the others a missing key must still raise. """ def __init__( @@ -115,13 +121,7 @@ def __init__( tokenizer=None, aux_hidden_states_optional: bool = False, ): - """Initialize with a list of .pt file paths. - - Args: - aux_hidden_states_optional: tolerate dumps written with - ``--no-aux-hidden-states``. Only modes that never read the aux planes - may set this; for the others a missing key must still raise. - """ + """Initialize with a list of .pt file paths.""" super().__init__() self.dumped_files = dumped_files self.answer_only_loss = answer_only_loss From 657ef87962f68450950dc6e8ad83c2caefcdbd3e Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 1 Oct 2026 09:29:13 +0000 Subject: [PATCH 4/4] [2/8] Add the external draft speculative-decoding mode Unlike Eagle/Medusa/DFlash, the converted module *is* the draft: an ordinary pretrained causal LM that owns its embeddings and lm_head. The base model is never grafted onto it and is not modified; it contributes only cached hidden states, which the base lm_head projects to teacher logits. The trained draft therefore exports as a plain HF checkpoint that downstream runtimes load with no mode-specific support. external_loss selects the objective: soft_ce reproduces the Eagle offline loss, tvd is total-variation distance, and tvd_deploy puts both sides through the serving filter first so the loss is the acceptance the deployed pair sees. Part of a five-PR series; this slice is the library mode only. The recipe and training entrypoint follow in [3/5]. Signed-off-by: Peter Hicks Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- modelopt/torch/speculative/__init__.py | 2 +- modelopt/torch/speculative/config.py | 53 ++++ .../torch/speculative/external/__init__.py | 20 ++ .../torch/speculative/external/conversion.py | 72 +++++ .../speculative/external/external_model.py | 45 +++ modelopt/torch/speculative/mode.py | 31 +- .../torch/speculative/plugins/__init__.py | 1 + .../torch/speculative/plugins/hf_external.py | 264 ++++++++++++++++++ .../speculative/plugins/test_hf_external.py | 83 ++++++ .../torch/speculative/test_external_draft.py | 174 ++++++++++++ .../test_external_draft_offline.py | 162 +++++++++++ 11 files changed, 905 insertions(+), 2 deletions(-) create mode 100644 modelopt/torch/speculative/external/__init__.py create mode 100644 modelopt/torch/speculative/external/conversion.py create mode 100644 modelopt/torch/speculative/external/external_model.py create mode 100644 modelopt/torch/speculative/plugins/hf_external.py create mode 100644 tests/gpu/torch/speculative/plugins/test_hf_external.py create mode 100644 tests/unit/torch/speculative/test_external_draft.py create mode 100644 tests/unit/torch/speculative/test_external_draft_offline.py diff --git a/modelopt/torch/speculative/__init__.py b/modelopt/torch/speculative/__init__.py index 4e1334edbf7..dee7ae650c6 100644 --- a/modelopt/torch/speculative/__init__.py +++ b/modelopt/torch/speculative/__init__.py @@ -15,6 +15,6 @@ """Speculative Decoding Optimizations.""" -from . import eagle, medusa, mode, plugins +from . import eagle, external, medusa, mode, plugins from .config import * from .speculative_decoding import * diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index c3f42a5ad14..39c3e3fc782 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -30,6 +30,7 @@ "EAGLE_MTP_DEFAULT_CFG", "DFlashConfig", "EagleConfig", + "ExternalDraftConfig", "MedusaConfig", "eagle3_default_config", "eagle_mtp_default_config", @@ -468,6 +469,58 @@ class MedusaConfig(ModeloptBaseConfig): ) +class ExternalDraftConfig(ModeloptBaseConfig): + """External (standalone) draft model config. + + The converted module is a pretrained causal LM used directly as the draft. It + keeps its own embeddings and lm_head; the base model contributes only cached + hidden states, projected by the base lm_head to obtain teacher logits. + """ + + external_offline: bool = ModeloptField( + default=True, + description=( + "Whether the draft consumes pre-computed base hidden states instead of running " + "the base model in-process. Only the offline path is implemented." + ), + ) + + external_loss: Literal["soft_ce", "tvd", "tvd_deploy"] = ModeloptField( + default="tvd", + description=( + "Training objective. ``soft_ce`` matches the Eagle offline objective. " + "``tvd`` is total-variation distance between the two untruncated " + "distributions. ``tvd_deploy`` puts both through the serving filter first " + "(``external_top_k`` then ``external_top_p``), making it the acceptance the " + "deployed pair will see: ``alpha = 1 - TVD``." + ), + ) + + external_top_k: int = ModeloptField( + default=20, + ge=0, + description=( + "Top-k of the serving filter, applied to *both* the base and the " + "draft by ``external_loss='tvd_deploy'``. 0 disables the top-k stage." + ), + ) + + external_top_p: float = ModeloptField( + default=0.95, + ge=0.0, + le=1.0, + description=( + "Top-p of the serving filter, applied to *both* the base and the " + "draft by ``external_loss='tvd_deploy'``. 0 or 1.0 disables the nucleus stage." + ), + ) + + external_report_acc: bool = ModeloptField( + default=True, + description=("Whether to report top-1 agreement with the base model during training."), + ) + + class EagleConfig(ModeloptBaseConfig): """Eagle config.""" diff --git a/modelopt/torch/speculative/external/__init__.py b/modelopt/torch/speculative/external/__init__.py new file mode 100644 index 00000000000..7af203db446 --- /dev/null +++ b/modelopt/torch/speculative/external/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + + +"""External Draft Optimization Method.""" + +from .conversion import * +from .external_model import * diff --git a/modelopt/torch/speculative/external/conversion.py b/modelopt/torch/speculative/external/conversion.py new file mode 100644 index 00000000000..6c40e91a5e1 --- /dev/null +++ b/modelopt/torch/speculative/external/conversion.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""External draft conversion/restore utilities.""" + +from torch import nn + +from modelopt.torch.opt.conversion import ModelLikeModule +from modelopt.torch.opt.dynamic import _DMRegistryCls +from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict + +from ..config import ExternalDraftConfig + +ExternalDraftDMRegistry = _DMRegistryCls(prefix="ExternalDraft") # global instance for the registry + +__all__ = [ + "ExternalDraftDMRegistry", + "convert_to_external_draft_model", + "restore_external_draft_model", +] + + +def convert_to_external_draft_model( + model: nn.Module, config: ExternalDraftConfig +) -> ConvertReturnType: + """Convert a pretrained causal LM into a trainable external draft as per `config`. + + Note that, unlike the other speculative modes, ``model`` here is the *draft* + rather than the base model. Nothing is grafted onto a target; the draft keeps + its own embeddings and lm_head. + """ + # initialize the true module if necessary + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + + original_cls = type(model) + if original_cls not in ExternalDraftDMRegistry: + for cls in ExternalDraftDMRegistry._registry: + if issubclass(original_cls, cls): + ExternalDraftDMRegistry.register({original_cls: "draft_model_class"})( + ExternalDraftDMRegistry[cls] + ) + break + + external_draft_model = ExternalDraftDMRegistry.convert(model) + external_draft_model.modify(config) + + # no metadata, all specified via config. + metadata = {} + + return external_draft_model, metadata + + +def restore_external_draft_model( + model: nn.Module, config: ExternalDraftConfig, metadata: MetadataDict +) -> nn.Module: + """Function for restoring a previously converted model to an external draft model.""" + # the metadata should be empty + assert not metadata, "No metadata expected!" + + return convert_to_external_draft_model(model, config)[0] diff --git a/modelopt/torch/speculative/external/external_model.py b/modelopt/torch/speculative/external/external_model.py new file mode 100644 index 00000000000..2ca706d1a46 --- /dev/null +++ b/modelopt/torch/speculative/external/external_model.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""External draft model to support speculative decoding with a standalone draft.""" + +from modelopt.torch.opt.dynamic import DynamicModule + +__all__ = ["ExternalDraftModel"] + + +class ExternalDraftModel(DynamicModule): + """Base external draft model. + + Unlike Eagle/Medusa/DFlash, the converted module *is* the draft: a pretrained + causal LM that owns its embeddings and lm_head. The base (target) model is + never grafted onto it; it contributes only cached hidden states, which are + projected by the target lm_head to obtain teacher logits. + """ + + def _setup(self): + self._register_temp_attribute("external_offline", True) + self._register_temp_attribute("external_loss", "tvd") + self._register_temp_attribute("external_report_acc", True) + self._register_temp_attribute("external_top_k", 20) + self._register_temp_attribute("external_top_p", 0.95) + + def modify(self, config): + """Base external draft modify function. Child class should implement the details.""" + self.external_offline = config.external_offline + self.external_loss = config.external_loss + self.external_report_acc = config.external_report_acc + self.external_top_k = config.external_top_k + self.external_top_p = config.external_top_p diff --git a/modelopt/torch/speculative/mode.py b/modelopt/torch/speculative/mode.py index ae965354a92..d393caf89b6 100644 --- a/modelopt/torch/speculative/mode.py +++ b/modelopt/torch/speculative/mode.py @@ -23,9 +23,10 @@ _ModeRegistryCls, ) -from .config import DFlashConfig, EagleConfig, MedusaConfig +from .config import DFlashConfig, EagleConfig, ExternalDraftConfig, MedusaConfig from .dflash.conversion import convert_to_dflash_model, restore_dflash_model from .eagle.conversion import convert_to_eagle_model, restore_eagle_model +from .external.conversion import convert_to_external_draft_model, restore_external_draft_model from .medusa.conversion import convert_to_medusa_model, restore_medusa_model SpeculativeDecodingModeRegistry = _ModeRegistryCls("speculative") @@ -113,3 +114,31 @@ def convert(self) -> ConvertEntrypoint: def restore(self) -> RestoreEntrypoint: """The mode's entrypoint for restoring a model.""" return restore_eagle_model + + +@SpeculativeDecodingModeRegistry.register_mode +class ExternalDraftModeDescriptor(ModeDescriptor): + """Class to describe the ``"external"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "external" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return ExternalDraftConfig + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return convert_to_external_draft_model + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_external_draft_model diff --git a/modelopt/torch/speculative/plugins/__init__.py b/modelopt/torch/speculative/plugins/__init__.py index bbdcac7553d..5f1c7a340ab 100644 --- a/modelopt/torch/speculative/plugins/__init__.py +++ b/modelopt/torch/speculative/plugins/__init__.py @@ -35,5 +35,6 @@ from .hf_domino import * from .hf_dspark import * from .hf_eagle import * + from .hf_external import * from .hf_lilicorr import * from .hf_medusa import * diff --git a/modelopt/torch/speculative/plugins/hf_external.py b/modelopt/torch/speculative/plugins/hf_external.py new file mode 100644 index 00000000000..0a77d4fb7ee --- /dev/null +++ b/modelopt/torch/speculative/plugins/hf_external.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Support for training a standalone (external) draft model with HuggingFace. + +Unlike Eagle/Medusa/DFlash, the draft here is an ordinary pretrained causal LM +that owns its embeddings and lm_head. The base model is not modified and nothing +is grafted onto it: it contributes only per-token hidden states, which are +projected by the base lm_head to obtain teacher logits. + +``external_loss`` selects the objective: ``soft_ce`` matches the Eagle offline +objective, ``tvd`` is total-variation distance between the two distributions, and +``tvd_deploy`` applies the serving filter to both sides first so the loss is the +acceptance the deployed pair sees. +""" + +import json +from pathlib import Path + +import torch +from transformers import PreTrainedModel + +from ..external.conversion import ExternalDraftDMRegistry +from ..external.external_model import ExternalDraftModel +from .modeling_final_norm import _maybe_apply_base_final_norm + +__all__ = ["HFExternalDraftModel"] + + +def _train_acc(acc): + """Shape a scalar accuracy the way ``EagleTrainerWithAccLog`` consumes it. + + The callback averages the collected values and then iterates them as + ``[parallel_draft][ttt_step]``, so a scalar has to arrive nested twice over; a + flatter value makes that iteration raise. This mode has one draft and one step. + """ + return [[acc.item()]] + + +@ExternalDraftDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) +class HFExternalDraftModel(ExternalDraftModel): + """An external draft model for HuggingFace models.""" + + def modify(self, config): + """Configure the draft and validate that it can score against this base.""" + super().modify(config) + + # Checked here rather than on the config: the field is part of modelopt_state, + # so rejecting it in the schema would make an already-saved checkpoint carrying + # the old default unrestorable. modify() runs on both convert and restore. + if not self.external_offline: + raise ValueError( + "external_offline=False is not supported: this mode trains against cached " + "base hidden states or a stored base policy, never a live base model." + ) + + if self.external_loss not in ("soft_ce", "tvd", "tvd_deploy"): + raise ValueError( + f"external_loss must be one of 'soft_ce', 'tvd', 'tvd_deploy', " + f"got {self.external_loss!r}." + ) + + def attach_base_lm_head(self, base_lm_head, base_final_norm=None) -> None: + """Attach the frozen base lm_head used to project cached hidden states to teacher logits.""" + for p in base_lm_head.parameters(): + p.requires_grad = False + if base_final_norm is not None: + for p in base_final_norm.parameters(): + p.requires_grad = False + object.__setattr__(self, "_base_lm_head", base_lm_head) + object.__setattr__(self, "_base_final_norm", base_final_norm) + + def _teacher_logits(self, hidden_states, base_model_outputs): + """Project dumped base hidden states through the base lm_head.""" + if getattr(self, "_base_lm_head", None) is None: + raise RuntimeError( + "Base lm_head is not attached; call attach_base_lm_head() before training. " + "Teacher logits cannot be derived from the dumped hidden states without it." + ) + h = hidden_states + # The base modules are held off the draft's module tree so they stay out of + # its state_dict, which also means .to(device) never reaches them. + if self._base_lm_head.weight.device != h.device: + self._base_lm_head.to(h.device) + if getattr(self, "_base_final_norm", None) is not None: + self._base_final_norm.to(h.device) + # Shared with the EAGLE/DFlash offline forwards: raises rather than feeding an + # un-normed hidden into lm_head when the producer captured a pre-norm hidden. + h = _maybe_apply_base_final_norm( + h, base_model_outputs, getattr(self, "_base_final_norm", None) + ) + h = h.to(self._base_lm_head.weight.dtype) + with torch.no_grad(): + return self._base_lm_head(h) + + def forward( + self, + input_ids=None, + attention_mask=None, + loss_mask=None, + base_model_outputs=None, + labels=None, + **kwargs, + ): + """Score the draft against cached base hidden states. + + Only ``base_model_hidden_states`` is consumed; ``aux_hidden_states`` exists in the + dump for the head-based modes and is not used here. + """ + if base_model_outputs is None: + raise ValueError( + "External draft training expects offline base_model_outputs in the batch; " + "online training against a live base model is not supported." + ) + draft_out = super().forward(input_ids=input_ids, attention_mask=attention_mask, **kwargs) + base_logits = self._teacher_logits( + base_model_outputs["base_model_hidden_states"], base_model_outputs + ) + if loss_mask is None: + # Derive from attention_mask so padding is not trained on, matching the + # other speculative plugins; all-ones only when there is no mask at all. + loss_mask = ( + attention_mask.clone().detach() + if attention_mask is not None + else torch.ones_like(input_ids) + ).to(draft_out.logits.dtype) + loss, acc = self.compute_loss(draft_out.logits, base_logits, loss_mask) + draft_out.loss = loss + if acc is not None: + draft_out.train_acc = _train_acc(acc) + return draft_out + + def save_pretrained(self, save_directory, *args, **kwargs): + """Save as an ordinary HF checkpoint. + + ``transformers`` writes ``type(self).__name__`` into ``config.architectures``, + which after conversion is the registry's dynamic subclass name (for example + ``ExternalDraftLlamaForCausalLM``). Nothing downstream resolves that: vLLM and + TRT-LLM look the string up in their own model registries and fail to load the + checkpoint. Conversion adds no modules to the draft -- it is still the pretrained + causal LM it started as -- so the original class name is the accurate one, and + restoring it is what keeps the promise that this mode exports a plain HF model. + """ + result = super().save_pretrained(save_directory, *args, **kwargs) + original = self.original_cls.__name__ + self.config.architectures = [original] + # transformers rewrites architectures during the save, so the correction has to + # come after it. Only the rank that wrote the file may rewrite it: under FSDP + # every rank enters save_pretrained, and a read-modify-write from all of them + # races on a shared filesystem. + if kwargs.get("is_main_process", True): + config_file = Path(save_directory) / "config.json" + if config_file.exists(): + config = json.loads(config_file.read_text()) + if config.get("architectures") != [original]: + config["architectures"] = [original] + config_file.write_text(json.dumps(config, indent=2, sort_keys=True) + "\n") + return result + + def validate_against_base(self, base_vocab_size: int) -> None: + """Check the draft indexes the same vocabulary as the base model. + + A mismatch cannot be caught downstream: training would optimise against + misaligned columns and yield a plausible loss curve with no acceptance. + """ + draft_vocab_size = int(self.config.vocab_size) + if draft_vocab_size != int(base_vocab_size): + raise ValueError( + "External draft and base model must share a vocabulary: draft " + f"vocab_size={draft_vocab_size}, base vocab_size={int(base_vocab_size)}." + ) + + def validate_tokenizer_against_base(self, draft_tokenizer, base_tokenizer, sample: int = 256): + """Check the two tokenizers agree on what the ids mean, not just how many there are. + + Equal ``vocab_size`` is not enough: different tokenizers are routinely padded to + the same width, and training against misaligned columns yields a falling loss + curve and no acceptance. + """ + base_vocab, draft_vocab = base_tokenizer.get_vocab(), draft_tokenizer.get_vocab() + checked = mismatched = 0 + for token, base_id in base_vocab.items(): + if base_id >= sample: + continue + checked += 1 + if draft_vocab.get(token) != base_id: + mismatched += 1 + if checked and mismatched: + raise ValueError( + f"External draft and base tokenizers disagree on {mismatched}/{checked} of " + "the first ids despite matching vocab_size. They are different tokenizers, " + "and training would optimise against misaligned columns." + ) + + @staticmethod + def _deployment_probs(logits, top_k, top_p): + """Probabilities after temperature-1 top-k then top-p, renormalised.""" + p = torch.softmax(logits.float(), dim=-1) + if top_k and top_k < p.shape[-1]: + kth = p.topk(top_k, dim=-1).values[..., -1:] + p = p * (p >= kth) + # Renormalise before the nucleus: serving applies top-p to the + # redistributed top-k mass, so skipping this makes top-p a no-op + # whenever the top-k mass is already below it. + p = p / p.sum(-1, keepdim=True).clamp_min(1e-9) + if top_p and top_p < 1.0: + srt, idx = p.sort(dim=-1, descending=True) + keep = (srt.cumsum(-1) - srt) < top_p + p = p * torch.zeros_like(keep).scatter(-1, idx, keep) + return p / p.sum(-1, keepdim=True).clamp_min(1e-9) + + def compute_loss(self, draft_logits, base_logits, loss_mask, base_predict_tok=None): + """Compute the configured objective against the base model's distribution. + + Args: + loss_mask: per-token mask; response-only when the dump used ``--answer-only-loss``. + + Returns: + ``(loss, accuracy)``; accuracy is ``None`` unless ``external_report_acc`` is set. + """ + seq_len = draft_logits.shape[1] + base_logits = base_logits[:, :seq_len] + mask = loss_mask[:, :seq_len].to(draft_logits.dtype) + denom = mask.sum() + 1e-5 + + if self.external_loss in ("tvd", "tvd_deploy"): + if self.external_loss == "tvd_deploy": + # Both sides through the serving filter, so the loss is the + # acceptance the deployed pair will actually see. + p = self._deployment_probs(base_logits, self.external_top_k, self.external_top_p) + q = self._deployment_probs(draft_logits, self.external_top_k, self.external_top_p) + else: + p = torch.softmax(base_logits.float(), dim=-1) + q = torch.softmax(draft_logits.float(), dim=-1) + loss = (0.5 * (p - q).abs().sum(-1) * mask).sum() / denom + else: + # Identical in form to hf_eagle._eagle_loss so switching modes does not + # silently change the objective. + base_softmax = torch.softmax(base_logits.float(), dim=-1) + draft_logsoft = torch.log_softmax(draft_logits.float(), dim=-1) + loss = -torch.sum(mask.unsqueeze(-1) * base_softmax * draft_logsoft) / denom + + accuracy = None + if self.external_report_acc: + with torch.no_grad(): + if base_predict_tok is None: + base_predict_tok = base_logits.argmax(dim=-1) + valid = mask.bool() + correct = (base_predict_tok == draft_logits.detach().argmax(dim=-1)) & valid + accuracy = correct.sum().float() / valid.sum().clamp_min(1).float() + + return loss, accuracy diff --git a/tests/gpu/torch/speculative/plugins/test_hf_external.py b/tests/gpu/torch/speculative/plugins/test_hf_external.py new file mode 100644 index 00000000000..2bd7e1705cb --- /dev/null +++ b/tests/gpu/torch/speculative/plugins/test_hf_external.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""GPU tests for the external (standalone) draft plugin. + +These tests require a CUDA GPU. CPU-only tests are in tests/unit/. + +The base lm_head is attached with ``object.__setattr__`` so it stays off the draft's +module tree and out of its optimiser and gradients. The cost of that is that +``model.cuda()`` does not move it, which no CPU test can observe: every tensor is on +the same device there. That is the class of bug these tests exist for. +""" + +import pytest +import torch +from _test_utils.torch.transformers_models import get_tiny_llama + +import modelopt.torch.speculative as mtsp + +SEQ_LEN = 16 + + +def _external_draft(**overrides): + model = get_tiny_llama() + mtsp.convert( + model, + [("external", {"external_offline": True, "external_loss": "tvd_deploy", **overrides})], + ) + return model + + +class TestExternalDraftDevicePlacement: + """The base lm_head lives off the module tree, so it needs its own device handling.""" + + def test_teacher_logits_with_base_head_left_on_cpu(self): + """Attaching a CPU head then moving the draft to GPU must still train.""" + draft = _external_draft() + base = get_tiny_llama() + # Deliberately attach before .cuda(), the order main.py uses. + draft.attach_base_lm_head( + base.get_output_embeddings(), getattr(getattr(base, "model", None), "norm", None) + ) + draft = draft.cuda() + + hidden = torch.randn( + 1, SEQ_LEN, draft.config.hidden_size, device="cuda", dtype=torch.bfloat16 + ) + logits = draft._teacher_logits(hidden, {}) + assert logits.device.type == "cuda" + assert logits.shape == (1, SEQ_LEN, draft.config.vocab_size) + + def test_base_head_stays_out_of_the_draft_parameters(self): + """Gradient isolation: the teacher must not appear in the draft's optimiser.""" + draft = _external_draft() + base = get_tiny_llama() + draft.attach_base_lm_head(base.get_output_embeddings()) + draft = draft.cuda() + + param_ids = {id(p) for p in draft.parameters()} + assert all(id(p) not in param_ids for p in base.get_output_embeddings().parameters()) + assert not any("_base_lm_head" in k for k in draft.state_dict()) + + +def test_dense_path_rejects_a_missing_teacher_head(): + """Without the head the dense objective cannot be formed; fail loudly, not silently.""" + draft = _external_draft().cuda() + with pytest.raises(RuntimeError, match="Base lm_head is not attached"): + draft._teacher_logits( + torch.randn(1, SEQ_LEN, draft.config.hidden_size, device="cuda", dtype=torch.bfloat16), + {}, + ) diff --git a/tests/unit/torch/speculative/test_external_draft.py b/tests/unit/torch/speculative/test_external_draft.py new file mode 100644 index 00000000000..72951f13738 --- /dev/null +++ b/tests/unit/torch/speculative/test_external_draft.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Unit tests for the ``external`` speculative decoding mode.""" + +import pytest +import torch + +import modelopt.torch.opt as mto +import modelopt.torch.speculative as mtsp + +transformers = pytest.importorskip("transformers") + +VOCAB = 64 +HIDDEN = 32 + + +def _tiny_causal_lm(vocab_size=VOCAB): + 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) + + +def _convert(model, **overrides): + config = {"external_offline": True, **overrides} + return mtsp.convert(model, [("external", config)]) + + +def test_convert_and_restore_roundtrip(tmp_path): + """convert -> save -> restore preserves the draft weights.""" + model = _convert(_tiny_causal_lm()) + before = {k: v.clone() for k, v in model.state_dict().items()} + + mto.save(model, tmp_path / "draft.pth") + restored = mto.restore(_tiny_causal_lm(), tmp_path / "draft.pth") + + after = restored.state_dict() + assert set(before) == set(after) + for k, v in before.items(): + assert torch.equal(v, after[k]), f"weight changed across save/restore: {k}" + + +def test_vocab_mismatch_raises_and_names_both_sizes(): + """A draft indexing a different vocabulary must fail loudly, not train.""" + model = _convert(_tiny_causal_lm(vocab_size=VOCAB)) + with pytest.raises(ValueError, match="share a vocabulary") as exc: + model.validate_against_base(base_vocab_size=999) + msg = str(exc.value) + assert str(VOCAB) in msg and "999" in msg + model.validate_against_base(base_vocab_size=VOCAB) + + +@pytest.mark.parametrize("bad", ["not_a_loss", "not_an_objective"]) +def test_invalid_objective_is_rejected(bad): + with pytest.raises(ValueError, match="external_loss"): + _convert(_tiny_causal_lm(), external_loss=bad) + + +def test_online_mode_is_rejected_at_convert_not_in_the_schema(): + """The flag is part of modelopt_state, so the schema must still accept it. + + Rejecting it on the config would make a checkpoint saved with the old default + unrestorable; the check belongs in modify(), which runs on convert and restore. + """ + from modelopt.torch.speculative.config import ExternalDraftConfig + + # the schema accepts it, so an existing checkpoint still deserialises + assert ExternalDraftConfig(external_offline=False).external_offline is False + with pytest.raises(ValueError, match="external_offline"): + _convert(_tiny_causal_lm(), external_offline=False) + + +def test_draft_does_not_share_base_parameters(): + """The point of the mode: the draft owns its embeddings and lm_head.""" + draft = _convert(_tiny_causal_lm()) + base = _tiny_causal_lm() + assert ( + draft.get_input_embeddings().weight.data_ptr() + != base.get_input_embeddings().weight.data_ptr() + ) + assert draft.lm_head.weight.data_ptr() != base.lm_head.weight.data_ptr() + + +def test_soft_ce_matches_eagle_formula_and_reports_accuracy(): + """``soft_ce`` must reproduce the Eagle offline objective exactly.""" + model = _convert(_tiny_causal_lm(), external_loss="soft_ce", external_report_acc=True) + torch.manual_seed(0) + draft_logits = torch.randn(2, 5, VOCAB) + base_logits = torch.randn(2, 5, VOCAB) + loss_mask = torch.ones(2, 5) + + loss, _ = model.compute_loss(draft_logits, base_logits, loss_mask) + expected = -torch.sum( + loss_mask[:, :, None] + * torch.softmax(base_logits.float(), dim=-1) + * torch.log_softmax(draft_logits.float(), dim=-1) + ) / (loss_mask.sum() + 1e-5) + assert torch.allclose(loss, expected, atol=1e-6) + + _, acc = model.compute_loss(base_logits.clone(), base_logits, loss_mask) + assert torch.allclose(acc, torch.tensor(1.0)) + + +def test_loss_mask_excludes_masked_positions(): + """Masked positions must not contribute; prompt tokens are masked in practice.""" + model = _convert(_tiny_causal_lm(), external_loss="soft_ce", external_report_acc=False) + torch.manual_seed(0) + draft_logits = torch.randn(1, 4, VOCAB) + base_logits = torch.randn(1, 4, VOCAB) + + mask = torch.tensor([[1.0, 1.0, 0.0, 0.0]]) + loss_masked, _ = model.compute_loss(draft_logits, base_logits, mask) + loss_subset, _ = model.compute_loss(draft_logits[:, :2], base_logits[:, :2], torch.ones(1, 2)) + assert torch.allclose(loss_masked, loss_subset, atol=1e-6) + + +@pytest.mark.parametrize("identical", [False, True]) +def test_tvd_deploy_is_filtered_tvd(identical): + """TVD charges the draft for mass off the teacher's support; both sides are filtered, + so identical logits score exactly zero.""" + model = _convert( + _tiny_causal_lm(), + external_loss="tvd_deploy", + external_top_k=8, + external_top_p=0.95, + external_report_acc=False, + ) + torch.manual_seed(0) + base_logits = torch.randn(1, 3, VOCAB) + draft_logits = base_logits.clone() if identical else torch.randn(1, 3, VOCAB) + + loss, _ = model.compute_loss(draft_logits, base_logits, torch.ones(1, 3)) + if identical: + assert loss.abs() < 1e-5, f"identical logits must score 0, got {loss}" + else: + p = model._deployment_probs(base_logits, 8, 0.95) + q = model._deployment_probs(draft_logits, 8, 0.95) + assert torch.allclose(loss, (0.5 * (p - q).abs().sum(-1)).mean(), atol=1e-5) + + +class _FakeTok: + def __init__(self, vocab): + self._v = vocab + + def get_vocab(self): + return self._v + + +def test_same_size_different_tokenizers_are_rejected(): + """Equal vocab_size is not enough: misaligned ids train against the wrong columns.""" + model = _convert(_tiny_causal_lm(vocab_size=VOCAB)) + base = _FakeTok({"a": 0, "b": 1, "c": 2}) + model.validate_tokenizer_against_base(_FakeTok({"a": 0, "b": 1, "c": 2}), base) + with pytest.raises(ValueError, match="different tokenizers"): + model.validate_tokenizer_against_base(_FakeTok({"a": 1, "b": 0, "c": 2}), base) diff --git a/tests/unit/torch/speculative/test_external_draft_offline.py b/tests/unit/torch/speculative/test_external_draft_offline.py new file mode 100644 index 00000000000..dcc66c0e4bf --- /dev/null +++ b/tests/unit/torch/speculative/test_external_draft_offline.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""End-to-end offline training path for the ``external`` speculative decoding mode.""" + +import pytest +import torch + +import modelopt.torch.speculative as mtsp + +transformers = pytest.importorskip("transformers") + +VOCAB = 64 +DRAFT_HIDDEN = 32 +BASE_HIDDEN = 48 +BATCH = 2 +SEQ = 6 + + +def _causal_lm(hidden, vocab_size=VOCAB): + 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) + + +def _external_draft(**overrides): + """A converted draft with the base lm_head attached, as main.py wires it.""" + base = _causal_lm(BASE_HIDDEN) + draft = _causal_lm(DRAFT_HIDDEN) + mtsp.convert(draft, [("external", {"external_offline": True, **overrides})]) + draft.validate_against_base(base.config.vocab_size) + draft.attach_base_lm_head(base.get_output_embeddings(), base.model.norm) + return draft, base + + +def _offline_batch(): + """A batch shaped like SpeculativeDecodingOfflineCollator output.""" + torch.manual_seed(0) + return { + "input_ids": torch.randint(0, VOCAB, (BATCH, SEQ)), + "attention_mask": torch.ones(BATCH, SEQ, dtype=torch.long), + "loss_mask": torch.tensor([[0, 0, 1, 1, 1, 1], [0, 1, 1, 1, 1, 0]], dtype=torch.float), + "base_model_outputs": { + "base_model_hidden_states": torch.randn(BATCH, SEQ, BASE_HIDDEN), + "aux_hidden_states": torch.randn(BATCH, SEQ, BASE_HIDDEN * 3), + }, + } + + +@pytest.mark.parametrize("loss_fn", ["soft_ce", "tvd_deploy"]) +def test_training_steps_reduce_loss(loss_fn): + """Highest-level check: the real offline path optimises its objective. + + Subsumes the forward-shape and finite-loss checks -- a broken forward cannot + produce eight decreasing steps. + """ + draft, _ = _external_draft(external_loss=loss_fn) + batch = _offline_batch() + opt = torch.optim.AdamW(draft.parameters(), lr=1e-3) + + first = draft(**batch).loss.item() + assert draft(**batch).logits.shape == (BATCH, SEQ, VOCAB) + for _ in range(8): + opt.zero_grad() + draft(**batch).loss.backward() + opt.step() + last = draft(**batch).loss.item() + assert last < first, f"{loss_fn} did not decrease: {first:.4f} -> {last:.4f}" + + +def test_gradient_reaches_draft_but_not_base(): + """The draft trains; the frozen base head must never accumulate gradient.""" + draft, base = _external_draft() + draft(**_offline_batch()).loss.backward() + + grads = [p.grad for p in draft.parameters() if p.requires_grad] + assert any(g is not None and g.abs().sum() > 0 for g in grads), "no gradient reached the draft" + assert base.get_output_embeddings().weight.grad is None, "gradient leaked into the base lm_head" + + +def test_forward_ignores_aux_hidden_states(): + """External drafts never read aux_hidden_states; its absence must not break them.""" + draft, _ = _external_draft() + batch = _offline_batch() + del batch["base_model_outputs"]["aux_hidden_states"] + assert torch.isfinite(draft(**batch).loss) + + +def test_missing_base_lm_head_raises(): + """Training without the base head must fail loudly, not silently skip the target.""" + draft = _causal_lm(DRAFT_HIDDEN) + mtsp.convert(draft, [("external", {"external_offline": True})]) + with pytest.raises(RuntimeError, match="attach_base_lm_head"): + draft(**_offline_batch()) + + +def test_online_batch_rejected(): + draft, _ = _external_draft() + batch = _offline_batch() + del batch["base_model_outputs"] + with pytest.raises(ValueError, match="base_model_outputs"): + draft(**batch) + + +def test_base_lm_head_absent_from_draft_state_dict(): + """The exported draft must not carry base weights.""" + draft, _ = _external_draft() + keys = draft.state_dict().keys() + assert not any("_base_lm_head" in k or "_base_final_norm" in k for k in keys) + + +def test_dense_tvd_uses_full_vocabulary_and_the_filter_changes_it(): + """The dense path reconstructs the base's whole distribution from hidden states, + so ``tvd`` is a TVD over the full vocabulary; ``tvd_deploy`` filters both sides.""" + batch = _offline_batch() + plain, _ = _external_draft(external_loss="tvd") + deploy, _ = _external_draft(external_loss="tvd_deploy") + + hidden = batch["base_model_outputs"]["base_model_hidden_states"] + with torch.no_grad(): + base_logits = plain._teacher_logits(hidden, batch["base_model_outputs"]) + assert base_logits.shape[-1] == VOCAB, "teacher logits must span the full vocabulary" + + a = plain(**batch).loss + b = deploy(**batch).loss + assert 0.0 <= float(a) <= 1.0, f"TVD out of range: {a}" + assert not torch.allclose(a, b), "the serving filter must change the loss" + + +def test_prenorm_without_a_norm_raises_rather_than_corrupting(): + """A producer declaring a pre-norm hidden with no norm attached must fail loudly. + + Silently skipping the norm feeds an un-normed residual to lm_head, which yields a + plausible loss curve against a garbage teacher. + """ + base = _causal_lm(BASE_HIDDEN) + draft = _causal_lm(DRAFT_HIDDEN) + mtsp.convert(draft, [("external", {"external_offline": True, "external_loss": "tvd"})]) + # No norm located for this base, which is what the buggy lookup produced. + draft.attach_base_lm_head(base.get_output_embeddings(), None) + hidden = torch.randn(1, 4, BASE_HIDDEN) + with pytest.raises(RuntimeError, match="base_hidden_prenorm"): + draft._teacher_logits(hidden, {"base_hidden_prenorm": True})