diff --git a/gptqmodel/models/auto.py b/gptqmodel/models/auto.py index fe9722c67..de75c58dd 100644 --- a/gptqmodel/models/auto.py +++ b/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, diff --git a/gptqmodel/models/definitions/__init__.py b/gptqmodel/models/definitions/__init__.py index 2760a3672..a2a335793 100644 --- a/gptqmodel/models/definitions/__init__.py +++ b/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 diff --git a/gptqmodel/models/definitions/gemma4_unified.py b/gptqmodel/models/definitions/gemma4_unified.py new file mode 100644 index 000000000..75ac50278 --- /dev/null +++ b/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) diff --git a/tests/test_gemma4_unified_support.py b/tests/test_gemma4_unified_support.py new file mode 100644 index 000000000..d86f48192 --- /dev/null +++ b/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)