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

keras/src/quantizers/awq_config.py

@@ -0,0 +1,140 @@
+from keras.src.api_export import keras_export
+from keras.src.quantizers.quantization_config import QuantizationConfig
+
+
+@keras_export("keras.quantizers.AWQConfig")
+class AWQConfig(QuantizationConfig):
+    """Configuration class for AWQ (Activation-aware Weight Quantization).
+
+    AWQ is a post-training quantization method that identifies and protects
+    salient weights based on activation magnitudes. It applies per-channel
+    scaling before quantization to minimize accuracy loss.
+
+    Methodology:
+    1. Collects activation statistics from calibration data
+    2. Identifies salient weight channels based on activation magnitudes
+    3. Searches for optimal per-channel scaling factors via grid search
+    4. Applies scaling before quantization to protect important weights
+
+    References:
+    - Original AWQ paper: "AWQ: Activation-aware Weight Quantization for
+      LLM Compression and Acceleration" (https://arxiv.org/abs/2306.00978)
+    - Reference implementation: https://github.com/mit-han-lab/llm-awq
+
+    Args:
+        dataset: The calibration dataset. It can be an iterable that yields
+            strings or pre-tokenized numerical tensors (e.g., a list of
+            strings, a generator, or a NumPy array). This data is used to
+            analyze activation patterns.
+        tokenizer: A tokenizer instance (or a similar callable) that is used
+            to process the `dataset`.
+        weight_bits: The number of bits for weight quantization. AWQ presently
+            only supports 4-bit quantization. Defaults to 4.
+        num_samples: The number of calibration data samples to use from the
+            dataset. Defaults to 128.
+        sequence_length: The sequence length to use for each calibration
+            sample. Defaults to 512.
+        group_size: The size of weight groups to quantize together. A
+            `group_size` of -1 indicates per-channel quantization.
+            Defaults to 128.
+        num_grid_points: The number of grid search points for finding optimal
+            per-channel scales. Higher values may find better scales but
+            take longer. Defaults to 20.
+        quantization_layer_structure: A dictionary defining the model's
+            quantization structure. It should contain:
+            - "pre_block_layers": list of layers to run before the first
+              block (e.g., embedding layer).
+            - "sequential_blocks": list of transformer blocks to quantize
+              sequentially.
+            If not provided, the model must implement
+            `get_quantization_layer_structure`.
+
+    Example:
+    ```python
+    from keras.quantizers import AWQConfig
+
+    # Create configuration for 4-bit AWQ quantization
+    config = AWQConfig(
+        dataset=calibration_data,          # Your calibration dataset
+        tokenizer=your_tokenizer,          # Tokenizer for text data
+        num_samples=128,                   # Number of calibration samples
+        sequence_length=512,               # Sequence length for each sample
+        group_size=128,                    # Weight grouping for quantization
+        num_grid_points=20,                # Grid search points for scale search
+    )
+
+    # Apply quantization to your model
+    model.quantize("awq", config=config)
+    ```
+
+    """
+
+    def __init__(
+        self,
+        dataset,
+        tokenizer,
+        *,
+        weight_bits: int = 4,
+        num_samples: int = 128,
+        sequence_length: int = 512,
+        group_size: int = 128,
+        num_grid_points: int = 20,
+        quantization_layer_structure: dict = None,
+    ):
+        super().__init__()
+        # AWQ only supports 4-bit quantization
+        if weight_bits != 4:
+            raise ValueError(
+                f"AWQ only supports 4-bit quantization. "
+                f"Received weight_bits={weight_bits}."
+            )
+        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 group_size < -1 or group_size == 0:
+            raise ValueError(
+                "Invalid group_size. Supported values are -1 (per-channel) "
+                f"or a positive integer, but got {group_size}."
+            )
+        if num_grid_points <= 0:
+            raise ValueError("num_grid_points must be a positive integer.")
+
+        self.dataset = dataset
+        self.tokenizer = tokenizer
+        self.weight_bits = weight_bits
+        self.num_samples = num_samples
+        self.sequence_length = sequence_length
+        self.group_size = group_size
+        self.num_grid_points = num_grid_points
+        self.quantization_layer_structure = quantization_layer_structure
+
+    @property
+    def mode(self):
+        return "awq"
+
+    def dtype_policy_string(self):
+        """Returns the dtype policy string for this configuration.
+
+        Returns:
+            A string representing the dtype policy, e.g. "awq/4/128".
+        """
+        return f"awq/{self.weight_bits}/{self.group_size}"
+
+    def get_config(self):
+        return {
+            # Dataset and Tokenizer are only required for one-time
+            # calibration and are not saved in the config.
+            "dataset": None,
+            "tokenizer": None,
+            "weight_bits": self.weight_bits,
+            "num_samples": self.num_samples,
+            "sequence_length": self.sequence_length,
+            "group_size": self.group_size,
+            "num_grid_points": self.num_grid_points,
+            "quantization_layer_structure": self.quantization_layer_structure,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)

