diff --git a/modelopt/torch/opt/dynamic.py b/modelopt/torch/opt/dynamic.py index 6c0e9dd3c5c..b7bd5618b58 100644 --- a/modelopt/torch/opt/dynamic.py +++ b/modelopt/torch/opt/dynamic.py @@ -1057,6 +1057,18 @@ def decorator(dm_class: type[_DMRegistryCls.T]) -> type[_DMRegistryCls.T]: return decorator + def register_subclass(self, nn_cls: type[nn.Module], key: str) -> None: + """Register ``nn_cls`` like its registered base class, if it has one and is not covered yet. + + Needed for subclasses that override ``forward``, which the automatic lookup skips. + """ + if nn_cls in self: + return + for base_cls in self._registry: + if issubclass(nn_cls, base_cls): + self.register({nn_cls: key})(self[base_cls]) + return + def unregister(self, nn_cls: type[nn.Module] | type[T]) -> None: """Unregister a previously registered dynamic base module and all its inherited modules. 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..9e792ea8ce9 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,50 @@ 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_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/dflash/conversion.py b/modelopt/torch/speculative/dflash/conversion.py index 2f7152ab576..51f13f98c3e 100644 --- a/modelopt/torch/speculative/dflash/conversion.py +++ b/modelopt/torch/speculative/dflash/conversion.py @@ -77,12 +77,7 @@ def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertRe "Expected 'dflash' (default), 'domino', 'dspark', 'dflash2' or 'lilicorr'." ) - original_cls = type(model) - if original_cls not in registry: - for cls in registry._registry: - if issubclass(original_cls, cls): - registry.register({original_cls: "base_model_class"})(registry[cls]) - break + registry.register_subclass(type(model), "base_model_class") dflash_model = registry.convert(model) dflash_model.modify(config) diff --git a/modelopt/torch/speculative/eagle/conversion.py b/modelopt/torch/speculative/eagle/conversion.py index b7cb59f6498..15f343d8f5f 100644 --- a/modelopt/torch/speculative/eagle/conversion.py +++ b/modelopt/torch/speculative/eagle/conversion.py @@ -32,12 +32,7 @@ def convert_to_eagle_model(model: nn.Module, config: EagleConfig) -> ConvertRetu # 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 EagleDMRegistry: - for cls in EagleDMRegistry._registry: - if issubclass(original_cls, cls): - EagleDMRegistry.register({original_cls: "base_model_class"})(EagleDMRegistry[cls]) - break + EagleDMRegistry.register_subclass(type(model), "base_model_class") # merge custom config with default config default_arch_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..583ce313d16 --- /dev/null +++ b/modelopt/torch/speculative/external/conversion.py @@ -0,0 +1,65 @@ +# 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 + + ExternalDraftDMRegistry.register_subclass(type(model), "draft_model_class") + + 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..29a8355771d --- /dev/null +++ b/modelopt/torch/speculative/external/external_model.py @@ -0,0 +1,43 @@ +# 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_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_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/medusa/conversion.py b/modelopt/torch/speculative/medusa/conversion.py index 61a5e2b8ea7..75ff7fcf98f 100644 --- a/modelopt/torch/speculative/medusa/conversion.py +++ b/modelopt/torch/speculative/medusa/conversion.py @@ -31,12 +31,7 @@ def convert_to_medusa_model(model: nn.Module, config: MedusaConfig) -> ConvertRe # 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 MedusaDMRegistry: - for cls in MedusaDMRegistry._registry: - if issubclass(original_cls, cls): - MedusaDMRegistry.register({original_cls: "base_model_class"})(MedusaDMRegistry[cls]) - break + MedusaDMRegistry.register_subclass(type(model), "base_model_class") medusa_model = MedusaDMRegistry.convert(model) medusa_model.modify( 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..97dbdf46c08 --- /dev/null +++ b/modelopt/torch/speculative/plugins/hf_external.py @@ -0,0 +1,318 @@ +# 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 copy +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) + + def attach_base_lm_head(self, base_lm_head, base_final_norm=None, device=None) -> None: + """Attach the base lm_head used to project cached hidden states to teacher logits. + + Copies before freezing. The modules passed in belong to the base model and its + head may be tied to its embeddings, so setting ``requires_grad=False`` on them + would silently freeze the base's output head for anything that trains it later. + Copying here rather than on first use also keeps it out of the training step, + where it would otherwise run on every rank mid-forward and briefly hold two + copies of a vocab x hidden matrix. + """ + head = copy.deepcopy(base_lm_head) + norm = copy.deepcopy(base_final_norm) if base_final_norm is not None else None + if device is not None: + head = head.to(device) + norm = norm.to(device) if norm is not None else None + for module in (head, norm): + if module is not None: + for p in module.parameters(): + p.requires_grad = False + object.__setattr__(self, "_base_lm_head", head) + object.__setattr__(self, "_base_final_norm", 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 + # These are our own copies (see attach_base_lm_head), so moving them is safe and + # reaches nothing the caller holds. The draft's module tree does not contain them, + # so .to(device) on the draft never moves them either. Each is checked on its own + # device: they can be attached from different places, and a norm left behind would + # fail inside _maybe_apply_base_final_norm rather than here. + if self._base_lm_head.weight.device != h.device: + object.__setattr__(self, "_base_lm_head", self._base_lm_head.to(h.device)) + norm = getattr(self, "_base_final_norm", None) + if norm is not None: + # A norm built with elementwise_affine=False has no parameters at all, and one + # whose scale is a buffer has none either; probing with a bare next() would + # raise StopIteration inside the forward and surface as an unrelated + # "generator raised StopIteration". + probe = next(norm.parameters(), None) + if probe is None: + probe = next(norm.buffers(), None) + if probe is not None and probe.device != h.device: + object.__setattr__(self, "_base_final_norm", 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, + **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, + is_main_process=True, + state_dict=None, + push_to_hub=False, + *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. + """ + # transformers uploads config.json from inside super().save_pretrained, before the + # correction below can run, so the published copy would carry the dynamic subclass + # name this override exists to remove. Refuse rather than publish a checkpoint + # nothing downstream can load. + if push_to_hub: + raise ValueError( + "push_to_hub=True is not supported for an external draft: config.json is " + "uploaded before this mode can restore the original architecture name, so " + "the pushed checkpoint would be unloadable. Save locally first, then upload " + "the corrected directory." + ) + # Every parameter up to push_to_hub is named rather than swallowed by *args: + # transformers' own callers pass these positionally, and a kwargs lookup would + # then see the default -- is_main_process would be True on every rank, and the + # push_to_hub guard above would never fire. The order mirrors + # PreTrainedModel.save_pretrained for the pinned range, which takes no + # save_function argument. + result = super().save_pretrained( + save_directory, is_main_process, state_dict, push_to_hub, *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 is_main_process: + 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): + """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. + + Compares the whole vocabulary rather than a low-id sample. Byte-level BPE + tokenizers place the 256 byte symbols at ids 0-255 in the same order, so two + genuinely different tokenizers agree on exactly the region a sample would check. + """ + base_vocab, draft_vocab = base_tokenizer.get_vocab(), draft_tokenizer.get_vocab() + # Both directions: walking only the base misses a draft that carries every base + # token at the right id and adds more of its own, whose ids still land on base + # columns that mean something else. + mismatched = { + token for token, bid in base_vocab.items() if draft_vocab.get(token) != bid + } | {token for token, did in draft_vocab.items() if base_vocab.get(token) != did} + if mismatched: + raise ValueError( + f"External draft and base tokenizers disagree on {len(mismatched)} tokens " + f"(base has {len(base_vocab)}, draft has {len(draft_vocab)}). 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. + + When top-k is active the nucleus runs on the k kept values rather than on a + full sort of the vocabulary: at most k entries are non-zero after the first + stage, so sorting all of V allocates several [B, S, V] temporaries -- values, + int64 indices, cumsum and mask -- to order a tail that is entirely zero. + """ + p = torch.softmax(logits.float(), dim=-1) + vocab = p.shape[-1] + if top_k and top_k < vocab: + vals, idx = p.topk(top_k, dim=-1) + # 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. + vals = vals / vals.sum(-1, keepdim=True).clamp_min(1e-9) + if top_p and top_p < 1.0: + vals = vals * ((vals.cumsum(-1) - vals) < top_p) + vals = vals / vals.sum(-1, keepdim=True).clamp_min(1e-9) + return torch.zeros_like(p).scatter_(-1, idx, vals) + 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] + if base_logits.shape[1] < seq_len: + raise ValueError( + f"dumped base logits cover {base_logits.shape[1]} positions but the draft was " + f"given {seq_len}; the dump and the training input disagree in length." + ) + 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..69c79f05892 --- /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_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/opt/test_dynamic.py b/tests/unit/torch/opt/test_dynamic.py index 6575c18c0ee..77820cfacdb 100644 --- a/tests/unit/torch/opt/test_dynamic.py +++ b/tests/unit/torch/opt/test_dynamic.py @@ -152,3 +152,25 @@ def cb2(m, w): manager.pop_attr("bmax") manager.pop_level() assert not manager, "manager should be empty!" + + +def test_register_subclass_inherits_the_registered_base(): + """Subclasses that override forward are skipped by the automatic lookup.""" + registry = _DMRegistryCls("test_subclass") + + @registry.register({nn.Linear: "linear"}) + class _DynLinear(DynamicModule): + def _setup(self): + pass + + class _Overriding(nn.Linear): + def forward(self, x): + return super().forward(x) + + assert _Overriding not in registry + registry.register_subclass(_Overriding, "linear") + assert _Overriding in registry + # idempotent, and a class with no registered base is left alone + registry.register_subclass(_Overriding, "linear") + registry.register_subclass(nn.LayerNorm, "linear") + assert nn.LayerNorm not in registry 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..c11d87e2ff3 --- /dev/null +++ b/tests/unit/torch/speculative/test_external_draft.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. + +"""Unit tests for the ``external`` speculative decoding mode.""" + +import json + +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): + return mtsp.convert(model, [("external", overrides)]) + + +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_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: + # Independent reference: plain per-row top-k then nucleus, no call into the + # code under test. Comparing against _deployment_probs would only assert that + # the function equals itself. + def _reference(logits, k, top_p): + out = torch.zeros_like(logits, dtype=torch.float) + flat, ref = logits.float().reshape(-1, logits.shape[-1]), out.reshape(-1, out.shape[-1]) + for row in range(flat.shape[0]): + probs = torch.softmax(flat[row], dim=-1) + vals, idx = torch.topk(probs, k) + vals = vals / vals.sum() + running = 0.0 + for j in range(k): + if running >= top_p: + vals[j] = 0.0 + running += float(vals[j]) + vals = vals / vals.sum() + ref[row].scatter_(0, idx, vals) + return out + + p = _reference(base_logits, 8, 0.95) + q = _reference(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_saved_config_names_the_original_class_not_the_dynamic_subclass(tmp_path): + """vLLM and TRT-LLM resolve config.architectures against their own registries. + + transformers writes type(self).__name__, which after conversion is the registry's + dynamic subclass, and nothing downstream knows that name. + """ + model = _convert(_tiny_causal_lm()) + assert "ExternalDraft" in type(model).__name__, "precondition: conversion renamed the class" + out = tmp_path / "ckpt" + model.save_pretrained(out) + written = json.loads((out / "config.json").read_text())["architectures"] + assert written == ["LlamaForCausalLM"], written + + +def test_non_main_process_does_not_rewrite_the_config(tmp_path): + """Every rank enters save_pretrained under FSDP; a read-modify-write from all races.""" + model = _convert(_tiny_causal_lm()) + out = tmp_path / "ckpt" + out.mkdir() + (out / "config.json").write_text(json.dumps({"architectures": ["Sentinel"]})) + # positionally, as transformers' own callers pass it + model.save_pretrained(out, False) + assert json.loads((out / "config.json").read_text())["architectures"] == ["Sentinel"] + + +def test_attaching_the_base_head_does_not_touch_the_callers_modules(): + """The head belongs to the base model and may be tied to its embeddings. + + Freezing or moving it in place would silently disable training of the base's output + head for anything that uses it later. + """ + base = _tiny_causal_lm() + head, norm = base.get_output_embeddings(), base.model.norm + draft = _convert(_tiny_causal_lm()) + draft.attach_base_lm_head(head, norm) + + assert all(p.requires_grad for p in head.parameters()), "caller's head was frozen" + assert all(p.requires_grad for p in norm.parameters()), "caller's norm was frozen" + assert draft._base_lm_head is not head, "attached head must be our own copy" + assert not any(p.requires_grad for p in draft._base_lm_head.parameters()) + + +def test_push_to_hub_is_refused_rather_than_publishing_a_broken_config(tmp_path): + """config.json is uploaded from inside super().save_pretrained, before we can fix it. + + Both call shapes: transformers declares push_to_hub as the fifth positional + parameter, so a kwargs-only guard would let the positional form through. + """ + model = _convert(_tiny_causal_lm()) + with pytest.raises(ValueError, match="push_to_hub"): + model.save_pretrained(tmp_path / "ckpt", push_to_hub=True) + with pytest.raises(ValueError, match="push_to_hub"): + model.save_pretrained(tmp_path / "ckpt", True, None, True) + + +def test_a_parameterless_final_norm_does_not_abort_the_forward(): + """elementwise_affine=False norms have no parameters; probing with next() raises.""" + base = _tiny_causal_lm() + norm = torch.nn.LayerNorm(base.config.hidden_size, elementwise_affine=False) + assert next(norm.parameters(), None) is None, "precondition: no parameters" + draft = _convert(_tiny_causal_lm()) + draft.attach_base_lm_head(base.get_output_embeddings(), norm) + + hidden = torch.randn(1, 4, base.config.hidden_size) + # would raise StopIteration -> "generator raised RuntimeError" inside the forward + draft._teacher_logits(hidden, {"base_model_hidden_states": hidden}) + + +def test_tokenizer_check_sees_tokens_only_the_draft_has(): + """Walking the base alone misses a draft that adds ids the base leaves unused.""" + model = _convert(_tiny_causal_lm(vocab_size=VOCAB)) + base = _FakeTok({"a": 0, "b": 1}) + superset = _FakeTok({"a": 0, "b": 1, "<|extra|>": 2}) + with pytest.raises(ValueError, match="different tokenizers"): + model.validate_tokenizer_against_base(superset, base) + + +def test_same_size_different_tokenizers_are_rejected(): + """Equal vocab_size is not enough: misaligned ids train against the wrong columns. + + Shaped like byte-level BPE: the 256 byte symbols occupy ids 0-255 in the same order + in both tokenizers, and they diverge only above that. A low-id sample would compare + exactly the region where unrelated tokenizers agree by construction and pass. + """ + model = _convert(_tiny_causal_lm(vocab_size=VOCAB)) + byte_syms = {chr(i): i for i in range(256)} + base = _FakeTok({**byte_syms, "Ġthe": 256, "Ġand": 257}) + model.validate_tokenizer_against_base(_FakeTok(dict(base.get_vocab())), base) + diverges_above_the_bytes = _FakeTok({**byte_syms, "Ġthe": 257, "Ġand": 256}) + with pytest.raises(ValueError, match="different tokenizers"): + model.validate_tokenizer_against_base(diverges_above_the_bytes, 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..e0cf35fec0f --- /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", {**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", {})]) + 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_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})