Adds Native GPTQ Layer, Runtime Execution, and Serialization Support (#21641)

* initial commit

* documentation

* dtype alignment

* cleanup model.quantize() + increase top1 tolerance for tests

* test fix

* cleanup layer discovery inside transformer blocks

* delete original kernel slightly earlier

* serialization support

* Fixed tests

* Clean up dense and einsum_dense code

* format

* fix serialization issues

* Added tests for vectorized groupwise (de)quantization

* convert g_idx to float32 since int32 cannot reside on GPU in TF

* addresses reviews

* encodes gptq config in GPTQDTypePolicy name

* gptqdtypepolicy fix

* set correct _name parameter for GPTQDTypePolicy

* Address reviews

* renames self.gptq to self._is_gptq_calibrated

* fix typo

Author: JyotinderSingh

keras/api/_tf_keras/keras/dtype_policies/__init__.py

@@ -11,6 +11,9 @@
 from keras.src.dtype_policies.dtype_policy import (
     FloatDTypePolicy as FloatDTypePolicy,
 )
+from keras.src.dtype_policies.dtype_policy import (
+    GPTQDTypePolicy as GPTQDTypePolicy,
+)
 from keras.src.dtype_policies.dtype_policy import (
     QuantizedDTypePolicy as QuantizedDTypePolicy,
 )

keras/api/dtype_policies/__init__.py

@@ -11,6 +11,9 @@
 from keras.src.dtype_policies.dtype_policy import (
     FloatDTypePolicy as FloatDTypePolicy,
 )
+from keras.src.dtype_policies.dtype_policy import (
+    GPTQDTypePolicy as GPTQDTypePolicy,
+)
 from keras.src.dtype_policies.dtype_policy import (
     QuantizedDTypePolicy as QuantizedDTypePolicy,
 )

keras/src/dtype_policies/__init__.py

@@ -4,6 +4,7 @@
 from keras.src.dtype_policies.dtype_policy import QUANTIZATION_MODES
 from keras.src.dtype_policies.dtype_policy import DTypePolicy
 from keras.src.dtype_policies.dtype_policy import FloatDTypePolicy
+from keras.src.dtype_policies.dtype_policy import GPTQDTypePolicy
 from keras.src.dtype_policies.dtype_policy import QuantizedDTypePolicy
 from keras.src.dtype_policies.dtype_policy import QuantizedFloat8DTypePolicy
 from keras.src.dtype_policies.dtype_policy_map import DTypePolicyMap
@@ -14,6 +15,7 @@
     QuantizedDTypePolicy,
     QuantizedFloat8DTypePolicy,
     DTypePolicyMap,
+    GPTQDTypePolicy,
 }
 ALL_OBJECTS_DICT = {cls.__name__: cls for cls in ALL_OBJECTS}
 

keras/src/dtype_policies/dtype_policy.py

@@ -288,6 +288,94 @@ def get_config(self):
         return config
 
 
