Adds Serialization Support for QuantizationConfig based quantized models (#21928)

* Adds serialization support for QuantizationConfig

* fix imports

* remove redundant config resolution logic

* remove redundant config resolution logic

* move deserialization to from_config + restructure tests

* from_config should call super

Author: JyotinderSingh

keras/src/layers/core/dense.py

@@ -12,8 +12,8 @@
 from keras.src.layers.input_spec import InputSpec
 from keras.src.layers.layer import Layer
 from keras.src.quantizers.quantization_config import QuantizationConfig
-from keras.src.quantizers.quantization_config import validate_and_resolve_config
 from keras.src.quantizers.quantizers import dequantize_with_sz_map
+from keras.src.saving import serialization_lib
 
 
 @keras_export("keras.layers.Dense")
@@ -94,6 +94,7 @@ def __init__(
         bias_constraint=None,
         lora_rank=None,
         lora_alpha=None,
+        quantization_config=None,
         **kwargs,
     ):
         if not isinstance(units, int) or units <= 0:
@@ -115,13 +116,18 @@ def __init__(
         self.lora_rank = lora_rank
         self.lora_alpha = lora_alpha if lora_alpha is not None else lora_rank
         self.lora_enabled = False
+        self.quantization_config = quantization_config
         self.input_spec = InputSpec(min_ndim=2)
         self.supports_masking = True
 
     def build(self, input_shape):
         kernel_shape = (input_shape[-1], self.units)
         if self.quantization_mode:
-            self.quantized_build(kernel_shape, mode=self.quantization_mode)
+            self.quantized_build(
+                kernel_shape,
+                mode=self.quantization_mode,
+                config=self.quantization_config,
+            )
         if self.quantization_mode not in ("int8", "int4", "gptq"):
             # If the layer is quantized to int8 or int4, `self._kernel` will be
             # added in `self._int8_build` or `_int4_build`. Therefore, we skip
@@ -330,12 +336,25 @@ def get_config(self):
             "bias_regularizer": regularizers.serialize(self.bias_regularizer),
             "kernel_constraint": constraints.serialize(self.kernel_constraint),
             "bias_constraint": constraints.serialize(self.bias_constraint),
+            "quantization_config": serialization_lib.serialize_keras_object(
+                self.quantization_config
+            ),
         }
         if self.lora_rank:
             config["lora_rank"] = self.lora_rank
             config["lora_alpha"] = self.lora_alpha
         return {**base_config, **config}
 
+    @classmethod
+    def from_config(cls, config):
+        config = config.copy()
+        config["quantization_config"] = (
+            serialization_lib.deserialize_keras_object(
+                config.get("quantization_config", None)
+            )
+        )
+        return super().from_config(config)
+
     @property
     def variable_serialization_spec(self):
         """Returns a dict mapping quantization modes to variable names in order.
@@ -777,27 +796,26 @@ def quantize(self, mode=None, type_check=True, config=None):
         if type_check and (type(self) is not Dense):
             raise self._not_implemented_error(self.quantize)
 
-        config = validate_and_resolve_config(mode, config)
-        mode = config.mode
+        self.quantization_config = config
 
         kernel_shape = self._kernel.shape
         if mode == "int8":
             weight_quantizer = QuantizationConfig.weight_quantizer_or_default(
-                config, quantizers.AbsMaxQuantizer(axis=0)
+                self.quantization_config, quantizers.AbsMaxQuantizer(axis=0)
             )
             kernel_value, kernel_scale = weight_quantizer(
                 self._kernel, to_numpy=True
             )
             kernel_scale = ops.squeeze(kernel_scale, axis=0)
             del self._kernel
             # Build variables for int8 mode
-            self.quantized_build(kernel_shape, mode, config)
+            self.quantized_build(kernel_shape, mode, self.quantization_config)
             self._kernel.assign(kernel_value)
             self.kernel_scale.assign(kernel_scale)
         elif mode == "int4":
             # 1. Quantize to int4 values (still int8 dtype, range [-8,7])
             weight_quantizer = QuantizationConfig.weight_quantizer_or_default(
-                config,
+                self.quantization_config,
                 quantizers.AbsMaxQuantizer(
                     axis=0, value_range=(-8, 7), output_dtype="int8"
                 ),
@@ -811,12 +829,12 @@ def quantize(self, mode=None, type_check=True, config=None):
             del self._kernel
             # Build variables using the original kernel shape; _int4_build will
             # compute the packed shape internally.
-            self.quantized_build(kernel_shape, mode, config)
+            self.quantized_build(kernel_shape, mode, self.quantization_config)
             # Assign packed values.
             self._kernel.assign(packed_kernel_value)
             self.kernel_scale.assign(kernel_scale)
         elif mode == "gptq":
-            self.quantized_build(kernel_shape, mode, config)
+            self.quantized_build(kernel_shape, mode, self.quantization_config)
         elif mode == "float8":
             self.quantized_build(kernel_shape, mode)
         else:
@@ -828,7 +846,7 @@ def quantize(self, mode=None, type_check=True, config=None):
 
             policy_name = mode
             if mode == "gptq":
-                policy_name = config.dtype_policy_string()
+                policy_name = self.quantization_config.dtype_policy_string()
             policy = dtype_policies.get(
                 f"{policy_name}_from_{self.dtype_policy.name}"
             )

keras/src/layers/core/dense_test.py

@@ -1025,3 +1025,123 @@ def test_gptq_kernel_packing(self):
 
         quantized_kernel_params = ops.prod(layer.quantized_kernel.shape)
         self.assertEqual(quantized_kernel_params, original_kernel_params // 2)
+
+    def _check_quantizer_config(
+        self, quantizer, valid_class, axis, value_range
+    ):
+        self.assertIsInstance(quantizer, valid_class)
+        self.assertEqual(quantizer.axis, axis)
+
+        # Normalize value_range to list
+        if value_range is not None:
+            self.assertAllEqual(quantizer.value_range, value_range)
+
+    def test_dense_int8_custom_quantizer(self):
+        """
+        Test custom quantizer serialization for dense layer.
+        """
+        # Setup
+        weight_range = (-127, 127)
+        act_range = (-5, 5)
+        config = Int8QuantizationConfig(
+            weight_quantizer=AbsMaxQuantizer(axis=0, value_range=weight_range),
+            activation_quantizer=AbsMaxQuantizer(
+                axis=-1, value_range=act_range
+            ),
+        )
+
+        # Build & Quantize
+        layer = layers.Dense(10)
+        layer.build((None, 5))
+        layer.quantize("int8", config=config)
+
+        # Serialize & Deserialize
+        serialized = layer.get_config()
+        new_layer = layers.Dense.from_config(serialized)
+
+        # Verify
+        self.assertIsInstance(
+            new_layer.quantization_config, Int8QuantizationConfig
+        )
+        self._check_quantizer_config(
+            new_layer.quantization_config.weight_quantizer,
+            AbsMaxQuantizer,
+            axis=(0,),
+            value_range=weight_range,
+        )
+        self._check_quantizer_config(
+            new_layer.quantization_config.activation_quantizer,
+            AbsMaxQuantizer,
+            axis=(-1,),
+            value_range=act_range,
+        )
+
+    def test_dense_int8_weight_only_quantizer(self):
+        """
+        Test custom quantizer serialization for dense layer with
+        weight-only quantization.
+        """
+        # Setup
+        config = Int8QuantizationConfig(
+            weight_quantizer=AbsMaxQuantizer(axis=0),
+            activation_quantizer=None,
+        )
+
+        # Build & Quantize
+        layer = layers.Dense(10)
+        layer.build((None, 5))
+        layer.quantize("int8", config=config)
+
+        # Serialize & Deserialize
+        serialized = layer.get_config()
+        new_layer = layers.Dense.from_config(serialized)
+
+        # Verify
+        self.assertIsInstance(
+            new_layer.quantization_config, Int8QuantizationConfig
+        )
+        self.assertIsInstance(
+            new_layer.quantization_config.weight_quantizer, AbsMaxQuantizer
+        )
+        self.assertIsNone(new_layer.quantization_config.activation_quantizer)
+
+    def test_dense_int4_custom_quantizer(self):
+        """
+        Test custom quantizer serialization for dense layer with
+        int4 quantization.
+        """
+        # Setup
+        weight_range = (-8, 7)
+        act_range = (-2, 2)
+        config = Int4QuantizationConfig(
+            weight_quantizer=AbsMaxQuantizer(axis=0, value_range=weight_range),
+            activation_quantizer=AbsMaxQuantizer(
+                axis=-1, value_range=act_range
+            ),
+        )
+
+        # Build & Quantize
+        layer = layers.Dense(10)
+        layer.build((None, 5))
+        layer.quantize("int4", config=config)
+
+        # Serialize & Deserialize
+        serialized = layer.get_config()
+        new_layer = layers.Dense.from_config(serialized)
+
+        # Verify
+        self.assertIsInstance(
+            new_layer.quantization_config, Int4QuantizationConfig
+        )
+        self._check_quantizer_config(
+            new_layer.quantization_config.weight_quantizer,
+            AbsMaxQuantizer,
+            axis=(0,),
+            value_range=weight_range,
+        )
+        self._check_quantizer_config(
+            new_layer.quantization_config.activation_quantizer,
+            AbsMaxQuantizer,
+            axis=(-1,),
+            value_range=act_range,
+        )

keras/src/layers/core/einsum_dense.py

@@ -18,6 +18,7 @@
 from keras.src.layers.layer import Layer
 from keras.src.quantizers.quantization_config import QuantizationConfig
 from keras.src.quantizers.quantizers import dequantize_with_sz_map
+from keras.src.saving import serialization_lib
 
 
 @keras_export("keras.layers.EinsumDense")
@@ -136,6 +137,7 @@ def __init__(
         lora_rank=None,
         lora_alpha=None,
         gptq_unpacked_column_size=None,
+        quantization_config=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -156,6 +158,7 @@ def __init__(
         self.lora_alpha = lora_alpha if lora_alpha is not None else lora_rank
         self.lora_enabled = False
         self.gptq_unpacked_column_size = gptq_unpacked_column_size
+        self.quantization_config = quantization_config
 
     def build(self, input_shape):
         shape_data = _analyze_einsum_string(
@@ -171,6 +174,7 @@ def build(self, input_shape):
             self.quantized_build(
                 kernel_shape,
                 mode=self.quantization_mode,
+                config=self.quantization_config,
             )
         # Skip creating a duplicate kernel variable when the layer is already
         # quantized to int8 or int4, because `quantized_build` has created the
@@ -394,6 +398,9 @@ def get_config(self):
             ),
             "kernel_constraint": constraints.serialize(self.kernel_constraint),
             "bias_constraint": constraints.serialize(self.bias_constraint),
+            "quantization_config": serialization_lib.serialize_keras_object(
+                self.quantization_config
+            ),
         }
         if self.lora_rank:
             config["lora_rank"] = self.lora_rank
@@ -402,6 +409,16 @@ def get_config(self):
             config["gptq_unpacked_column_size"] = self.gptq_unpacked_column_size
         return {**base_config, **config}
 
+    @classmethod
+    def from_config(cls, config):
+        config = config.copy()
+        config["quantization_config"] = (
+            serialization_lib.deserialize_keras_object(
+                config.get("quantization_config", None)
+            )
+        )
+        return super().from_config(config)
+
     @property
     def variable_serialization_spec(self):
         """Returns a dict mapping quantization modes to variable names in order.
@@ -465,6 +482,10 @@ def _int8_build(self, kernel_shape, config=None):
                 quantizers.AbsMaxQuantizer(axis=self._input_reduced_axes),
             )
         )
