diff --git a/optimum/gptq/quantizer.py b/optimum/gptq/quantizer.py index 2bb5222fe3..7bb3bab3e4 100644 --- a/optimum/gptq/quantizer.py +++ b/optimum/gptq/quantizer.py @@ -46,14 +46,28 @@ ) from accelerate.hooks import remove_hook_from_module +# The bridge is optional for GPTQModel releases that predate this API. +_gptqmodel_load_post_init = None +_gptqmodel_load_prepare_model = None + if is_gptqmodel_available(): from gptqmodel import BACKEND, QuantizeConfig from gptqmodel.quantization import FORMAT, GPTQ, METHOD from gptqmodel.utils.importer import hf_select_quant_linear_v2 - from gptqmodel.utils.model import hf_convert_gptq_v1_to_v2_format, hf_convert_gptq_v2_to_v1_format + from gptqmodel.utils.model import ( + hf_convert_gptq_v1_to_v2_format, + hf_convert_gptq_v2_to_v1_format, + ) from gptqmodel.utils.model import hf_gptqmodel_post_init as gptq_post_init from gptqmodel.version import __version__ as gptqmodel_version + try: + from gptqmodel.utils.model import hf_gptqmodel_post_init_for_load as _gptqmodel_load_post_init + from gptqmodel.utils.model import hf_gptqmodel_prepare_model_for_load as _gptqmodel_load_prepare_model + except ImportError: + _gptqmodel_load_post_init = None + _gptqmodel_load_prepare_model = None + logger = getLogger(__name__) @@ -254,6 +268,19 @@ def convert_model(self, model: nn.Module, **kwargs): Model to be converted """ + # Preserve native GPTQModel module manifests and per-module settings. + if _gptqmodel_load_prepare_model is not None: + context = _gptqmodel_load_prepare_model( + model, + checkpoint_files=kwargs.get("checkpoint_files"), + device_map=kwargs.get("device_map"), + backend=self.backend, + dtype=kwargs.get("dtype"), + ) + if context is not None: + model._gptqmodel_load_context = context + return model + if self.block_name_to_quantize is None: self.block_name_to_quantize = get_block_name_with_pattern(model) block_name = self.block_name_to_quantize @@ -613,6 +640,14 @@ def post_init_model(self, model): The input model """ + # Kernel post-init must run after checkpoint tensors reach their devices. + context = getattr(model, "_gptqmodel_load_context", None) + if context is not None: + try: + return _gptqmodel_load_post_init(model, context=context) + finally: + del model._gptqmodel_load_context + class StoreAttr(object): pass diff --git a/tests/gptq/test_quantization.py b/tests/gptq/test_quantization.py index cbf6aad669..1415d68b9b 100644 --- a/tests/gptq/test_quantization.py +++ b/tests/gptq/test_quantization.py @@ -13,12 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gc +import re import tempfile import unittest +from types import SimpleNamespace +from unittest.mock import call, patch import torch from parameterized import parameterized -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GPTQConfig from transformers.testing_utils import slow from optimum.gptq import GPTQQuantizer, load_quantized_model @@ -32,6 +36,7 @@ if is_gptqmodel_available(): from gptqmodel import GPTQModel + from gptqmodel.nn_modules.qlinear import BaseQuantLinear from gptqmodel.quantization import FORMAT, METHOD from gptqmodel.utils.importer import hf_select_quant_linear_v2 @@ -186,7 +191,7 @@ class GPTQTestActOrder(GPTQTest): # `act_group_aware` == `True` requires `desc_act` == `False` when both are explicitly set desc_act = True act_group_aware = False - expected_quantized_perplexity = 34 + expected_quantized_perplexity = 33 def test_serialization(self): """ @@ -258,6 +263,187 @@ def __init__(self): self.assertEqual(model.layer.qzero_format(), 2) +@require_gptqmodel +class GPTQNativeLoadBridgeTest(unittest.TestCase): + @patch("optimum.gptq.quantizer._gptqmodel_load_prepare_model") + def test_load_context_is_scoped_per_model(self, prepare_model): + quantizer = GPTQQuantizer(bits=4) + original_quantize_config = quantizer.quantizeConfig + first_model = torch.nn.Module() + second_model = torch.nn.Module() + first_context = SimpleNamespace(name="first") + second_context = SimpleNamespace(name="second") + prepare_model.side_effect = [first_context, second_context] + + first_result = quantizer.convert_model( + first_model, + checkpoint_files=["model.safetensors"], + device_map={"": "cpu"}, + dtype=torch.float16, + ) + second_result = quantizer.convert_model( + second_model, + checkpoint_files=["model.safetensors"], + device_map={"": "cpu"}, + dtype=torch.float16, + ) + + self.assertIs(first_result, first_model) + self.assertIs(second_result, second_model) + self.assertIs(first_model._gptqmodel_load_context, first_context) + self.assertIs(second_model._gptqmodel_load_context, second_context) + self.assertFalse(hasattr(quantizer, "_gptqmodel_load_context")) + self.assertIs(quantizer.quantizeConfig, original_quantize_config) + self.assertFalse(hasattr(quantizer, "quant_linear")) + prepare_model.assert_has_calls([ + call( + first_model, + checkpoint_files=["model.safetensors"], + device_map={"": "cpu"}, + backend=quantizer.backend, + dtype=torch.float16, + ), + call( + second_model, + checkpoint_files=["model.safetensors"], + device_map={"": "cpu"}, + backend=quantizer.backend, + dtype=torch.float16, + ), + ]) + + with patch( + "optimum.gptq.quantizer._gptqmodel_load_post_init", + side_effect=lambda model, context: model, + ) as post_init: + self.assertIs(quantizer.post_init_model(first_model), first_model) + self.assertIs(quantizer.post_init_model(second_model), second_model) + + self.assertEqual( + post_init.call_args_list, + [ + call(first_model, context=first_context), + call(second_model, context=second_context), + ], + ) + self.assertFalse(hasattr(first_model, "_gptqmodel_load_context")) + self.assertFalse(hasattr(second_model, "_gptqmodel_load_context")) + + +@slow +@require_gptqmodel +class GPTQNativeLoadBridgeIntegrationTest(unittest.TestCase): + model_id = "ModelCloud/Phi-tiny-MoE-instruct-GPTQ-W4-MixedGroup-G32-G128" + num_hidden_layers = 32 + num_local_experts = 16 + global_group_size = 128 + + @classmethod + def expected_quantized_modules(cls): + # Q/K/V and expert gate/up use G32; output and expert down projections keep the global G128. + expected = {} + for layer_index in range(cls.num_hidden_layers): + layer = f"model.layers.{layer_index}" + for projection in ("q_proj", "k_proj", "v_proj"): + expected[f"{layer}.self_attn.{projection}"] = (4, 32) + expected[f"{layer}.self_attn.o_proj"] = (4, cls.global_group_size) + for expert_index in range(cls.num_local_experts): + expert = f"{layer}.mlp.experts.{expert_index}" + expected[f"{expert}.gate_proj"] = (4, 32) + expected[f"{expert}.up_proj"] = (4, 32) + expected[f"{expert}.down_proj"] = (4, cls.global_group_size) + return expected + + @classmethod + def load_model(cls, model_id): + # Use the generic CPU kernel so this test isolates loading rather than optimized-kernel shape limits. + return AutoModelForCausalLM.from_pretrained( + model_id, + device_map={"": "cpu"}, + dtype=torch.float16, + quantization_config=GPTQConfig(bits=4, backend="torch"), + ) + + @staticmethod + def checkpoint_tensor_names(model_id): + from huggingface_hub import hf_hub_download + from safetensors import safe_open + + # Resolve the public Hub fixture through the standard cache; no machine-local fixture path is required. + checkpoint_file = hf_hub_download(repo_id=model_id, filename="model.safetensors") + with safe_open(checkpoint_file, framework="pt", device="cpu") as checkpoint: + return set(checkpoint.keys()) + + def test_all_dynamic_moe_projections_load_correctly_only_with_delegate(self): + from optimum.gptq import quantizer as optimum_quantizer + + if optimum_quantizer._gptqmodel_load_prepare_model is None: + self.skipTest("requires the GPTQModel native load delegate") + + expected_quantized_modules = self.expected_quantized_modules() + expected_dynamic = { + f"+:^{re.escape(name)}$": {"bits": bits, "group_size": group_size} + for name, (bits, group_size) in expected_quantized_modules.items() + if group_size != self.global_group_size + } + self.assertEqual(len(expected_quantized_modules), 1664) + self.assertEqual(len(expected_dynamic), 1120) + expected_dense_linear_modules = { + *(f"model.layers.{layer_index}.mlp.router" for layer_index in range(self.num_hidden_layers)), + "lm_head", + } + + # Check tensors on disk so incomplete weights cannot pass by exposing only correct config metadata. + checkpoint_tensor_names = self.checkpoint_tensor_names(self.model_id) + for component in ("qweight", "qzeros", "scales", "g_idx"): + suffix = f".{component}" + actual_names = { + name.removesuffix(suffix) for name in checkpoint_tensor_names if name.endswith(suffix) + } + self.assertEqual(actual_names, set(expected_quantized_modules)) + + # The delegate must rebuild every per-expert module and preserve its mixed group size. + model = self.load_model(self.model_id) + modules = dict(model.named_modules()) + actual_quantized_modules = { + name: (module.bits, module.group_size) + for name, module in modules.items() + if isinstance(module, BaseQuantLinear) + } + self.assertEqual(actual_quantized_modules, expected_quantized_modules) + + actual_dense_linear_modules = { + name for name, module in modules.items() if isinstance(module, torch.nn.Linear) + } + self.assertEqual(actual_dense_linear_modules, expected_dense_linear_modules) + self.assertEqual(model.config.quantization_config.group_size, self.global_group_size) + self.assertEqual(model.config.quantization_config.dynamic, expected_dynamic) + self.assertEqual( + [ + name + for name, tensor in (*model.named_parameters(), *model.named_buffers()) + if tensor.is_meta + ], + [], + ) + + with torch.inference_mode(): + output = model(input_ids=torch.tensor([[1, 42, 314, 2718, 7, 11]], dtype=torch.long)) + self.assertEqual(output.logits.shape, (1, 6, 32064)) + self.assertTrue(torch.isfinite(output.logits).all()) + + del model + gc.collect() + # The legacy path returns a model but silently loses G32 overrides and per-expert modules. + with patch("optimum.gptq.quantizer._gptqmodel_load_prepare_model", None): + legacy_model = self.load_model(self.model_id) + + legacy_modules = dict(legacy_model.named_modules()) + legacy_q_proj = legacy_modules["model.layers.0.self_attn.q_proj"] + self.assertEqual((legacy_q_proj.bits, legacy_q_proj.group_size), (4, self.global_group_size)) + self.assertNotIn("model.layers.0.mlp.experts.0.down_proj", legacy_modules) + + class GPTQUtilsTest(unittest.TestCase): """ Test utilities