fix(gptq): delegate native GPTQModel loads

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
 

tests/gptq/test_quantization.py

@@ -15,6 +15,8 @@
 
 import tempfile
 import unittest
+from types import SimpleNamespace
+from unittest.mock import call, 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,73 @@ 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"))
+
+
 class GPTQUtilsTest(unittest.TestCase):
     """
     Test utilities