+        # If the config provided a default AbsMaxQuantizer, we need to
+        # override the axis to match the equation's reduction axes.
+        if isinstance(self.inputs_quantizer, quantizers.AbsMaxQuantizer):
+            self.inputs_quantizer.axis = tuple(self._input_reduced_axes)
         self._kernel = self.add_weight(
             name="kernel",
             shape=kernel_shape,
@@ -614,6 +635,10 @@ def _int4_build(self, kernel_shape, config=None):
                 quantizers.AbsMaxQuantizer(axis=self._input_reduced_axes),
             )
         )
+        # If the config provided a default AbsMaxQuantizer, we need to
+        # override the axis to match the equation's reduction axes.
+        if isinstance(self.inputs_quantizer, quantizers.AbsMaxQuantizer):
+            self.inputs_quantizer.axis = tuple(self._input_reduced_axes)
 
         # Choose the axis to perform int4 packing - use the first reduced axis
         # for the kernel (analogous to the input dimension of a Dense layer).
@@ -980,14 +1005,16 @@ def quantize(self, mode=None, type_check=True, config=None):
         if type_check and (type(self) is not EinsumDense):
             raise self._not_implemented_error(self.quantize)
 
+        self.quantization_config = config
+
         kernel_shape = self._kernel.shape
         if mode in ("int8", "int4", "gptq"):
             self._set_quantization_info()
 
         if mode == "int8":
             # Quantize `self._kernel` to int8 and compute corresponding scale
             weight_quantizer = QuantizationConfig.weight_quantizer_or_default(
-                config,
+                self.quantization_config,
                 quantizers.AbsMaxQuantizer(axis=self._kernel_reduced_axes),
             )
             kernel_value, kernel_scale = weight_quantizer(
@@ -998,7 +1025,7 @@ def quantize(self, mode=None, type_check=True, config=None):
         elif mode == "int4":
             # Quantize to int4 values (stored in int8 dtype, range [-8, 7])
             weight_quantizer = QuantizationConfig.weight_quantizer_or_default(
-                config,
+                self.quantization_config,
                 quantizers.AbsMaxQuantizer(
                     axis=self._kernel_reduced_axes,
                     value_range=(-8, 7),
@@ -1017,7 +1044,7 @@ def quantize(self, mode=None, type_check=True, config=None):
             )
             kernel_value = packed_kernel_value
             del self._kernel
-        self.quantized_build(kernel_shape, mode, config)
+        self.quantized_build(kernel_shape, mode, self.quantization_config)
 
         # Assign values to the newly created variables.
         if mode in ("int8", "int4"):
@@ -1028,7 +1055,7 @@ def quantize(self, mode=None, type_check=True, config=None):
         if self.dtype_policy.quantization_mode is None:
             policy_name = mode
             if mode == "gptq":
-                policy_name = config.dtype_policy_string()
+                policy_name = self.quantization_config.dtype_policy_string()
             policy = dtype_policies.get(
                 f"{policy_name}_from_{self.dtype_policy.name}"
             )