From 428a3b9d2ec02760b4e40fd7f64ba515926d96a2 Mon Sep 17 00:00:00 2001 From: Anish Kumar Date: Thu, 22 Oct 2026 20:35:13 -0700 Subject: [PATCH 1/2] fix(speculative): EAGLE3 on heterogeneous configs; first-batch hidden-state check Gemma 4 declares num_key_value_heads and head_dim per layer, and transformers >= 5.15 raises AmbiguousGlobalPerLayerAttributeError on a global read. FakeBaseModel.from_source read them globally and HFEagleModel.modify built the draft config with the base's own config class, so training failed before step 1. _cfg_attr() reads such attributes from the last layer for the fake base, and a heterogeneous base gets a LlamaConfig draft with the EAGLE default attention shape (a layer-type rope map is reduced to the last layer); homogeneous bases are unchanged. The offline path only checked streamed hidden states by shape, which can agree by accident (gpt-oss-20b: cache block size 4 == four captured planes). The plugin now scores the first batch with scored tokens (top-1 agreement between the base lm_head on the final plane and the next token); rank 0 prints HIDDEN_STATE_SANITY and writes the record to SPECDEC_HIDDEN_STATE_SANITY_PATH when set. Rank 0 is the initialized process-group rank (modelopt.torch.utils.distributed), not the RANK environment variable. Signed-off-by: Anish Kumar --- CHANGELOG.rst | 2 + .../torch/speculative/plugins/hf_eagle.py | 54 ++++- .../speculative/plugins/modeling_fakebase.py | 22 ++- .../_test_utils/torch/transformers_models.py | 26 +++ .../speculative/plugins/test_fakebase.py | 21 ++ .../plugins/test_heterogeneous_eagle.py | 187 ++++++++++++++++++ 6 files changed, 307 insertions(+), 5 deletions(-) create mode 100644 tests/unit/torch/speculative/plugins/test_heterogeneous_eagle.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 582697eb7d6..d41c451a9fc 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -63,6 +63,8 @@ Changelog - Fix the Puzzletron memory objective pricing the KV cache with ``hidden_size // num_attention_heads`` instead of the model's ``head_dim``, which misestimated it for models such as Mistral-Small, GPT-OSS and Gemma-3 and for every candidate with a pruned hidden size. Recompute subblock statistics produced with earlier versions. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. +- Fix EAGLE3 training on targets with per-layer (heterogeneous) configs such as Gemma 4 under transformers 5.15 and later. +- Offline EAGLE training reports a first-batch hidden-state check (``HIDDEN_STATE_SANITY``). 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/speculative/plugins/hf_eagle.py b/modelopt/torch/speculative/plugins/hf_eagle.py index 910d5907870..5ba59066a97 100644 --- a/modelopt/torch/speculative/plugins/hf_eagle.py +++ b/modelopt/torch/speculative/plugins/hf_eagle.py @@ -17,15 +17,17 @@ import contextlib import copy +import json import os from typing import Any import torch from torch.nn.attention.flex_attention import BlockMask, create_block_mask -from transformers import Cache, DynamicCache, PreTrainedModel +from transformers import Cache, DynamicCache, LlamaConfig, PreTrainedModel from transformers.models.llama.modeling_llama import LlamaDecoderLayer from transformers.utils import ModelOutput +import modelopt.torch.utils.distributed as dist from modelopt.torch.utils import print_rank_0, warn_rank_0 from ...export.plugins.hf_spec_export import EagleExporter, SpeculativeDecodingExporter @@ -63,6 +65,48 @@ def default_eagle_aux_layer_ids(num_layers: int) -> list[int]: return sorted({1, max(0, num_layers // 2 - 1), max(0, num_layers - 4)}) +def _report_hidden_state_sanity(base_logits, input_ids, loss_mask): + """Check that streamed hidden states belong to the tokens they arrive with. + + Rank 0 prints ``HIDDEN_STATE_SANITY`` and writes it to ``SPECDEC_HIDDEN_STATE_SANITY_PATH`` + when set. Returns ``None`` when the batch scores no token. + """ + with torch.no_grad(): + predicted = base_logits[:, :-1].argmax(dim=-1) + target = input_ids[:, 1:].to(predicted.device) + if loss_mask is not None: + valid = loss_mask[:, 1:].to(predicted.device).bool() + else: + valid = torch.ones_like(target, dtype=torch.bool) + tokens = int(valid.sum()) + if tokens == 0: + return None + agreement = float(((predicted == target) & valid).sum()) / tokens + record = { + "base_top1_next_token_agreement": round(agreement, 4), + "tokens": tokens, + "batch_shape": list(input_ids.shape), + } + print_rank_0("HIDDEN_STATE_SANITY " + json.dumps(record)) + path = os.environ.get("SPECDEC_HIDDEN_STATE_SANITY_PATH") + if path and dist.is_master(): + with open(path, "w") as handle: + json.dump(record, handle) + return record + + +def _reduce_per_layer_rope(arch_config, base_cfg): + """Reduce a ``rope_parameters`` keyed by layer type to the last base layer's entry.""" + rope = arch_config.get("rope_parameters") + if isinstance(rope, dict) and rope and all(isinstance(v, dict) for v in rope.values()): + selected = rope.get((getattr(base_cfg, "layer_types", None) or [None])[-1]) + if selected is None: + del arch_config["rope_parameters"] + else: + arch_config["rope_parameters"] = dict(selected) + arch_config.pop("rope_scaling", None) # the legacy key would win in from_dict + + @EagleDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) class HFEagleModel(EagleModel): """Eagle Model Class for huggingface models.""" @@ -363,6 +407,11 @@ def modify( # Use the base model's config class so fields like max_position_embeddings are declared # before transformers' rope standardization runs in __post_init__. base_config_cls = type(self._base_llm_config) + if getattr(self._base_llm_config, "is_heterogeneous", False): + # A per-layer base config (Gemma 4) cannot describe the single llama-style draft + # layer, which keeps the EAGLE default attention shape as for every other base. + base_config_cls = LlamaConfig + _reduce_per_layer_rope(arch_config, self._base_llm_config) self.eagle_config = base_config_cls.from_dict(arch_config) self.eagle_config.eagle_decoder_type = self.eagle_decoder_type self.eagle_config.draft_vocab_size = getattr( @@ -714,6 +763,9 @@ def forward( base_outputs.out_hiddens, kwargs["base_model_outputs"], self._base_model_norm ) base_outputs.logits = self._base_model_lm_head(out_hiddens) + if not getattr(self, "_hidden_state_sanity_reported", False): + record = _report_hidden_state_sanity(base_outputs.logits, input_ids, loss_mask) + self._hidden_state_sanity_reported = record is not None past_key_values = None else: with self._nvtx_range("base_model_forward"): diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index 85a3fe3a9f4..3e9bcf0a604 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py @@ -92,6 +92,20 @@ def find_final_norm_path(model) -> str | None: "language_model.backbone", ] _VLM_CONFIG_ATTRS = ["text_config", "llm_config"] + + +def _cfg_attr(cfg, name, default=None, layer=-1): + """Read ``name`` from ``cfg``; take it from one layer when transformers declares it per layer.""" + try: + return getattr(cfg, name, default) + except Exception as exc: # transformers' AmbiguousGlobalPerLayerAttributeError, by name + if type(exc).__name__ != "AmbiguousGlobalPerLayerAttributeError": + raise + per_layer = getattr(cfg, "per_layer_config") + count = len(per_layer) + return getattr(per_layer[(count + layer) % count], name, default) + + _SAFETENSORS_INDEX_FILENAME = "model.safetensors.index.json" # Single-file safetensors names to try, in order. Mistral native checkpoints # use ``consolidated.safetensors`` instead of the HF-standard ``model.safetensors``. @@ -232,10 +246,10 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM max_position_embeddings=getattr(base_cfg, "max_position_embeddings", None), dtype=getattr(base_cfg, "dtype", torch.bfloat16), tie_word_embeddings=getattr(base_cfg, "tie_word_embeddings", False), - num_attention_heads=getattr(base_cfg, "num_attention_heads", None), - num_key_value_heads=getattr(base_cfg, "num_key_value_heads", None), - intermediate_size=getattr(base_cfg, "intermediate_size", None), - rms_norm_eps=getattr(base_cfg, "rms_norm_eps", 1e-6), + num_attention_heads=_cfg_attr(base_cfg, "num_attention_heads", None), + num_key_value_heads=_cfg_attr(base_cfg, "num_key_value_heads", None), + intermediate_size=_cfg_attr(base_cfg, "intermediate_size", None), + rms_norm_eps=_cfg_attr(base_cfg, "rms_norm_eps", 1e-6), # Shared with the exporter: where a config keeps rope_theta depends on the # transformers version, and reading it wrong is silent until serve time. rope_theta=_get_rope_theta(base_cfg), diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index bdf910981d2..8076d6ccc49 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -334,6 +334,32 @@ def create_tiny_gemma3vl_dir( ) +##### Gemma4 (text-only; per-layer attention config) ##### +def get_tiny_gemma4(**config_kwargs) -> PreTrainedModel: + """Tiny Gemma 4 text model whose head_dim and kv-head count differ between layer types.""" + set_seed(SEED) + + # Lazy import: Gemma4TextConfig only exists in transformers >= 5.15. + from transformers import Gemma4TextConfig + + kwargs = { + "dtype": torch.bfloat16, + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 6, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "num_global_key_value_heads": 1, + "attention_k_eq_v": True, # without it the global kv-head count is ignored + "head_dim": 8, + "global_head_dim": 16, + "max_position_embeddings": 32, + "vocab_size": 32, + } + kwargs.update(config_kwargs) + return AutoModelForCausalLM.from_config(Gemma4TextConfig(**kwargs)) + + ##### Qwen3.5-VL (dense or MoE, hybrid GatedDeltaNet + gated attention) ##### QWEN3_5_VL_REF = "Qwen/Qwen3.5-0.8B" diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index 44cc19ed90c..282e646c995 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/tests/unit/torch/speculative/plugins/test_fakebase.py @@ -23,6 +23,7 @@ pytest.importorskip("transformers") import transformers +from _test_utils.torch.transformers_models import get_tiny_gemma4 from modelopt.torch.speculative.plugins.modeling_fakebase import ( FakeBaseConfig, @@ -65,12 +66,32 @@ def fake_checkpoint(tmp_path, fake_config): return tmp_path +@pytest.fixture +def heterogeneous_config(fake_config, monkeypatch): + """Swap the fake config for a tiny per-layer (Gemma 4) text config of the same shape.""" + pytest.importorskip("transformers.models.gemma4") + cfg = get_tiny_gemma4(hidden_size=_HIDDEN_SIZE, vocab_size=_VOCAB_SIZE).config + if not getattr(cfg, "is_heterogeneous", False): + pytest.skip("transformers without per-layer Gemma 4 attributes") + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **kw: cfg) + return cfg + + def test_fakebase_local_happy_path(fake_checkpoint): model = FakeBaseModel.from_source(str(fake_checkpoint)) assert model.lm_head.weight.shape == torch.Size([_VOCAB_SIZE, _HIDDEN_SIZE]) assert model.embed_tokens.weight.shape == torch.Size([_VOCAB_SIZE, _HIDDEN_SIZE]) +def test_fakebase_heterogeneous_config_reads_the_last_layer(fake_checkpoint, heterogeneous_config): + """Attributes transformers declares per layer are taken from the last layer, not raised on.""" + model = FakeBaseModel.from_source(str(fake_checkpoint)) + last = heterogeneous_config.per_layer_config[-1] + assert model.config.num_attention_heads == heterogeneous_config.num_attention_heads + assert model.config.num_key_value_heads == last.num_key_value_heads + assert model.lm_head.weight.shape == torch.Size([_VOCAB_SIZE, _HIDDEN_SIZE]) + + def test_fakebase_missing_index_raises(tmp_path, fake_config): with pytest.raises(FileNotFoundError, match="safetensors"): FakeBaseModel.from_source(str(tmp_path)) diff --git a/tests/unit/torch/speculative/plugins/test_heterogeneous_eagle.py b/tests/unit/torch/speculative/plugins/test_heterogeneous_eagle.py new file mode 100644 index 00000000000..1bac341e1e0 --- /dev/null +++ b/tests/unit/torch/speculative/plugins/test_heterogeneous_eagle.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""EAGLE on heterogeneous (per-layer) targets and the first-batch hidden-state check.""" + +import json +from copy import deepcopy + +import pytest +import torch + +pytest.importorskip("transformers") +from _test_utils.torch.transformers_models import get_tiny_gemma4 +from transformers import LlamaConfig + +import modelopt.torch.speculative as mtsp +import modelopt.torch.utils.distributed as dist +from modelopt.torch.speculative.config import EAGLE3_DEFAULT_CFG +from modelopt.torch.speculative.plugins.hf_eagle import ( + _reduce_per_layer_rope, + _report_hidden_state_sanity, +) +from modelopt.torch.speculative.plugins.modeling_fakebase import _cfg_attr + +_PER_LAYER_ROPE = { + "full_attention": {"rope_theta": 1e6, "rope_type": "default"}, + "sliding_attention": {"rope_theta": 1e4, "rope_type": "default"}, +} + + +def _tiny_gemma4(): + """Tiny Gemma 4 text model whose head_dim and kv heads differ between layer types.""" + pytest.importorskip("transformers.models.gemma4") + model = get_tiny_gemma4() + if not getattr(model.config, "is_heterogeneous", False): + pytest.skip("transformers without per-layer Gemma 4 attributes") + return model + + +def _convert_offline(model, **arch_overrides): + """Convert ``model`` in place to an offline EAGLE-3 draft with a tiny recipe.""" + config = deepcopy(EAGLE3_DEFAULT_CFG["config"]) + config["eagle_architecture_config"].update( + { + "draft_vocab_size": model.config.vocab_size, + "intermediate_size": 32, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "head_dim": 2, + **arch_overrides, + } + ) + config["eagle_offline"] = True + config["eagle_use_torch_compile"] = False + mtsp.convert(model, mode=[("eagle", config)]) + return model + + +def _offline_batch(model, seq_len=8): + """A streamed-base batch: final and aux planes in the shape the offline collator emits.""" + hidden = model.config.hidden_size + num_aux = len(model.eagle_config.eagle_aux_hidden_state_layer_ids) + dtype = model.lm_head.weight.dtype + return { + "input_ids": torch.randint(0, model.config.vocab_size, (2, seq_len)), + "base_model_outputs": { + "base_model_hidden_states": torch.randn(2, seq_len, hidden, dtype=dtype), + "aux_hidden_states": torch.randn(2, seq_len, hidden * num_aux, dtype=dtype), + }, + } + + +def _aligned_logits(input_ids, vocab): + """Logits whose argmax at position t is the token at t + 1.""" + logits = torch.zeros(*input_ids.shape, vocab) + logits[:, :-1].scatter_(2, input_ids[:, 1:, None], 1.0) + return logits + + +def test_cfg_attr(): + """Homogeneous configs read as getattr; per-layer attributes come from the requested layer.""" + llama = LlamaConfig(hidden_size=32, num_attention_heads=4, num_key_value_heads=2) + assert _cfg_attr(llama, "num_key_value_heads") == 2 + assert _cfg_attr(llama, "missing", "fallback") == "fallback" + cfg = _tiny_gemma4().config + with pytest.raises(Exception, match="per-layer"): + getattr(cfg, "head_dim") + first, last = cfg.per_layer_config[0], cfg.per_layer_config[-1] + assert first.num_key_value_heads != last.num_key_value_heads + assert _cfg_attr(cfg, "num_key_value_heads") == last.num_key_value_heads + assert _cfg_attr(cfg, "head_dim") == last.head_dim + assert _cfg_attr(cfg, "head_dim", layer=0) == first.head_dim + + +def test_per_layer_rope_is_reduced_to_the_last_layer(): + """A flat rope_parameters is kept; a layer-type map is reduced and survives from_dict.""" + cfg = _tiny_gemma4().config + flat = {"rope_theta": 5e5, "rope_type": "default"} + arch = {"rope_parameters": dict(flat)} + _reduce_per_layer_rope(arch, cfg) + assert arch["rope_parameters"] == flat + arch = { + "rope_scaling": {"rope_type": "default", "rope_theta": 1e4}, + "rope_parameters": dict(_PER_LAYER_ROPE), + } + _reduce_per_layer_rope(arch, cfg) + selected = _PER_LAYER_ROPE[cfg.layer_types[-1]] + assert arch["rope_parameters"] == selected + assert LlamaConfig.from_dict(arch).rope_parameters["rope_theta"] == selected["rope_theta"] + + +@pytest.mark.parametrize("per_layer_rope", [False, True]) +def test_convert_builds_a_llama_draft_on_a_heterogeneous_base(per_layer_rope): + """The draft keeps the recipe's attention shape; its RoPE is the recipe's or the reduced map.""" + overrides = {"rope_parameters": dict(_PER_LAYER_ROPE)} if per_layer_rope else {} + model = _convert_offline(_tiny_gemma4(), **overrides) + draft = model.eagle_config + assert isinstance(draft, LlamaConfig) + assert (draft.num_key_value_heads, draft.head_dim) == (16, 2) + if per_layer_rope: + expected = _PER_LAYER_ROPE[model.config.layer_types[-1]]["rope_theta"] + else: + expected = EAGLE3_DEFAULT_CFG["config"]["eagle_architecture_config"]["rope_theta"] + assert draft.rope_parameters["rope_theta"] == expected + assert len(model.eagle_module.layers) == 1 + + +def test_offline_forward_reports_the_first_scored_batch(tmp_path, monkeypatch, capsys): + """An unscored batch leaves the check pending; the first scored batch reports exactly once.""" + model = _convert_offline(_tiny_gemma4()).train() + path = tmp_path / "sanity.json" + monkeypatch.setenv("SPECDEC_HIDDEN_STATE_SANITY_PATH", str(path)) + torch.manual_seed(0) + batch = _offline_batch(model) + + model(**batch, loss_mask=torch.zeros_like(batch["input_ids"])) + assert not model._hidden_state_sanity_reported and not path.exists() + + model(**batch) + assert model._hidden_state_sanity_reported + assert json.loads(path.read_text())["tokens"] == batch["input_ids"][:, 1:].numel() + assert capsys.readouterr().out.count("HIDDEN_STATE_SANITY") == 1 + + path.unlink() + model(**batch) + assert "HIDDEN_STATE_SANITY" not in capsys.readouterr().out and not path.exists() + + +def test_hidden_state_sanity(tmp_path, monkeypatch, capsys): + """Aligned planes score 1.0, shifted near 0; no scored token stays pending; only rank 0 reports.""" + torch.manual_seed(0) + input_ids = torch.randint(0, 16, (1, 64)) + loss_mask = torch.ones_like(input_ids) + loss_mask[:, :8] = 0 + path = tmp_path / "sanity.json" + monkeypatch.setenv("SPECDEC_HIDDEN_STATE_SANITY_PATH", str(path)) + monkeypatch.setenv("RANK", "1") # a launcher variable; the process-group rank decides + logits = _aligned_logits(input_ids, 16) + + record = _report_hidden_state_sanity(logits, input_ids, loss_mask) + assert record["base_top1_next_token_agreement"] == 1.0 + assert record["tokens"] == int(loss_mask[:, 1:].sum()) + assert json.loads(path.read_text()) == record + assert "HIDDEN_STATE_SANITY " + json.dumps(record) in capsys.readouterr().out + + shifted = _report_hidden_state_sanity(torch.roll(logits, 3, dims=1), input_ids, loss_mask) + assert shifted["base_top1_next_token_agreement"] < 0.5 + assert _report_hidden_state_sanity(logits, input_ids, torch.zeros_like(input_ids)) is None + + monkeypatch.delenv("RANK") + monkeypatch.setattr(dist, "rank", lambda group=None: 1) # an initialized non-zero rank + path.unlink() + capsys.readouterr() + assert _report_hidden_state_sanity(logits, input_ids, None)["tokens"] == 63 + assert capsys.readouterr().out == "" and not path.exists() From d03b93f1113f9cdd851136736f343fc736c32c24 Mon Sep 17 00:00:00 2001 From: Anish Kumar Date: Mon, 26 Oct 2026 13:51:26 -0700 Subject: [PATCH 2/2] fix(speculative): score the hidden-state sanity check from the streamed planes When the producer also supplies base_model_logits, the first-batch check scored those logits, so a misaligned hidden-state stream could still report full agreement. Derive the scored logits from base_model_hidden_states (with the pre-norm handling) on that one batch; training keeps using the supplied logits. Adds a test with aligned supplied logits and random planes, which must score low. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Anish Kumar --- modelopt/torch/speculative/plugins/hf_eagle.py | 14 ++++++++++++-- .../plugins/test_heterogeneous_eagle.py | 15 +++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_eagle.py b/modelopt/torch/speculative/plugins/hf_eagle.py index 5ba59066a97..e34247cdc67 100644 --- a/modelopt/torch/speculative/plugins/hf_eagle.py +++ b/modelopt/torch/speculative/plugins/hf_eagle.py @@ -756,15 +756,25 @@ def forward( # Parse base model outputs forwarded from teacher assert "base_model_outputs" in kwargs base_outputs = EagleBaseModelOutput.from_offline_dict(kwargs["base_model_outputs"]) + streamed_logits = None if base_outputs.logits is None: # Re-apply the base final norm when the producer captured a pre-norm hidden # (vLLM streaming); fails loud if pre-norm is declared but no norm was located. out_hiddens = _maybe_apply_base_final_norm( base_outputs.out_hiddens, kwargs["base_model_outputs"], self._base_model_norm ) - base_outputs.logits = self._base_model_lm_head(out_hiddens) + streamed_logits = base_outputs.logits = self._base_model_lm_head(out_hiddens) if not getattr(self, "_hidden_state_sanity_reported", False): - record = _report_hidden_state_sanity(base_outputs.logits, input_ids, loss_mask) + # Score the streamed hidden states, never producer-supplied logits: those can be + # correct while the hidden-state stream is misaligned, which is what this catches. + if streamed_logits is None: + out_hiddens = _maybe_apply_base_final_norm( + base_outputs.out_hiddens, + kwargs["base_model_outputs"], + self._base_model_norm, + ) + streamed_logits = self._base_model_lm_head(out_hiddens) + record = _report_hidden_state_sanity(streamed_logits, input_ids, loss_mask) self._hidden_state_sanity_reported = record is not None past_key_values = None else: diff --git a/tests/unit/torch/speculative/plugins/test_heterogeneous_eagle.py b/tests/unit/torch/speculative/plugins/test_heterogeneous_eagle.py index 1bac341e1e0..f5d5afd1975 100644 --- a/tests/unit/torch/speculative/plugins/test_heterogeneous_eagle.py +++ b/tests/unit/torch/speculative/plugins/test_heterogeneous_eagle.py @@ -158,6 +158,21 @@ def test_offline_forward_reports_the_first_scored_batch(tmp_path, monkeypatch, c assert "HIDDEN_STATE_SANITY" not in capsys.readouterr().out and not path.exists() +def test_offline_sanity_scores_the_streamed_hidden_states(tmp_path, monkeypatch): + """Producer-supplied logits do not take part: aligned logits with random planes score low.""" + model = _convert_offline(_tiny_gemma4()).train() + path = tmp_path / "sanity.json" + monkeypatch.setenv("SPECDEC_HIDDEN_STATE_SANITY_PATH", str(path)) + torch.manual_seed(0) + batch = _offline_batch(model) + batch["base_model_outputs"]["base_model_logits"] = _aligned_logits( + batch["input_ids"], model.config.vocab_size + ).to(model.lm_head.weight.dtype) + + model(**batch) + assert json.loads(path.read_text())["base_top1_next_token_agreement"] < 0.5 + + def test_hidden_state_sanity(tmp_path, monkeypatch, capsys): """Aligned planes score 1.0, shifted near 0; no scored token stays pending; only rank 0 reports.""" torch.manual_seed(0)