keras-nightly 3.12.0.dev2025092403__py3-none-any.whl → 3.14.0.dev2026010104__py3-none-any.whl

keras/src/quantizers/gptq_core.py

@@ -6,10 +6,13 @@ from absl import logging
 
 from keras.src import ops
 from keras.src import utils as keras_utils
+from keras.src.dtype_policies.dtype_policy import GPTQDTypePolicy
+from keras.src.dtype_policies.dtype_policy_map import DTypePolicyMap
 from keras.src.layers import Dense
 from keras.src.layers import EinsumDense
-from keras.src.layers import Embedding
 from keras.src.quantizers.gptq import GPTQ
+from keras.src.quantizers.gptq_config import GPTQConfig
+from keras.src.quantizers.utils import should_quantize_layer
 
 
 @contextmanager
@@ -190,38 +193,6 @@ def get_dataloader(
     return samples.astype(np.int32)[:, None, :]
 
 
-def _get_backbone_layers(model):
-    """Extract embedding and transformer layers from a KerasHub model."""
-    backbone = model.backbone
-    if not hasattr(backbone, "transformer_layers"):
-        raise ValueError(
-            "The model's backbone does not have a 'transformer_layers' "
-            "attribute. Please ensure you are using a standard KerasHub "
-            "transformer model."
-        )
-    transformer_blocks = backbone.transformer_layers
-
-    embedding_layer = None
-    if hasattr(backbone, "token_embedding"):
-        embedding_layer = backbone.token_embedding
-    elif hasattr(backbone, "embedding"):
-        embedding_layer = backbone.embedding
-    return embedding_layer, transformer_blocks
-
-
-def _get_custom_layers(model):
-    """Heuristic for extracting embedding + transformer blocks from a custom
-    model."""
-    embedding_layer = None
-    transformer_blocks = []
-    for layer in model.layers:
-        if isinstance(layer, Embedding) and embedding_layer is None:
-            embedding_layer = layer
-        elif getattr(layer, "_layers", None):  # container-like block
-            transformer_blocks.append(layer)
-    return embedding_layer, transformer_blocks
-
-
 def find_layers_in_block(block):
     """
     Finds all Dense and EinsumDense layers in a transformer block.
@@ -239,39 +210,31 @@ def find_layers_in_block(block):
     return found_layers
 
 
-def apply_gptq_layerwise(model, dataloader, config):
+def apply_gptq_layerwise(dataloader, config, structure, filters=None):
     """Applies GPTQ quantization layer-by-layer to a Keras model.
 
-    This function is designed to work with common transformer architectures,
-    like those provided by KerasHub. It automatically discovers the model's
-    structure by first looking for the standard format: a `model.backbone`
-    attribute that contains a `transformer_layers` list.
-
-    If a standard backbone is not found, it falls back to a heuristic for
-    custom models, where it assumes the first `keras.layers.Embedding` layer
-    is the input embedding and any subsequent container layers are the
-    transformer blocks to be quantized.
+    This function uses the provided `structure` to identify pre-quantization
+    layers and sequential blocks.
 
     The core logic operates as follows:
-    1.  It automatically detects the model's structure, identifying the main
-        embedding layer and a sequence of transformer blocks.
-    2.  It processes the model sequentially, one block at a time. For each
+
+    1.  It processes the model sequentially, one block at a time. For each
         block, it uses temporary hooks to capture the input activations of
         each target layer during a forward pass with the calibration data.
-    3.  These captured activations are used to compute the Hessian matrix for
+    2.  These captured activations are used to compute the Hessian matrix for
         each layer's weights.
-    4.  The GPTQ algorithm is then applied to each layer to find the optimal
+    3.  The GPTQ algorithm is then applied to each layer to find the optimal
         quantized weights that minimize the error introduced.
-    5.  The output activations from the current block are then used as the
+    4.  The output activations from the current block are then used as the
         input for the next block, ensuring that quantization errors are
         accounted for throughout the model.
 
     Args:
-        model: The Keras model instance to be quantized. The function will
-            attempt to automatically discover its structure.
-        dataloader: An iterable providing calibration data. Each item should
-            be a batch of token IDs suitable for the model's embedding layer.
+        dataloader: An iterable providing calibration data.
         config: A GPTQConfiguration object.
+        structure: A dictionary with keys "pre_block_layers" and
+            "sequential_blocks".
+        filters: Optional filters to exclude layers from quantization.
 
     Raises:
         ValueError: If the function cannot automatically find an embedding
@@ -281,30 +244,23 @@ def apply_gptq_layerwise(model, dataloader, config):
     num_samples = config.num_samples
 
     logging.info("Starting model quantization...")
-    embedding_layer = None
-    transformer_blocks = []
-    if hasattr(model, "backbone"):
-        logging.info("Detected KerasHub model structure.")
-        embedding_layer, transformer_blocks = _get_backbone_layers(model)
-    else:
-        logging.info("Detected custom model structure.")
-        embedding_layer, transformer_blocks = _get_custom_layers(model)
 
-    if embedding_layer is None:
-        raise ValueError(
-            "Could not automatically find an embedding layer in the model."
-        )
+    pre_layers = structure.get("pre_block_layers", [])
+    transformer_blocks = structure.get("sequential_blocks", [])
+
     if not transformer_blocks:
         raise ValueError(
-            "Could not automatically find any transformer-like blocks to "
-            "quantize."
+            "No sequential blocks found in the provided structure to quantize."
         )
 
-    # Initial inputs are the outputs of the token embedding layer
-    inputs = [
-        embedding_layer(ops.convert_to_tensor(batch, dtype="int32"))
-        for batch in dataloader
-    ]
+    # Initial inputs are the outputs of the pre-block layers
+    inputs = []
+    for batch in dataloader:
+        batch = ops.convert_to_tensor(batch, dtype="int32")
+        for layer in pre_layers:
+            batch = layer(batch)
+        inputs.append(batch)
+
     num_samples = min(num_samples, len(inputs))
 
     progbar = keras_utils.Progbar(target=len(transformer_blocks))
@@ -313,10 +269,19 @@ def apply_gptq_layerwise(model, dataloader, config):
         logging.info(f"Quantizing Block {block_idx}")
         sub_layers_map = find_layers_in_block(block)
 
+        # Filter out layers that are not quantized with GPTQ
+        final_sub_layers_map = {}
+        for name, layer in sub_layers_map.items():
+            if not should_quantize_layer(layer, filters):
+                continue
+
+            final_sub_layers_map[name] = layer
+
+        sub_layers_map = final_sub_layers_map
+
         if not sub_layers_map:
             logging.info(
-                f"  No Dense or EinsumDense layers found in block {block_idx}. "
-                "Skipping."
+                f"  No quantizable layers found in block {block_idx}. Skipping."
             )
         else:
             logging.info(f"Found layers: {list(sub_layers_map.keys())}")
@@ -354,11 +319,30 @@ def apply_gptq_layerwise(model, dataloader, config):
     logging.info("Quantization process complete.")
 
 
-def gptq_quantize(model, config):
+def gptq_quantize(config, quantization_layer_structure, filters=None):
     """
-    Top-level function to quantize a Keras model using GPTQ.
+    Quantizes the model using GPTQ.
+
+    Args:
+        config: The GPTQ configuration.
+        quantization_layer_structure: A dictionary describing the model's layer
+        structure for quantization.
+        filters: Optional filters to exclude layers from quantization.
     """
-    logging.info("Starting GPTQ quantization process...")
+    if config.dataset is None or config.tokenizer is None:
+        raise ValueError(
+            "GPTQ quantization requires a dataset and a tokenizer. "
+            "Please provide them in the `GPTQConfig`."
+        )
+
+    if quantization_layer_structure is None:
+        raise ValueError(
+            "For 'gptq' mode, a valid quantization structure must be provided "
+            "either via `config.quantization_layer_structure` or by overriding "
+            "`model.get_quantization_layer_structure(mode)`. The structure "
+            "should be a dictionary with keys 'pre_block_layers' and "
+            "'sequential_blocks'."
+        )
 
     # Load all data needed from the generator/source in a single call.
     total_samples_to_request = config.num_samples
@@ -373,4 +357,92 @@ def gptq_quantize(model, config):
     # is now a NumPy array, which can be sliced and reused.
     calibration_dataloader = dataloader[: config.num_samples]
 
-    apply_gptq_layerwise(model, calibration_dataloader, config)
+    apply_gptq_layerwise(
+        calibration_dataloader,
+        config,
+        quantization_layer_structure,
+        filters=filters,
+    )
+
+
+def get_group_size_for_layer(layer, 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(layer.dtype_policy, GPTQDTypePolicy):
+        return layer.dtype_policy.group_size
+    elif isinstance(layer.dtype_policy, DTypePolicyMap):
+        policy = layer.dtype_policy[layer.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."
+        )
+
+
+def get_weight_bits_for_layer(layer, config):
+    """Determine the number of weight bits for GPTQ quantization.
+
+    The number of weight bits 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 weight bits should be specified in the `dtype_policy`.
+
+    Args:
+        config: An optional configuration object that may contain the
+            `weight_bits` attribute.
+    Returns:
+        int. The determined number of weight bits for GPTQ quantization.
+    Raises:
+        ValueError: If the weight bits is not specified in either the
+            `config` or the `dtype_policy`.
+    """
+    if config and isinstance(config, GPTQConfig):
+        return config.weight_bits
+    elif isinstance(layer.dtype_policy, GPTQDTypePolicy):
+        return layer.dtype_policy.weight_bits
+    elif isinstance(layer.dtype_policy, DTypePolicyMap):
+        policy = layer.dtype_policy[layer.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.weight_bits
+    else:
+        raise ValueError(
+            "For GPTQ quantization, the weight_bits must be specified"
+            "either through a `dtype_policy` of type "
+            "`GPTQDTypePolicy` or the `config` argument."
+        )

keras/src/quantizers/quantization_config.py

@@ -0,0 +1,232 @@
+from keras.src.api_export import keras_export
+from keras.src.dtype_policies import QUANTIZATION_MODES
+from keras.src.saving import serialization_lib
+
+
+@keras_export("keras.quantizers.QuantizationConfig")
+class QuantizationConfig:
+    """Base class for quantization configs.
+
+    Subclasses must implement the `mode` property and the `get_config` and
+    `from_config` class methods.
+
+    Args:
+        weight_quantizer: Quantizer for weights.
+        activation_quantizer: Quantizer for activations.
+    """
+
+    def __init__(self, weight_quantizer=None, activation_quantizer=None):
+        self.weight_quantizer = weight_quantizer
+        self.activation_quantizer = activation_quantizer
+
+    @property
+    def mode(self):
+        raise NotImplementedError(
+            "Subclasses must implement this property. Do not instantiate "
+            "QuantizationConfig directly."
+        )
+
+    def get_config(self):
+        return {
+            "weight_quantizer": serialization_lib.serialize_keras_object(
+                self.weight_quantizer
+            ),
+            "activation_quantizer": serialization_lib.serialize_keras_object(
+                self.activation_quantizer
+            ),
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        weight_quantizer = serialization_lib.deserialize_keras_object(
+            config.get("weight_quantizer")
+        )
+        activation_quantizer = serialization_lib.deserialize_keras_object(
+            config.get("activation_quantizer")
+        )
+        return cls(
+            weight_quantizer=weight_quantizer,
+            activation_quantizer=activation_quantizer,
+        )
+
+    @staticmethod
+    def weight_quantizer_or_default(config, default):
+        if config is not None and config.weight_quantizer is not None:
+            return config.weight_quantizer
+        return default
+
+    @staticmethod
+    def activation_quantizer_or_default(config, default):
+        if config is not None:
+            return config.activation_quantizer
+        return default
+
+
+@keras_export("keras.quantizers.Int8QuantizationConfig")
+class Int8QuantizationConfig(QuantizationConfig):
+    """Int8 quantization config.
+
+    Args:
+        weight_quantizer: Quantizer for weights.
+        activation_quantizer: Quantizer for activations. If "default", uses
+            AbsMaxQuantizer with axis=-1.
+    """
+
+    def __init__(self, weight_quantizer=None, activation_quantizer="default"):
+        from keras.src.quantizers.quantizers import AbsMaxQuantizer
+
+        if activation_quantizer == "default":
+            activation_quantizer = AbsMaxQuantizer()
+        super().__init__(weight_quantizer, activation_quantizer)
+        if self.weight_quantizer is not None:
+            if self.weight_quantizer.output_dtype != "int8":
+                raise ValueError(
+                    "Int8QuantizationConfig requires a weight_quantizer "
+                    "with output_dtype='int8'. Received: "
+                    f"output_dtype={self.weight_quantizer.output_dtype}"
+                )
+
+    @property
+    def mode(self):
+        return "int8"
+
+
+@keras_export("keras.quantizers.Int4QuantizationConfig")
+class Int4QuantizationConfig(QuantizationConfig):
+    """Int4 quantization config.
+
+    Args:
+        weight_quantizer: Quantizer for weights.
+        activation_quantizer: Quantizer for activations. If "default", uses
+            AbsMaxQuantizer with axis=-1.
+    """
+
+    def __init__(self, weight_quantizer=None, activation_quantizer="default"):
+        from keras.src.quantizers.quantizers import AbsMaxQuantizer
+
+        if activation_quantizer == "default":
+            activation_quantizer = AbsMaxQuantizer()
+        super().__init__(weight_quantizer, activation_quantizer)
+        if self.weight_quantizer is not None:
+            if self.weight_quantizer.value_range != (-8, 7):
+                raise ValueError(
+                    "Int4QuantizationConfig requires a weight_quantizer "
+                    "with value_range=(-8, 7). Received: "
+                    f"value_range={self.weight_quantizer.value_range}"
+                )
+
+            if self.weight_quantizer.output_dtype != "int8":
+                raise ValueError(
+                    "Int4QuantizationConfig requires a weight_quantizer "
+                    "with output_dtype='int8'. Received: "
+                    f"output_dtype={self.weight_quantizer.output_dtype}"
+                )
+
+    @property
+    def mode(self):
+        return "int4"
+
+
+@keras_export("keras.quantizers.Float8QuantizationConfig")
+class Float8QuantizationConfig(QuantizationConfig):
+    """FP8 quantization config.
+
+    FP8 mixed-precision training does not support user defined quantizers.
+    This config is only used to indicate that FP8 mixed-precision training
+    should be used.
+    """
+
+    def __init__(self):
+        super().__init__(None, None)
+
+    @property
+    def mode(self):
+        return "float8"
+
+    def get_config(self):
+        return {}
+
+    @classmethod
+    def from_config(cls, config):
+        return cls()
+
+
+def validate_and_resolve_config(mode, config):
+    """Validate and resolve quantization config.
+
+    This function validates the quantization config and resolves the mode.
+    If mode is not provided, it is inferred from the config.
+    If config is not provided, a default config is inferred from the mode.
+
+    Args:
+        mode: Quantization mode.
+        config: Quantization config.
+    """
+    # 1. Backwards Compatibility: Handle string shortcuts.
+    if isinstance(config, str):
+        mode = config
+        config = None
+
+    _validate_mode(mode)
+
+    # 2. Resolve "mode" into a Config object.
+    if config is None:
+        if mode == "int8":
+            config = Int8QuantizationConfig()
+        elif mode == "int4":
+            config = Int4QuantizationConfig()
+        elif mode == "float8":
+            config = Float8QuantizationConfig()
+        elif mode == "gptq":
+            raise ValueError(
+                "For GPTQ, you must pass a `GPTQConfig` object in the "
+                "`config` argument."
+            )
+        else:
+            if mode is not None:
+                raise ValueError(
+                    f"Invalid quantization mode. Received: mode={mode}"
+                )
+            raise ValueError(
+                "You must provide either `mode` or `config` to `quantize`."
+            )
+    else:
+        if not isinstance(config, QuantizationConfig):
+            raise ValueError(
+                "Argument `config` must be an instance of "
+                "`QuantizationConfig`. "
+                f"Received: config={config} (of type {type(config)})"
+            )
+
+    # 3. Validation: Prevent contradictions.
+    if mode is not None and config.mode != mode:
+        raise ValueError(
+            f"Contradictory arguments: mode='{mode}' but "
+            f"config.mode='{config.mode}'"
+        )
+
+    # Ensure mode is consistent.
+    mode = config.mode
+
+    # Ensure the mode derived from the config is valid.
+    _validate_mode(mode)
+
+    if mode == "gptq":
+        from keras.src.quantizers.gptq_config import GPTQConfig
+
+        if not isinstance(config, GPTQConfig):
+            raise ValueError(
+                "Mode 'gptq' requires a valid `config` argument of type "
+                f"`GPTQConfig`. Received: {type(config)}"
+            )
+
+    return config
+
+
+def _validate_mode(mode):
+    """Validates quantization mode."""
+    if mode is not None and mode not in QUANTIZATION_MODES:
+        raise ValueError(
+            "Invalid quantization mode. "
+            f"Expected one of {QUANTIZATION_MODES}. Received: mode={mode}"
+        )