keras/src/quantizers/awq_core.py

@@ -0,0 +1,217 @@
+"""AWQ core functionality for layer-wise quantization.
+
+This module provides the orchestration logic for applying AWQ quantization
+to transformer models in a layer-by-layer fashion.
+"""
+
+from contextlib import contextmanager
+
+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 AWQDTypePolicy
+from keras.src.dtype_policies.dtype_policy_map import DTypePolicyMap
+from keras.src.quantizers.awq import AWQ
+from keras.src.quantizers.awq_config import AWQConfig
+from keras.src.quantizers.gptq_core import find_layers_in_block
+from keras.src.quantizers.gptq_core import get_dataloader
+from keras.src.quantizers.utils import should_quantize_layer
+
+
+@contextmanager
+def stream_activations(layers_map, awq_objects):
+    """Context manager to capture activations for AWQ calibration.
+
+    Temporarily patches layer.call methods to capture activation statistics
+    for computing per-channel scaling factors.
+
+    Args:
+        layers_map: Dict[str, Layer]. Mapping from layer names to layers.
+        awq_objects: Dict[str, AWQ]. Mapping from names to AWQ instances.
+
+    Yields:
+        None: The patched state is active only within the `with` block.
+    """
+    original_calls = {}
+
+    def create_hook(name, original_call_func):
+        def hook(*args, **kwargs):
+            inp = args[0] if args else kwargs["inputs"]
+            num_features = awq_objects[name].rows
+            input_2d = ops.reshape(inp, (-1, num_features))
+            awq_objects[name].update_activation_magnitudes(input_2d)
+            return original_call_func(*args, **kwargs)
+
+        return hook
+
+    try:
+        for name, layer in layers_map.items():
+            original_calls[name] = layer.call
+            layer.call = create_hook(name, layer.call)
+        yield
+    finally:
+        for name, layer in layers_map.items():
+            layer.call = original_calls[name]
+
+
+def apply_awq_layerwise(dataloader, config, structure, filters=None):
+    """Apply AWQ quantization layer-by-layer to a Keras model.
+
+    This function processes the model sequentially, one block at a time:
+    1. Captures activation statistics through calibration data forward pass
+    2. Uses activation magnitudes to determine weight saliency
+    3. Finds optimal per-channel scales via grid search
+    4. Quantizes weights with AWQ scaling
+
+    Args:
+        dataloader: Calibration data as numpy array.
+        config: AWQConfig instance.
+        structure: Dict with 'pre_block_layers' and 'sequential_blocks'.
+        filters: Optional layer filters.
+    """
+    num_samples = config.num_samples
+    logging.info("Starting AWQ quantization...")
+
+    pre_layers = structure.get("pre_block_layers", [])
+    transformer_blocks = structure.get("sequential_blocks", [])
+
+    if not transformer_blocks:
+        raise ValueError(
+            "No sequential blocks found in the structure to quantize."
+        )
+
+    # Process inputs through pre-block layers (e.g., embedding)
+    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))
+
+    for block_idx, block in enumerate(transformer_blocks):
+        logging.info(f"Quantizing Block {block_idx}")
+        sub_layers_map = find_layers_in_block(block)
+
+        # Apply filters
+        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 quantizable layers found in block {block_idx}. Skipping."
+            )
+        else:
+            logging.info(f"Found layers: {list(sub_layers_map.keys())}")
+
+            # Create AWQ objects for each layer
+            awq_objects = {
+                name: AWQ(layer, config)
+                for name, layer in sub_layers_map.items()
+            }
+
+            # Capture activation statistics
+            with stream_activations(sub_layers_map, awq_objects):
+                for sample_idx in range(num_samples):
+                    current_input = inputs[sample_idx]
+                    if len(current_input.shape) == 2:
+                        current_input = ops.expand_dims(current_input, axis=0)
+                    _ = block(current_input)
+
+            # Quantize each layer
+            for name, awq_object in awq_objects.items():
+                logging.info(f"Quantizing {name}...")
+                awq_object.quantize_layer()
+                awq_object.free()
+
+            del awq_objects
+
+        # Generate inputs for next block
+        if block_idx < len(transformer_blocks) - 1:
+            logging.info(f"Generating inputs for block {block_idx + 1}...")
+            next_block_inputs = []
+            for sample_idx in range(num_samples):
+                current_input = inputs[sample_idx]
+                if len(current_input.shape) == 2:
+                    current_input = ops.expand_dims(current_input, axis=0)
+                output = block(current_input)[0]
+                next_block_inputs.append(output)
+            inputs = next_block_inputs
+
+        progbar.update(current=block_idx + 1)
+
+    logging.info("AWQ quantization complete.")
+
+
+def awq_quantize(config, quantization_layer_structure, filters=None):
+    """Main entry point for AWQ quantization.
+
+    Args:
+        config: AWQConfig instance.
+        quantization_layer_structure: Model structure dictionary.
+        filters: Optional layer filters.
+    """
+    if config.dataset is None or config.tokenizer is None:
+        raise ValueError(
+            "AWQ quantization requires a dataset and tokenizer. "
+            "Please provide them in the AWQConfig."
+        )
+
+    if quantization_layer_structure is None:
+        raise ValueError(
+            "For 'awq' 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 calibration data
+    dataloader = get_dataloader(
+        config.tokenizer,
+        config.sequence_length,
+        config.dataset,
+        num_samples=config.num_samples,
+    )
+
+    apply_awq_layerwise(
+        dataloader[: config.num_samples],
+        config,
+        quantization_layer_structure,
+        filters=filters,
+    )
+
+
+def get_group_size_for_layer(layer, config):
+    """Get group size from config or dtype policy.
+
+    Args:
+        layer: The layer to get group size for.
+        config: Optional AWQConfig instance.
+
+    Returns:
+        int: The group size for quantization.
+
+    Raises:
+        ValueError: If group size cannot be determined.
+    """
+    if config and isinstance(config, AWQConfig):
+        return config.group_size
+    elif isinstance(layer.dtype_policy, AWQDTypePolicy):
+        return layer.dtype_policy.group_size
+    elif isinstance(layer.dtype_policy, DTypePolicyMap):
+        policy = layer.dtype_policy[layer.path]
+        if isinstance(policy, AWQDTypePolicy):
+            return policy.group_size
+    raise ValueError(
+        "For AWQ quantization, group_size must be specified "
+        "through AWQConfig or AWQDTypePolicy."
+    )

keras/src/quantizers/gptq.py

@@ -1,5 +1,4 @@
 import types
-from functools import partial
 
 from keras.src import ops
 from keras.src import quantizers
@@ -296,12 +295,7 @@ class GPTQ:
             # For EinsumDense, we determine the effective 2D dimensions.
             self.kernel_shape = layer.kernel.shape
             shape = list(self.kernel_shape)
-            try:
-                d_model_dim_index = shape.index(max(shape))
-            except ValueError:
-                raise TypeError(
-                    f"Could not determine hidden dimension from shape {shape}"
-                )
+            d_model_dim_index = shape.index(max(shape))
 
             if d_model_dim_index == 0:  # QKV projection case
                 in_features, heads, head_dim = shape
@@ -471,7 +465,7 @@ class GPTQ:
             group_size=self.config.group_size,
             activation_order=self.config.activation_order,
             order_metric=ops.diagonal(hessian_matrix),
-            compute_scale_zero=partial(self.quantizer.find_params, weight=True),
+            compute_scale_zero=self.quantizer.find_params,
         )
         quantized = ops.cast(
             quantized, self.original_layer.quantized_kernel.dtype

keras/src/quantizers/gptq_config.py

@@ -1,8 +1,9 @@
 from keras.src.api_export import keras_export
+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.
 
@@ -131,6 +132,12 @@ 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__(
@@ -146,7 +153,9 @@ class GPTQConfig:
         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}. "
@@ -174,6 +183,32 @@ class GPTQConfig:
         self.group_size = group_size
         self.symmetric = symmetric
         self.activation_order = activation_order
+        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.

keras/src/quantizers/gptq_core.py

@@ -10,9 +10,9 @@ 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
@@ -131,7 +131,7 @@ def get_dataloader(
     pieces = []
     if isinstance(dataset_list[0], str):
         for i, s in enumerate(dataset_list):
-            toks = np.asarray(tokenizer.tokenize(s)).reshape(-1)
+            toks = ops.convert_to_numpy(tokenizer.tokenize(s)).reshape(-1)
             pieces.append(toks)
             # avoid windows that span document boundaries
             if eos_id is not None and i < len(dataset_list) - 1:
@@ -193,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.
@@ -242,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
@@ -284,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))
@@ -316,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())}")
@@ -357,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
@@ -376,7 +357,12 @@ 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):