[MODEL] support gemma4_unified

gptqmodel/models/auto.py

@@ -97,6 +97,7 @@
 from .definitions.gemma3 import Gemma3ForConditionalGenerationGPTQ, Gemma3QModel  # noqa: E402
 from .definitions.gemma3n import Gemma3nForConditionalGenerationGPTQ, Gemma3nTextQModel  # noqa: E402
 from .definitions.gemma4 import Gemma4ForConditionalGenerationGPTQ, Gemma4TextQModel  # noqa: E402
+from .definitions.gemma4_unified import Gemma4UnifiedForConditionalGenerationGPTQ  # noqa: E402
 from .definitions.glm import GlmQModel  # noqa: E402
 from .definitions.glm4_moe import GLM4MoEGPTQ  # noqa: E402
 from .definitions.glm4_moe_lite import Glm4MoeLiteQModel  # noqa: E402
@@ -255,6 +256,7 @@
     "gemma3n": Gemma3nForConditionalGenerationGPTQ,
     "gemma4_text": Gemma4TextQModel,
     "gemma4": Gemma4ForConditionalGenerationGPTQ,
+    "gemma4_unified": Gemma4UnifiedForConditionalGenerationGPTQ,
     "phi": PhiQModel,
     "phi3": Phi3QModel,
     "phi4mm": Phi4MMGPTQ,

gptqmodel/models/definitions/__init__.py

@@ -30,6 +30,7 @@
 from .gemma3 import Gemma3QModel
 from .gemma3n import Gemma3nForConditionalGenerationGPTQ, Gemma3nTextQModel
 from .gemma4 import Gemma4ForConditionalGenerationGPTQ, Gemma4TextQModel
+from .gemma4_unified import Gemma4UnifiedForConditionalGenerationGPTQ
 from .glm import GlmQModel
 from .glmasr import GlmASRGPTQ
 from .glm_ocr import GlmOCRGPTQ

gptqmodel/models/definitions/gemma4_unified.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
+# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
+# SPDX-License-Identifier: Apache-2.0
+# Contact: qubitium@modelcloud.ai, x.com/qubitium
+
+import torch
+
+from ...utils.device import get_device
+from ...utils.model import get_module_by_name_prefix, move_to, nested_move_to
+from ..base import BaseQModel
+
+
+def _prepare_gemma4_unified_replay_kwargs(model_def, layer, layer_input, additional_inputs, target_device):
+    """Refresh Gemma 4 unified rotary kwargs per layer during cached replay.
+
+    Gemma 4 unified builds one rope (cos, sin) per attention ``layer_type`` and hands each
+    decoder layer the single tuple for its own type, so replaying a layer in isolation needs
+    that tuple regenerated for the layer's ``layer_type`` (sliding vs full). This mirrors the
+    Gemma 4 rope replay but without any per-layer-input handling, which this variant lacks.
+    """
+
+    rotary_path = getattr(model_def, "rotary_embedding", None)
+    if not rotary_path or not layer_input:
+        return additional_inputs
+
+    rotary, _ = get_module_by_name_prefix(model_def.model, [rotary_path])
+    if rotary is None:
+        return additional_inputs
+
+    layer_type = getattr(getattr(layer, "self_attn", None), "layer_type", None)
+    if layer_type is None:
+        return additional_inputs
+
+    hidden_states = layer_input[0]
+    seq_len = hidden_states.shape[1] if hidden_states.dim() >= 2 else hidden_states.shape[0]
+    batch_dim = hidden_states.shape[0] if hidden_states.dim() >= 2 else 1
+
+    position_ids = additional_inputs.get("position_ids")
+    if position_ids is None or position_ids.shape[-1] != seq_len:
+        position_ids = torch.arange(seq_len, device=target_device, dtype=torch.long).unsqueeze(0).expand(batch_dim, -1)
+        additional_inputs["position_ids"] = position_ids
+
+    try:
+        rotary_device = get_device(rotary)
+    except Exception:
+        rotary_device = position_ids.device
+
+    rotary_position_ids = move_to(position_ids, device=rotary_device)
+    rotary_input = torch.empty(1, device=rotary_device, dtype=hidden_states.dtype)
+    additional_inputs["position_embeddings"] = nested_move_to(
+        rotary(rotary_input, rotary_position_ids, layer_type),
+        device=target_device,
+    )
+
+    return additional_inputs
+
+
+class Gemma4UnifiedForConditionalGenerationGPTQ(BaseQModel):
+    """Quantization definition for Gemma 4 unified (multimodal) checkpoints.
+
+    Gemma 4 unified reuses the composite decoder layout (per-projection q/k/v norms and the
+    dual pre/post feed-forward norms) but, unlike the per-layer-input Gemma 4 variants, has no
+    per-layer input gate/projection, so those module-tree entries and the per-layer-input
+    capture hooks are intentionally absent. The sliding/full rope boundaries still need the
+    per-layer replay refresh below.
+    """
+
+    layer_modules_strict = False
+    support_batch_quantize = False
+    pre_lm_head_norm_module = "model.language_model.norm"
+    rotary_embedding = "model.language_model.rotary_emb"
+
+    module_tree = [
+        "model",
+        "language_model",
+        "layers",
+        "#",
+        {
+            "input_layernorm": ("input_layernorm:!",),
+            "self_attn": (
+                "q_norm:!",
+                "q_proj:0",
+                "k_norm:!",
+                "k_proj:0",
+                "v_norm:!",
+                "v_proj:0",
+                "o_proj:1",
+            ),
+            "post_attention_layernorm": ("post_attention_layernorm:!",),
+            "pre_feedforward_layernorm": ("pre_feedforward_layernorm:!",),
+            "mlp": ("gate_proj:0", "up_proj:0", "down_proj:1"),
+            "post_feedforward_layernorm": ("post_feedforward_layernorm:!",),
+        },
+    ]
+
+    def prepare_layer_replay_kwargs(self, layer, layer_input, additional_inputs, target_device):
+        """Refresh Gemma 4 unified rope kwargs during cached layer replay."""
+
+        return _prepare_gemma4_unified_replay_kwargs(self, layer, layer_input, additional_inputs, target_device)

