fix(gptq): delegate native GPTQModel loads

optimum/gptq/quantizer.py

@@ -46,14 +46,30 @@
     )
     from accelerate.hooks import remove_hook_from_module
 
+# The bridge is optional for GPTQModel releases that predate this API.
+hf_gptqmodel_post_init_for_load = None
+hf_gptqmodel_prepare_model_for_load = 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,
+            hf_gptqmodel_prepare_model_for_load,
+        )
+    except ImportError:
+        hf_gptqmodel_post_init_for_load = None
+        hf_gptqmodel_prepare_model_for_load = None
+
 logger = getLogger(__name__)
 
 
@@ -159,6 +175,7 @@ def __init__(
         self.quant_method = QuantizationMethod.GPTQ
         self.cache_block_outputs = cache_block_outputs
         self.modules_in_block_to_quantize = modules_in_block_to_quantize
+        self._gptqmodel_load_context = None
 
         self.quantizeConfig = QuantizeConfig(
             bits=self.bits,
@@ -254,6 +271,21 @@ def convert_model(self, model: nn.Module, **kwargs):
                 Model to be converted
 
         """
+        # Preserve native GPTQModel module manifests and per-module settings.
+        if hf_gptqmodel_prepare_model_for_load is not None:
+            context = hf_gptqmodel_prepare_model_for_load(
+                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:
+                self._gptqmodel_load_context = context
+                self.quantizeConfig = context.quantize_config
+                self.quant_linear = context.quant_linear
+                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 +645,13 @@ def post_init_model(self, model):
                 The input model
         """
 
+        # Kernel post-init must run after checkpoint tensors reach their devices.
+        if self._gptqmodel_load_context is not None:
+            return hf_gptqmodel_post_init_for_load(
+                model,
+                context=self._gptqmodel_load_context,
+            )
+
         class StoreAttr(object):
             pass
 

tests/gptq/test_quantization.py

@@ -15,6 +15,8 @@
 
 import tempfile
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
 import torch
 from parameterized import parameterized
@@ -186,7 +188,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 +260,40 @@ def __init__(self):
         self.assertEqual(model.layer.qzero_format(), 2)
 
 
+@require_gptqmodel
+class GPTQNativeLoadBridgeTest(unittest.TestCase):
+    @patch("optimum.gptq.quantizer.hf_gptqmodel_prepare_model_for_load")
+    def test_convert_and_post_init_delegate_to_gptqmodel_bridge(self, prepare_model):
+        quantizer = GPTQQuantizer(bits=4)
+        model = torch.nn.Module()
+        context = SimpleNamespace(
+            quantize_config=quantizer.quantizeConfig,
+            quant_linear=object(),
+        )
+        prepare_model.return_value = context
+
+        result = quantizer.convert_model(
+            model,
+            checkpoint_files=["model.safetensors"],
+            device_map={"": "cpu"},
+            dtype=torch.float16,
+        )
+
+        self.assertIs(result, model)
+        self.assertIs(quantizer._gptqmodel_load_context, context)
+        prepare_model.assert_called_once_with(
+            model,
+            checkpoint_files=["model.safetensors"],
+            device_map={"": "cpu"},
+            backend=quantizer.backend,
+            dtype=torch.float16,
+        )
+
+        with patch("optimum.gptq.quantizer.hf_gptqmodel_post_init_for_load", return_value=model) as post_init:
+            self.assertIs(quantizer.post_init_model(model), model)
+            post_init.assert_called_once_with(model, context=context)
+
+
 class GPTQUtilsTest(unittest.TestCase):
     """
     Test utilities