Adds 4-bit packing support to GPTQ Quantization (#21686)

* updates int4 packing logic to handle both int8 and uint8

* adds gptq int4 packing support to einsum dense

* adds gptq int4 packing support to einsum dense

* Adds tests

* cleanup

* fix bugz

* kernel property should return unpacked kernel

* fixes dtype assertion for torch

* addresses reviews

Author: JyotinderSingh

keras/src/layers/core/dense.py

@@ -9,11 +9,8 @@
 from keras.src import quantizers
 from keras.src import regularizers
 from keras.src.api_export import keras_export
-from keras.src.dtype_policies.dtype_policy import GPTQDTypePolicy
-from keras.src.dtype_policies.dtype_policy_map import DTypePolicyMap
 from keras.src.layers.input_spec import InputSpec
 from keras.src.layers.layer import Layer
-from keras.src.quantizers.gptq_config import GPTQConfig
 from keras.src.quantizers.quantizers import dequantize_with_sz_map
 
 
@@ -143,22 +140,47 @@ def build(self, input_shape):
 
     @property
     def kernel(self):
+        from keras.src.quantizers import gptq_core
+
         if not self.built:
             raise AttributeError(
                 "You must build the layer before accessing `kernel`."
             )
-        if (
-            getattr(self, "is_gptq_calibrated", False)
-            and self.quantization_mode == "gptq"
-        ):
-            return self.quantized_kernel
-        kernel = self._kernel
-        if self.quantization_mode == "int4":
-            kernel = quantizers.unpack_int4(kernel, self._orig_input_dim)
+
+        mode = self.quantization_mode
+        is_gptq = mode == "gptq"
+        is_int4 = mode == "int4"
+        calibrated = bool(getattr(self, "is_gptq_calibrated", False))
+        gptq_bits = (
+            gptq_core.get_weight_bits_for_layer(self, None) if is_gptq else None
+        )
+
+        # Decide the source tensor first (packed vs already-quantized vs plain
+        # kernel)
+        if is_gptq and calibrated and gptq_bits != 4:
+            # calibrated GPTQ, not 4-bit, no unpacking needed
+            kernel = self.quantized_kernel
+        else:
+            # Start with the stored kernel
+            kernel = getattr(self, "_kernel", None)
+
+            # Handle int4 unpacking cases in one place
+            if is_int4:
+                kernel = quantizers.unpack_int4(kernel, self._orig_input_dim)
+            elif is_gptq and calibrated and gptq_bits == 4:
+                kernel = quantizers.unpack_int4(
+                    self.quantized_kernel,
+                    orig_len=self.units,
+                    axis=0,
+                    dtype="uint8",
+                )
+
+        # Apply LoRA once at the end.
         if self.lora_enabled:
-            return kernel + (self.lora_alpha / self.lora_rank) * ops.matmul(
+            kernel = kernel + (self.lora_alpha / self.lora_rank) * ops.matmul(
                 self.lora_kernel_a, self.lora_kernel_b
             )
+
         return kernel
 
     def call(self, inputs, training=None):
@@ -414,23 +436,33 @@ def _int8_build(self, kernel_shape):
         )
 
     def _gptq_build(self, kernel_shape, config):
+        from keras.src.quantizers import gptq_core
+
         # Ensures the forward pass uses the original high-precision kernel
         # until calibration has been performed.
         self.is_gptq_calibrated = False
         self.kernel_shape = kernel_shape
+
+        weight_bits = gptq_core.get_weight_bits_for_layer(self, config)
+        # For 4-bit weights, we pack two values per byte.
+        units = (
+            (kernel_shape[1] + 1) // 2 if weight_bits == 4 else kernel_shape[1]
+        )
+
         self.quantized_kernel = self.add_weight(
             name="kernel",
-            shape=(kernel_shape[1], kernel_shape[0]),
+            shape=(units, kernel_shape[0]),
             initializer="zeros",
             dtype="uint8",
             trainable=False,
         )
 
-        group_size = self._get_gptq_group_size(config)
-        if group_size == -1:
-            n_groups = 1
-        else:
-            n_groups = math.ceil(self.kernel_shape[0] / group_size)
+        group_size = gptq_core.get_group_size_for_layer(self, config)
+        n_groups = (
+            1
+            if group_size == -1
+            else math.ceil(self.kernel_shape[0] / group_size)
+        )
         self.kernel_scale = self.add_weight(
             name="kernel_scale",
             shape=(self.units, n_groups),
@@ -453,18 +485,31 @@ def _gptq_build(self, kernel_shape, config):
         )
 
     def _gptq_call(self, inputs, training=False):
+        from keras.src.quantizers import gptq_core
+
         if not self.is_gptq_calibrated:
             W = self._kernel
         else:
+            should_unpack = (
+                gptq_core.get_weight_bits_for_layer(self, config=None) == 4
+            )
             W = (
-                ops.transpose(
-                    dequantize_with_sz_map(
-                        self.quantized_kernel,
-                        self.kernel_scale,
-                        self.kernel_zero,
-                        self.g_idx,
-                    )
-                ),
+                quantizers.unpack_int4(
+                    self.quantized_kernel,
+                    orig_len=self.units,
+                    axis=0,
+                    dtype="uint8",
+                )
+                if should_unpack
+                else self.quantized_kernel
+            )
+            W = ops.transpose(
+                dequantize_with_sz_map(
+                    W,
+                    self.kernel_scale,
+                    self.kernel_zero,
+                    self.g_idx,
+                )
             )
 
         y = ops.matmul(inputs, W)
@@ -875,43 +920,3 @@ def _get_kernel_with_merged_lora(self):
         else:
             kernel_value = requantized_kernel
         return kernel_value, kernel_scale
-
-    def _get_gptq_group_size(self, config):
-        """Determine the group size for GPTQ quantization.
-
-        The group size can be specified either through the `config` argument
-        or through the `dtype_policy` if it is of type `GPTQDTypePolicy`.
-
-        The config argument is usually available when quantizing the layer
-        via the `quantize` method. If the layer was deserialized from a
-        saved model, the group size should be specified in the `dtype_policy`.
-
-        Args:
-            config: An optional configuration object that may contain the
-                `group_size` attribute.
-        Returns:
-            int. The determined group size for GPTQ quantization.
-        Raises:
-            ValueError: If the group size is not specified in either the
-                `config` or the `dtype_policy`.
-        """
-        if config and isinstance(config, GPTQConfig):
-            return config.group_size
-        elif isinstance(self.dtype_policy, GPTQDTypePolicy):
-            return self.dtype_policy.group_size
-        elif isinstance(self.dtype_policy, DTypePolicyMap):
-            policy = self.dtype_policy[self.path]
-            if not isinstance(policy, GPTQDTypePolicy):
-                # This should never happen based on how we set the
-                # quantization mode, but we check just in case.
-                raise ValueError(
-                    "Expected a `dtype_policy` of type `GPTQDTypePolicy`."
-                    f"Got: {type(policy)}"
-                )
-            return policy.group_size
-        else:
-            raise ValueError(
-                "For GPTQ quantization, the group_size must be specified"
-                "either through a `dtype_policy` of type "
-                "`GPTQDTypePolicy` or the `config` argument."
-            )

keras/src/layers/core/dense_test.py

@@ -898,3 +898,37 @@ def test_legacy_load_own_variables(self):
         self.assertAllClose(layer.kernel_amax_history, float8_store["5"])
         self.assertAllClose(layer.outputs_grad_scale, float8_store["6"])
         self.assertAllClose(layer.outputs_grad_amax_history, float8_store["7"])
+
+    def test_int4_gptq_kernel_returns_unpacked_form(self):
+        """Test that the `kernel` property returns the unpacked int4 GPTQ
+        kernel."""
+        layer = layers.Dense(units=2)
+        layer.build((None, 2))
+        layer.quantize(
+            "gptq",
+            config=GPTQConfig(
+                dataset=None, tokenizer=None, weight_bits=4, group_size=8
+            ),
+        )
+        layer.is_gptq_calibrated = True  # Bypass calibration check
+        packed_kernel = layer.quantized_kernel
+        self.assertAllClose(
+            layer.kernel, quantizers.unpack_int4(packed_kernel, 2)
+        )
+
+    def test_gptq_kernel_packing(self):
+        """Validates that 4-bit GPTQ packing reduces the kernel size."""
+        layer = layers.Dense(units=16, use_bias=False)
+        layer.build((None, 8))
+
+        original_kernel_params = ops.prod(layer._kernel.shape)
+
+        layer.quantize(
+            "gptq",
+            config=GPTQConfig(
+                dataset=None, tokenizer=None, weight_bits=4, group_size=8
+            ),
+        )
+
+        quantized_kernel_params = ops.prod(layer.quantized_kernel.shape)
+        self.assertEqual(quantized_kernel_params, original_kernel_params // 2)