Introduces QuantizationConfig for fine-grained quantization control

[JyotinderSingh]

Description

This PR introduces keras.quantizers.QuantizationConfig, a configuration class that allows users to customize the quantization behavior for weights and activations during Post-Training Quantization (PTQ).

Previously, calling model.quantize(mode="int8") applied a fixed quantization strategy (typically AbsMaxQuantizer for both weights and activations). With this change, users can instantiate a configuration object (e.g. Int8QuantizationConfig) to define specific quantizers for weights and activations, or disable activation quantization entirely for weight-only quantization.

Key Changes

QuantizationConfig Class

• Introduced keras.src.quantizers.quantization_config.py.
• Added subclasses for specific modes: Int8QuantizationConfig, Int4QuantizationConfig, and Float8QuantizationConfig.
• Holds weight_quantizer and activation_quantizer instances with serialization support (get_config/from_config).

Model & Layer Updates

• Model: Updated Model.quantize() to accept an optional QuantizationConfig object, allowing custom configurations beyond the default string modes.
• Layers: Refactored Dense, EinsumDense, Embedding, and ReversibleEmbedding.
  • quantized_build and quantize methods now accept and parse the provided config.
  • Added logic to handle activation_quantizer=None, enabling weight-only quantization.
  • Updated grad_fn in Dense and EinsumDense to correctly handle cases where inputs are not quantized (i.e., when inputs_scale is effectively 1.0).

Quantizer Updates

• Updated AbsMaxQuantizer.__call__ to accept to_numpy=True to support direct weight manipulation during the quantization step.

Usage Examples

1. Default Quantization (Int8)
   Applies the default AbsMaxQuantizer to both weights and activations.

model.quantize("int8")

2. Weight-Only Quantization (Int8)
   Disable activation quantization by setting the activation quantizer to None.

from keras.quantizers import Int8QuantizationConfig, AbsMaxQuantizer

config = Int8QuantizationConfig(
    weight_quantizer=AbsMaxQuantizer(axis=0),
    activation_quantizer=None 
)

model.quantize(config=config)

3. Custom Quantization Parameters
   Customize the value range or other parameters for specific quantizers.

config = Int8QuantizationConfig(
    # Restrict range for symmetric quantization
    weight_quantizer=AbsMaxQuantizer(axis=0, value_range=(-127, 127)),
    activation_quantizer=AbsMaxQuantizer(axis=-1, value_range=(-127, 127))
)

model.quantize(config=config)

Tests

• Added QuantizationConfigTest for serialization/deserialization.
• Updated layer tests (DenseTest, EinsumDenseTest, EmbeddingTest, ReversibleEmbeddingTest) to verify behavior with QuantizationConfig, specifically testing weight-only scenarios.
• Updated GPTQTest to ensure config validation logic remains correct.

[gemini-code-assist]

Summary of Changes

Hello @JyotinderSingh, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the quantization capabilities within Keras by introducing a flexible QuantizationConfig API. This new configuration object allows users to specify custom quantizers for both weights and activations, moving beyond predefined quantization settings. The change integrates this new API into key layers like Dense, EinsumDense, Embedding, and ReversibleEmbedding, and refactors GPTQConfig to build upon this new base. This provides a more extensible and user-friendly approach to model quantization, allowing for more tailored optimization strategies.

Highlights

• New Quantization API: Introduced a new QuantizationConfig API that allows for customizable weight and activation quantizers, providing greater flexibility in defining quantization schemes.
• Layer Integration: Core layers such as Dense, EinsumDense, Embedding, and ReversibleEmbedding have been updated to accept and utilize the new QuantizationConfig object, enabling fine-grained control over their quantization behavior.
• GPTQConfig Inheritance: The GPTQConfig now inherits from QuantizationConfig, streamlining the quantization configuration hierarchy and ensuring consistency across different quantization methods.
• Model Quantization Method Update: The quantize method in the Model class has been enhanced to accept the config argument for all quantization modes, removing the previous restriction that it was only supported for GPTQ mode.
• Weight-Only Quantization Support: Added explicit handling for weight-only quantization scenarios where inputs_quantizer might be None, ensuring correct scaling in such cases.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a QuantizationConfig to provide a more flexible and customizable quantization API. This is a significant improvement, allowing users to specify their own quantizers for weights and activations, and enabling features like weight-only quantization. The changes are well-implemented across various layers including Dense, EinsumDense, Embedding, and ReversibleEmbedding, as well as the model-level quantize method. The new QuantizationConfig class is well-designed with serialization support, and the accompanying tests are comprehensive. I have a couple of suggestions for minor code improvements to reduce redundancy and enhance clarity.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 90.54054% with 21 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.33%. Comparing base (f0a48a6) to head (27da291).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/quantizers/quantization_config.py	90.21%	4 Missing and 5 partials ⚠️
keras/src/layers/core/dense.py	76.00%	4 Missing and 2 partials ⚠️
keras/api/_tf_keras/keras/quantizers/__init__.py	0.00%	4 Missing ⚠️
keras/src/layers/core/reversible_embedding.py	92.30%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21896      +/-   ##
==========================================
+ Coverage   76.30%   76.33%   +0.03%     
==========================================
  Files         580      581       +1     
  Lines       60029    60186     +157     
  Branches     9432     9461      +29     
==========================================
+ Hits        45803    45944     +141     
- Misses      11750    11759       +9     
- Partials     2476     2483       +7     

Flag	Coverage Δ
keras	76.20% <89.63%> (+0.03%)	⬆️
keras-jax	62.18% <86.93%> (+0.05%)	⬆️
keras-numpy	57.39% <86.03%> (+0.07%)	⬆️
keras-openvino	34.30% <31.53%> (+<0.01%)	⬆️
keras-torch	63.28% <86.03%> (+0.05%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[JyotinderSingh]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces QuantizationConfig to provide a more structured and fine-grained control over quantization settings. The changes are well-implemented across various layers and the new configuration class is well-designed. I've found a couple of minor issues related to an unused parameter and an outdated docstring that should be addressed.