keras-nightly 3.12.0.dev2025083103__py3-none-any.whl → 3.14.0.dev2026011604__py3-none-any.whl

keras/src/quantizers/gptq_config.py

@@ -1,11 +1,9 @@
-from absl import logging
-
 from keras.src.api_export import keras_export
-from keras.src.quantizers.gptq_core import quantize_model
+from keras.src.quantizers.quantization_config import QuantizationConfig
 
 
 @keras_export("keras.quantizers.GPTQConfig")
-class GPTQConfig:
+class GPTQConfig(QuantizationConfig):
     """Configuration class for the GPTQ (Gradient-based Post-Training
     Quantization) algorithm.
 
@@ -134,36 +132,88 @@ class GPTQConfig:
         activation_order: (bool, optional) If `True`, reorders weight columns
             based on activation magnitude, which can improve quantization
             accuracy. Defaults to `False`.
+        quantization_layer_structure: (dict, optional) A dictionary defining the
+            model's quantization structure. It should contain:
+            - "pre_block_layers": list of layers to run before the first block.
+            - "sequential_blocks": list of blocks to be quantized sequentially.
+            If not provided, the model must implement
+            `get_quantization_layer_structure`.
     """
 
     def __init__(
         self,
         dataset,
         tokenizer,
+        *,
         weight_bits: int = 4,
         num_samples: int = 128,
+        per_channel: bool = True,
         sequence_length: int = 512,
         hessian_damping: float = 0.01,
         group_size: int = 128,
         symmetric: bool = False,
         activation_order: bool = False,
+        quantization_layer_structure: dict = None,
     ):
+        super().__init__()
+        if weight_bits not in [2, 3, 4, 8]:
+            raise ValueError(
+                f"Unsupported weight_bits {weight_bits}. "
+                "Supported values are 2, 3, 4, and 8."
+            )
+        if num_samples <= 0:
+            raise ValueError("num_samples must be a positive integer.")
+        if sequence_length <= 0:
+            raise ValueError("sequence_length must be a positive integer.")
+        if hessian_damping < 0 or hessian_damping > 1:
+            raise ValueError("hessian_damping must be between 0 and 1.")
+        if group_size < -1 or group_size == 0:
+            raise ValueError(
+                "Invalid group_size. Supported values are -1 (whole-tensor) "
+                "or a positive integer, "
+                f"but got {group_size}."
+            )
         self.dataset = dataset
         self.tokenizer = tokenizer
         self.num_samples = num_samples
+        self.per_channel = per_channel
         self.sequence_length = sequence_length
         self.hessian_damping = hessian_damping
         self.weight_bits = weight_bits
         self.group_size = group_size
         self.symmetric = symmetric
         self.activation_order = activation_order
-
-    def quantize(self, model):
-        """
-        Applies GPTQ quantization to the provided model using this
-        configuration.
+        self.quantization_layer_structure = quantization_layer_structure
+
+    def get_config(self):
+        return {
+            # Dataset and Tokenizer are only required for a one-time
+            # calibration and are not saved in the config.
+            "dataset": None,
+            "tokenizer": None,
+            "weight_bits": self.weight_bits,
+            "num_samples": self.num_samples,
+            "per_channel": self.per_channel,
+            "sequence_length": self.sequence_length,
+            "hessian_damping": self.hessian_damping,
+            "group_size": self.group_size,
+            "symmetric": self.symmetric,
+            "activation_order": self.activation_order,
+            "quantization_layer_structure": self.quantization_layer_structure,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+    @property
+    def mode(self):
+        return "gptq"
+
+    def dtype_policy_string(self):
+        """Returns the dtype policy string for this configuration.
+
+        Returns:
+            A string representing the dtype policy, e.g. "gptq_4bit".
         """
-        logging.info("Initiating quantization from GPTQConfig...")
-        # The core logic is now delegated to gptqutils, which will handle
-        # the dynamic imports and data loading.
-        quantize_model(model=model, config=self)
+        return f"gptq/{self.weight_bits}/{self.group_size}"