+@keras_export("keras.dtype_policies.GPTQDTypePolicy")
+class GPTQDTypePolicy(QuantizedDTypePolicy):
+    """Quantized dtype policy for GPTQ quantization.
+
+    This policy helps propagate quantization settings for GPTQ
+    when loading a GPTQ quantized model in Keras format.
+
+    Args:
+        mode: The quantization mode. This should be a string in the format
+            `"gptq/<weight_bits>/<group_size>"`.
+            -   `"gptq"`: The identifier for the quantization algorithm.
+            -   `<weight_bits>`: Number of bits to quantize weights to.
+                Supported values are 2, 3, 4, and 8.
+            -   `<group_size>`: The group size for quantization. Supported
+                values are -1 (for whole-tensor quantization) or any
+                positive integer. Typically a smaller group size leads
+                to better accuracy but slower speed.
+            Example: `"gptq/4/128"`.
+        source_name: The source dtype policy name, e.g. "float32".
+    """
+
+    def __init__(
+        self,
+        mode,
+        source_name=None,
+    ):
+        parts = mode.split("/")
+        expected_format = "'gptq/<weight_bits>/<group_size>'"
+
+        # Validate format
+        if len(parts) != 3 or parts[0] != "gptq":
+            raise ValueError(
+                "Invalid mode for GPTQDTypePolicy. Expected format "
+                f"{expected_format}, but got '{mode}'."
+            )
+
+        # Validate and cast weight_bits and group_size
+        try:
+            weight_bits = int(parts[1])
+            group_size = int(parts[2])
+        except ValueError:
+            raise ValueError(
+                "Invalid mode for GPTQDTypePolicy. <weight_bits> and "
+                "<group_size> must be integers. Expected format "
+                f"{expected_format}, but got '{mode}'."
+            )
+
+        # Validate supported values
+        if weight_bits not in [2, 3, 4, 8]:
+            raise ValueError(
+                "Invalid weight_bits in mode. Supported values are "
+                f"2, 3, 4, and 8, but got {weight_bits} from '{mode}'."
+            )
+
+        if group_size < -1 or group_size == 0:
+            raise ValueError(
+                "Invalid group_size in mode. Supported values are "
+                "-1 (whole-tensor) or a positive integer, "
+                f"but got {group_size} from '{mode}'."
+            )
+
+        base_mode = parts[0]
+        super().__init__(
+            mode=base_mode,
+            source_name=source_name,
+        )
+
+        self._name = f"{mode}_from_{source_name}"
+        self.mode = base_mode
+        self.weight_bits = weight_bits
+        self.group_size = group_size
+
+    def __eq__(self, other):
+        if super().__eq__(other) is False:
+            return False
+        return (
+            self.weight_bits == other.weight_bits
+            and self.group_size == other.group_size
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        # Reconstruct the full mode string for serialization
+        mode = f"{self.mode}/{self.weight_bits}/{self.group_size}"
+        config.update({"mode": mode})
+        return config
+
+
 @keras_export(
     [
         "keras.config.set_dtype_policy",
@@ -352,6 +440,8 @@ def _get_quantized_dtype_policy_by_str(policy):
     mode, source_name = split_name
     if policy.startswith("int8") or policy.startswith("int4"):
         return QuantizedDTypePolicy(mode, source_name)
+    elif policy.startswith("gptq"):
+        return GPTQDTypePolicy(mode, source_name)
     elif policy.startswith("float8"):
         return QuantizedFloat8DTypePolicy(mode, source_name)
     else:

keras/src/dtype_policies/dtype_policy_test.py

@@ -9,6 +9,7 @@
 from keras.src.dtype_policies.dtype_policy import QuantizedFloat8DTypePolicy
 from keras.src.dtype_policies.dtype_policy import dtype_policy
 from keras.src.dtype_policies.dtype_policy import set_dtype_policy
+from keras.src.quantizers.gptq_config import GPTQConfig
 from keras.src.testing import test_case
 
 
@@ -691,3 +692,55 @@ def test_set_policy_none(self):
         """Test setting the policy to None."""
         with self.assertRaisesRegex(ValueError, "Invalid `policy` argument"):
             set_dtype_policy(None)
+
+
+class GPTQConfigErrorHandlingTest(test_case.TestCase):
+    """Test error handling in GPTQConfig."""
+
+    def test_invalid_weight_bits(self):
+        with self.assertRaisesRegex(ValueError, "Unsupported weight_bits"):
+            GPTQConfig(
+                dataset=None,
+                tokenizer=None,
+                weight_bits=5,
+            )
+
+    def test_negative_num_samples(self):
+        with self.assertRaisesRegex(
+            ValueError, "num_samples must be a positive integer."
+        ):
+            GPTQConfig(
+                dataset=None,
+                tokenizer=None,
+                num_samples=-10,
+            )
+
+    def test_zero_sequence_length(self):
+        with self.assertRaisesRegex(
+            ValueError, "sequence_length must be a positive integer."
+        ):
+            GPTQConfig(
+                dataset=None,
+                tokenizer=None,
+                sequence_length=0,
+            )
+
+    def test_invalid_hessian_damping(self):
+        with self.assertRaisesRegex(
+            ValueError, "hessian_damping must be between 0 and 1."
+        ):
+            GPTQConfig(
+                dataset=None,
+                tokenizer=None,
+                hessian_damping=1.5,
+            )
+
+    def test_invalid_group_size(self):
+        with self.assertRaisesRegex(
+            ValueError, "Invalid group_size. Supported values are -1"
+        ):
+            GPTQConfig(
+                dataset=None,
+                tokenizer=None,
+                group_size=0,
+            )