tests/test_gemma4_unified_support.py

@@ -0,0 +1,78 @@
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from gptqmodel.models import auto
+from gptqmodel.models.definitions.gemma4_unified import Gemma4UnifiedForConditionalGenerationGPTQ
+
+
+def test_gemma4_unified_model_type_selects_definition(monkeypatch):
+    fake_config = SimpleNamespace(model_type="gemma4_unified")
+
+    monkeypatch.setattr(auto, "resolve_trust_remote_code", lambda path, trust_remote_code=False: trust_remote_code)
+    monkeypatch.setattr(auto.AutoConfig, "from_pretrained", lambda *args, **kwargs: fake_config)
+
+    assert auto.check_and_get_model_definition("/tmp/gemma4-unified") is Gemma4UnifiedForConditionalGenerationGPTQ
+
+
+def test_gemma4_unified_module_tree_excludes_per_layer_input_paths():
+    layer_modules = Gemma4UnifiedForConditionalGenerationGPTQ.simple_layer_modules(
+        model_config=SimpleNamespace(),
+        quantize_config=SimpleNamespace(dynamic=None),
+    )
+    flat_modules = {name for block in layer_modules for name in block}
+
+    assert Gemma4UnifiedForConditionalGenerationGPTQ.layer_modules_strict is False
+    for proj in ("q_proj", "k_proj", "v_proj", "o_proj"):
+        assert f"self_attn.{proj}" in flat_modules
+    for proj in ("gate_proj", "up_proj", "down_proj"):
+        assert f"mlp.{proj}" in flat_modules
+    # Unlike the per-layer-input Gemma 4 variants, gemma4_unified has no per-layer adapters.
+    assert "per_layer_input_gate" not in flat_modules
+    assert "per_layer_projection" not in flat_modules
+
+
+def test_gemma4_unified_replay_kwargs_refresh_position_embeddings_per_layer_type():
+    class _FakeRotary(nn.Module):
+        def forward(self, x, position_ids, layer_type=None):
+            marker = 7.0 if layer_type == "full_attention" else 3.0
+            shape = (position_ids.shape[0], position_ids.shape[1], 1)
+            value = torch.full(shape, marker, dtype=x.dtype, device=x.device)
+            return value, value + 1
+
+    class _LanguageModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.rotary_emb = _FakeRotary()
+
+    class _Core(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.language_model = _LanguageModel()
+
+    class _Wrapper(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = _Core()
+
+    model_def = object.__new__(Gemma4UnifiedForConditionalGenerationGPTQ)
+    nn.Module.__init__(model_def)
+    model_def.model = _Wrapper()
+
+    hidden_states = torch.randn(1, 4, 8)
+    for layer_type, marker in (("sliding_attention", 3.0), ("full_attention", 7.0)):
+        layer = SimpleNamespace(self_attn=SimpleNamespace(layer_type=layer_type))
+        refreshed = model_def.prepare_layer_replay_kwargs(
+            layer=layer,
+            layer_input=[hidden_states],
+            additional_inputs={
+                "position_ids": torch.arange(4).unsqueeze(0),
+                "position_embeddings": ("stale",),
+            },
+            target_device=torch.device("cpu"),
+        )
+        cos, sin = refreshed["position_embeddings"]
+        assert cos.shape == (1, 4, 1)
+        assert torch.all(cos == marker)
+        assert torch.all(sin == marker + 1)