fix(models): resolve verifier_norm class by verifier family

src/speculators/models/dflash/core.py

@@ -25,6 +25,7 @@
     conditional_torch_compile,
     flatten_rope_parameters,
     resolve_target_layer_ids,
+    resolve_verifier_norm_class,
 )
 
 logger = logging.getLogger(__name__)
@@ -123,7 +124,10 @@ def __init__(
             config.transformer_layer_config.hidden_size,
             eps=config.transformer_layer_config.rms_norm_eps,  # type: ignore[arg-type]
         )
-        self.verifier_norm = Qwen3RMSNorm(
+        # Must apply the verifier's own final-norm convention (`x * (1 + w)`
+        # for the Gemma/Qwen3.5 families, `x * w` otherwise) or the
+        # reconstructed verifier targets are silently mis-scaled.
+        self.verifier_norm = resolve_verifier_norm_class(config)(
             config.transformer_layer_config.hidden_size,
             eps=config.transformer_layer_config.rms_norm_eps,  # type: ignore[arg-type]
         )

src/speculators/models/utils.py

@@ -1,13 +1,67 @@
+import json
 import logging
 import warnings
 from copy import deepcopy
 from functools import partial
+from pathlib import Path
 
 import torch
 from transformers import AutoConfig, PretrainedConfig
+from transformers.models.gemma3.modeling_gemma3 import Gemma3RMSNorm
+from transformers.models.qwen3.modeling_qwen3 import Qwen3RMSNorm
 
 logger = logging.getLogger(__name__)
 
+# Verifier families whose final norm follows the Gemma convention
+# `x_norm * (1 + w)` instead of the plain RMSNorm `x_norm * w`. Matched as a
+# prefix against the verifier's HF `model_type` / `text_config.model_type`
+# (read from its config.json when resolvable) with a fallback to lowercased
+# `architectures`. Qwen3.5 is in here because `Qwen3_5RMSNorm` is an alias of
+# vLLM's `GemmaRMSNorm` (see vllm/model_executor/models/qwen3_5.py) and
+# transformers' `Qwen3_5RMSNorm.forward` computes `output * (1.0 + weight)`.
+GEMMA_STYLE_FINAL_NORM_PREFIXES = ("gemma", "qwen3_5")
+
+
+def uses_gemma_style_final_norm(config) -> bool:  # noqa: ANN001
+    """Whether the verifier's final norm is `x_norm * (1 + w)` rather than `x_norm * w`."""
+    verifier = getattr(getattr(config, "speculators_config", None), "verifier", None)
+    if verifier is None:
+        return False
+
+    candidates: list[str] = []
+
+    # Prefer the verifier's own config.json (model_type is the reliable signal,
+    # and the multimodal wrapper keeps the family under text_config).
+    name_or_path = getattr(verifier, "name_or_path", None)
+    if name_or_path:
+        config_file = Path(name_or_path) / "config.json"
+        if config_file.is_file():
+            try:
+                hf_config = json.loads(config_file.read_text())
+            except (OSError, ValueError):
+                hf_config = {}
+            candidates.append(str(hf_config.get("model_type", "")))
+            candidates.append(str(hf_config.get("text_config", {}).get("model_type", "")))
+
+    # Fall back to the architectures recorded in the speculator config.
+    candidates.extend(str(a) for a in (getattr(verifier, "architectures", None) or []))
+
+    return any(
+        c.lower().startswith(GEMMA_STYLE_FINAL_NORM_PREFIXES) for c in candidates if c
+    )
+
+
+def resolve_verifier_norm_class(config) -> type:  # noqa: ANN001
+    """The RMSNorm class matching the verifier's final-norm weight convention.
+
+    Generalizes #892's Gemma3 selection: the frozen ``verifier_norm`` must
+    apply the same gain convention the verifier was trained under, or the
+    reconstructed targets are silently mis-scaled.
+    """
+    if uses_gemma_style_final_norm(config):
+        return Gemma3RMSNorm
+    return Qwen3RMSNorm
+
 
 def conditional_torch_compile(func=None, *args, **kwargs):
     if func is None:

tests/unit/models/test_verifier_norm_convention.py

@@ -0,0 +1,120 @@
+"""Verifier final-norm class resolution.
+
+Some verifier families store the final RMSNorm weight in the Gemma
+convention — the applied gain is ``1 + w`` — while a plain ``Qwen3RMSNorm``
+applies gain ``w``. This family includes Gemma itself and the Qwen3.5/Qwen3.8
+models, whose ``Qwen3_5RMSNorm`` is an alias of vLLM's ``GemmaRMSNorm``.
+The frozen ``verifier_norm`` must therefore be constructed from the class
+matching the verifier's convention, or the reconstructed verifier targets
+are silently mis-scaled.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import cast
+
+import pytest
+import torch
+from transformers.models.gemma3.modeling_gemma3 import Gemma3RMSNorm
+from transformers.models.qwen3.modeling_qwen3 import Qwen3RMSNorm
+
+from speculators.models.dflash.core import DFlashDraftModel
+from speculators.models.utils import resolve_verifier_norm_class, uses_gemma_style_final_norm
+
+from .test_checkpoint_key_ownership import _fake_verifier, _make_fake_loader, _make_model
+
+
+def _point_at_fake_verifier(
+    model, tmp_path, model_type: str, architectures: list[str], text_model_type: str = ""
+):
+    verifier_dir = tmp_path / model_type / text_model_type / "_".join(architectures)
+    verifier_dir.mkdir(parents=True)
+    raw = {"model_type": model_type, "architectures": architectures}
+    if text_model_type:  # multimodal wrapper carrying the family in text_config
+        raw["text_config"] = {"model_type": text_model_type}
+    (verifier_dir / "config.json").write_text(json.dumps(raw))
+    model.config.speculators_config.verifier.name_or_path = str(verifier_dir)
+    model.config.speculators_config.verifier.architectures = architectures
+
+
+@pytest.mark.parametrize(
+    ("model_type", "text_model_type", "architectures", "expected"),
+    [
+        ("qwen3_5", "", ["Qwen3_5ForConditionalGeneration"], True),
+        ("qwen3", "qwen3_5_text", ["Qwen3_5ForConditionalGeneration"], True),
+        ("gemma3_text", "", ["Gemma3ForConditionalGeneration"], True),
+        ("qwen3", "", ["Gemma3ForCausalLM"], True),  # architectures fallback
+        ("qwen3", "", ["Qwen3ForCausalLM"], False),
+        ("llama", "", ["LlamaForCausalLM"], False),
+    ],
+    ids=[
+        "qwen3_5",
+        "qwen3_5_text_nested",
+        "gemma3",
+        "gemma3_arch_fallback",
+        "qwen3_negative",
+        "llama_negative",
+    ],
+)
+def test_detection_by_model_type_and_architectures(
+    tmp_path, model_type: str, text_model_type: str, architectures: list[str], expected: bool
+):
+    model = _make_model(DFlashDraftModel, draft_vocab_size=64)
+    _point_at_fake_verifier(model, tmp_path, model_type, architectures, text_model_type)
+    assert uses_gemma_style_final_norm(model.config) is expected
+    assert (
+        resolve_verifier_norm_class(model.config) is Gemma3RMSNorm
+        if expected
+        else resolve_verifier_norm_class(model.config) is Qwen3RMSNorm
+    )
+
+
+def test_plain_construction_unchanged():
+    """A plain-convention verifier (the default dummy) keeps Qwen3RMSNorm."""
+    model = _make_model(DFlashDraftModel, draft_vocab_size=64)
+    assert isinstance(model.verifier_norm, Qwen3RMSNorm)
+
+
+@pytest.mark.parametrize("model_type", ["qwen3_5", "gemma3_text"], ids=["qwen3_5", "gemma3"])
+def test_gemma_style_construction_swaps_class(
+    tmp_path, monkeypatch: pytest.MonkeyPatch, model_type: str
+):
+    """Gemma-convention verifiers construct verifier_norm as Gemma3RMSNorm
+    and load the raw checkpoint weight unchanged (the +1 lives in forward)."""
+    model = _make_model(DFlashDraftModel, draft_vocab_size=64)
+    _point_at_fake_verifier(model, tmp_path, model_type, ["XForCausalLM"])
+    rebuilt = DFlashDraftModel(model.config)
+    assert isinstance(rebuilt.verifier_norm, Gemma3RMSNorm)
+
+    # The weight must load verbatim: the convention is applied in forward,
+    # not baked into the parameter.
+    fake = _fake_verifier()
+    fake["model.norm.weight"] = torch.randn(16)
+    rebuilt.save_pretrained(tmp_path / "draft")
+    monkeypatch.setattr(
+        "speculators.utils.loading.load_model_layers",
+        _make_fake_loader(fake),
+    )
+    loaded = cast(
+        "DFlashDraftModel",
+        DFlashDraftModel.from_pretrained(tmp_path / "draft", local_files_only=True),
+    )
+    assert isinstance(loaded.verifier_norm, Gemma3RMSNorm)
+    assert torch.equal(loaded.verifier_norm.weight, fake["model.norm.weight"])
+
+
+def test_gemma3_rmsnorm_matches_folded_qwen3_rmsnorm():
+    """Reference equivalence: Gemma3RMSNorm(w) == Qwen3RMSNorm(w + 1).
+
+    The two fix mechanisms for the convention mismatch (class swap vs
+    folding +1 into the loaded weight) must produce the same function.
+    """
+    torch.manual_seed(0)
+    w = torch.randn(16)
+    x = torch.randn(8, 16)
+    gemma = Gemma3RMSNorm(16, eps=1e-6)
+    gemma.weight.data = w.clone()
+    qwen_folded = Qwen3RMSNorm(16, eps=1e-6)
+    qwen_folded.weight.data = w + 1.0
+    torch.testing.assert_close(gemma(x), qwen_folded(x), rtol=1e-5, atol=1e-5)