fix(speculative): EAGLE3 on heterogeneous per-layer configs, and a first-batch hidden-state sanity check

CHANGELOG.rst

@@ -96,6 +96,8 @@ Changelog
 - Fix the column-major ``DequantizeLinear`` surgery leaving a negative or default ``axis`` attribute unchanged while transposing the weight it refers to, which changed the model's output. Re-run the surgery on affected models.
 - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``.
 - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization.
+- 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)
 ^^^^^^^^^^^^^^^^^^^

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(
@@ -707,13 +756,26 @@ 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):
+                # 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:
             with self._nvtx_range("base_model_forward"):

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),

tests/_test_utils/torch/transformers_models.py

@@ -333,6 +333,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"
 

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))