Adds quantization documentation

Author: JyotinderSingh

guides/md/quantization/overview.md

@@ -0,0 +1,149 @@
+# Quantization in Keras
+
+**Author:** [Jyotinder Singh](https://x.com/Jyotinder_Singh)<br>
+**Date created:** 2025/10/09<br>
+**Last modified:** 2025/10/09<br>
+**Description:** Overview of quantization in Keras (int8, float8, int4, GPTQ).
+
+
+<img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> [**View in Colab**](https://colab.research.google.com/github/keras-team/keras-io/blob/master/guides/ipynb/quantization/overview.ipynb)  <span class="k-dot">•</span><img class="k-inline-icon" src="https://github.com/favicon.ico"/> [**GitHub source**](https://github.com/keras-team/keras-io/blob/master/guides/quantization/overview.py)
+
+---
+
+## Introduction
+
+Modern large models are often **memory- and bandwidth-bound**: most inference time is spent moving tensors between memory and compute units rather than doing math. Quantization reduces the number of bits used to represent the model's weights and (optionally) activations, which:
+
+* Shrinks model size and VRAM/RAM footprint.
+* Increases effective memory bandwidth (fewer bytes per value).
+* Can improve throughput and sometimes latency on supporting hardware with low-precision kernels.
+
+Keras provides first-class **post-training quantization (PTQ)** workflows which support pretrained models and expose a uniform API at both the model and layer level.
+
+At a high level, Keras supports:
+
+* Joint weight + activation PTQ in `int4`, `int8`, and `float8`.
+* Weight-only PTQ via **GPTQ** (2/3/4/8-bit) to maximize compression with minimal accuracy impact, especially for large language models (LLMs).
+
+> **Terminology**
+>
+> * *Scale / zero-point:* Quantization maps real values `x` to integers `q` using a scale (and optionally a zero-point). Symmetric schemes use only a scale.
+> * *Per-channel vs per-tensor:* A separate scale per output channel (e.g., per hidden unit) usually preserves accuracy better than a single scale for the whole tensor.
+> * *Calibration:* A short pass over sample data to estimate activation ranges (e.g., max absolute value).
+
+---
+
+## Quantization Modes
+
+Keras currently focuses on the following numeric formats. Each mode can be applied selectively to layers or to the whole model via the same API.
+
+* **`int8` (8-bit integer)**: **joint weight + activation** PTQ.
+
+  * **How it works:** Values are linearly mapped to 8-bit integers with per-channel scales. Activations are calibrated using dynamic quantization (see note below).
+  * **Why use it:** Good accuracy for many architectures; broad hardware support.
+  * **What to expect:** ~4x smaller than FP32 parameters (~2x vs FP16) and lower activation bandwidth, with small accuracy loss on many tasks. Throughput gains depend on kernel availability and memory bandwidth.
+
+* **`float8` (FP8: E4M3 / E5M2 variants)**: Low-precision floating-point useful for training and inference on FP8-capable hardware.
+
+  * **How it works:** Values are quantized to FP8 with a dynamic scale. Fused FP8 kernels on supported devices yield speedups.
+  * **Why use it:** Mixed-precision training/inference with hardware acceleration while keeping floating-point semantics (since underflow/overflow characteristics differ from int).
+  * **What to expect:** Competitive speed and memory reductions where FP8 kernels are available; accuracy varies by model, but is usually acceptable for most tasks.
+
+* **`int4`**: Ultra-low-bit **weights** for aggressive compression; activations remain in higher precision (int8).
+
+  * **How it works:** Two signed 4-bit "nibbles" are packed per int8 byte. Keras uses symmetric per-output-channel scales to dequantize efficiently inside matmul.
+  * **Why use it:** Significant VRAM/storage savings for LLMs with acceptable accuracy when combined with robust per-channel scaling.
+  * **What to expect:** ~8× smaller than FP32 (~4× vs FP16) for weights; throughput gains depend on kernel availability and memory bandwidth. Competitive accuracy deltas for encoder-only architectures, may show larger regressions on decoder-only models.
+
+* **`GPTQ` (weight-only 2/3/4/8 bits)**: *Second-order, post-training* method minimizing layer output error.
+
+  * **How it works (brief):** For each weight block (group), GPTQ solves a local least-squares problem using a Hessian approximation built from a small calibration set, then quantizes to low bit-width. The result is a packed weight tensor plus per-group parameters (e.g., scales).
+  * **Why use it:** Strong accuracy retention at very low bit-widths without retraining; ideal for rapid LLM compression.
+  * **What to expect:** Large storage/VRAM savings with small perplexity/accuracy deltas on many decoder-only models when calibrated on task-relevant samples.
+
+> **Implementation notes**
+>
+> * For `int4`, Keras packs signed 4-bit values (range ≈ [−8, 7]) and stores per-channel scales such as `kernel_scale`. Dequantization happens on the fly, and matmuls use 8-bit (unpacked) kernels.
+> * Activation scaling for `int4` / `int8` / `float8` uses **AbsMax calibration** by default (range set by the maximum absolute value observed). Alternative calibration methods (e.g., percentile) may be added in future releases.
+> * Per-channel scaling is the default for weights where supported, because it materially improves accuracy at negligible overhead.
+
+---
+
+## Quantizing Keras Models
+
+Quantization is applied explicitly after layers or models are built. The API is designed to be predictable: you call quantize, the graph is rewritten,the weights are replaced, and you can immediately run inference or save the model.
+
+Typical workflow:
+
+1. **Build / load your FP model.** Train if needed. Ensure `build()` or a forward pass has materialized weights.
+2. **(GPTQ only)** Keras may run a short calibration pass to collect activation ranges (you can pass a small, representative dataset).
+3. **Invoke quantization.** Call `model.quantize("<mode>")` or `layer.quantize("<mode>")` with `"int8"`, `"int4"`, `"float8"`, or `"gptq"` (weight-only).
+4. **Use or save.** Run inference, or `model.save(...)`. Quantization state (packed weights, scales, metadata) is preserved on save/load.
+
+### Model Quantization
+
+```python
+import keras
+import numpy as np
+
+# Sample training data
+x_train = keras.ops.array(np.random.rand(100, 10))
+y_train = keras.ops.array(np.random.rand(100, 1))
+
+# Build the model
+model = keras.Sequential([
+    keras.layers.Dense(32, activation="relu", input_shape=(10,)),
+    keras.layers.Dense(1)
+])
+
+# Compile and fit the model
+model.compile(optimizer="adam", loss="mean_squared_error")
+model.fit(x_train, y_train, epochs=1, verbose=0)
+
+# Quantize the model
+model.quantize("int8")
+```
+
+**What this does:** Quantizes the weights of the supported layers, and re-wires their forward paths to be compatible with the quantized kernels and quantization scales.
+
+**Note**: Throughput gains depend on backend/hardware kernels; in cases where kernels fall back to dequantized matmul, you still get memory savings but smaller speedups.
+
+### Layer-wise Quantization
+
+The Keras quantization framework allows you to quantize each layer separately, without having to quantize the entire model using the same unified API.
+
+```python
+from keras import layers
+
+input_shape = (10,)
+layer = layers.Dense(32, activation="relu", input_shape=input_shape)
+layer.build(input_shape)
+
+layer.quantize("int4")  # or "int8", "float8", etc.
+```
+
+**When to use layer-wise quantization**
+
+* To keep numerically sensitive blocks (e.g., small residual paths, logits) at higher precision while quantizing large projection layers.
+* To mix modes (e.g., attention projections in int4, feed-forward in int8) and measure trade-offs incrementally.
+* Always validate on a small eval set after each step; mixing precisions across residual connections can shift distributions.
+
+---
+
+## Layer & model coverage
+
+Keras supports the following core layers in its quantization framework:
+
+* `Dense`
+* `EinsumDense`
+* `Embedding` (available in KerasHub)
+* `ReversibleEmbedding` (available in KerasHub)
+
+Any composite layers that are built from the above (for example, `MultiHeadAttention`, `GroupedQueryAttention`, feed-forward blocks in Transformers) inherit quantization support by construction. This covers the majority of modern encoder-only and decoder-only Transformer architectures.
+
+Since all KerasHub models subclass `keras.Model`, they automatically support the `model.quantize(...)` API. In practice, this means you can take a popular LLM preset, call a single function to obtain an int8/int4/GPTQ-quantized variant, and then save or serve it—without changing your training code.
+
+> **Practical guidance**
+>
+> * For GPTQ, use a calibration set that matches your inference domain (a few hundred to a few thousand tokens is often enough to see strong retention).
+> * Measure both **VRAM** and **throughput/latency**: memory savings are immediate; speedups depend on the availability of fused low-precision kernels on your device.

guides/quantization/overview.py

@@ -0,0 +1,168 @@
+"""
+Title: Quantization in Keras
+Author: Keras Team
+Date created: 2025/10/09
+Last modified: 2025/10/09
+Description: Overview of weight & activation quantization workflows in Keras (int8, float8, int4, GPTQ) and integration patterns.
+Accelerator: None
+"""
+
+"""
+## Introduction
+
+Modern large models are memory- and bandwidth-bound. Quantization reduces model footprint
+and improves inference throughput by representing weights and (optionally) activations
+with fewer bits than the original floating-point formats.
+
+At a high level, Keras provides first-class support for multiple quantization modes:
+
+* Weight-only post-training quantization (e.g., GPTQ 4-bit) for rapid compression.
+* Joint weight & activation post-training quantization (PTQ) in `int8` or `float8`.
+* Quantization-aware training (QAT) paths (where supported) to recover accuracy when
+  naive PTQ would cause excessive degradation.
+
+Two dimensions to keep in mind:
+
+* What is being quantized? (weights only vs weights + activations)
+* When is quantization applied? (post-training vs during training via QAT)
+"""
+
+"""
+## Supported quantization modes
+
+Keras currently focuses on the following numeric formats:
+
+* int8 (8-bit integer) — Common baseline for activation + weight PTQ/QAT.
+* float8 (e.g., E4M3 / E5M2 variants) — Mixed-precision friendly low-bit floating formats
+  enabling smoother accuracy retention vs int8 for some transformer blocks.
+* int4 (packed) — Ultra-low-bit weight representation for large language models.
+  Activations generally remain in a higher precision (e.g., float16 / bfloat16 / int8).
+* GPTQ (weight-only 4-bit) — Post-training, block-wise second-order approximated
+  quantization minimizing output error. Ideal for rapid compression of pretrained LLMs.
+
+Behind the scenes, the `int4` path packs two 4-bit nibbles into each signed int8 byte.
+A per-output-channel (e.g., per hidden unit) `kernel_scale` tensor rescales the packed
+values back to a floating domain during matmul. Activation scaling (for int8/float8 modes)
+uses an AbsMax (max absolute value) calibration by default, but alternative calibration
+strategies may be pluggable in future updates.
+"""
+
+"""
+## Integration paths: policy-first vs imperative
+
+Keras offers two complementary integration styles so you can adopt quantization with
+minimal friction depending on how you construct models.
+
+### 1. Policy-first (declarative) approach
+
+Specify a quantization policy via `dtype` (or `compute_dtype`) when instantiating layers
+or updating them later. Examples:
+
+```python
+from keras import layers
+
+# Declare an int4 quantization transformation from float32 pretrained weights
+dense_int4 = layers.Dense(4096, activation="gelu", dtype="int4_from_float32")
+
+# Mixed path: start from mixed bfloat16 weights, produce int8 quantized variant
+dense_int8 = layers.Dense(4096, dtype="int8_from_mixed_bfloat16")
+```
+
+Semantics of a policy string:
+
+```
+<target_quant_dtype>_from_<source_dtype>
+```
+
+The layer will (lazily) quantize its weights the first time it's built or when weights
+are assigned.
+
+You can also retroactively apply a policy:
+
+```python
+layer.dtype = "int4_from_float32"  # triggers (re)quantization on next build/assign
+```
+
+### 2. Imperative (post-construction) approach
+
+Build your model/layer normally, then invoke an explicit quantization method:
+
+```python
+model = build_large_lm()  # returns a compiled or uncompiled keras.Model
+model.build(input_shape)  # ensure weights are created
+model.quantize("int4")    # in-place weight transformation (and bookkeeping)
+```
+
+For an individual layer:
+
+```python
+layer.build(input_shape)
+layer.quantize("gptq_int4")  # or "int8", "float8", etc.
+```
+
+This style is convenient for applying quantization to existing checkpoints without
+rewriting code.
+"""
+
+"""
+## Layer & model coverage
+
+Core dense projection layers are supported:
+
+* `Dense`
+* `EinsumDense` (including LoRA-compatible adaptations)
+
+At the model level, KerasHub provides helpers/presets (for popular LLM architectures)
+that expose a single `model.quantize(<mode>)` convenience method targeting all eligible
+submodules.
+
+Additional layer types (e.g., attention, convolution) may be incrementally covered as
+kernel implementations mature.
+"""
+
+"""
+## Under the hood: packed int4
+
+For `int4` / GPTQ weight-only quantization:
+
+1. Original float weights (e.g., float32 or bfloat16) are partitioned per output channel.
+2. A scale (and optionally zero-point / offset for asymmetric schemes) is computed per channel.
+3. Each weight is divided by its channel scale, clamped to the 4-bit signed range, and
+   rounded to nearest integer.
+4. Two 4-bit values are packed into one int8 storage byte (low nibble + high nibble).
+5. During a forward matmul, packed bytes are unpacked on-the-fly (SIMD / fused kernel),
+   multiplied by their channel scale, and accumulated in higher precision (e.g., float16).
+6. Activation quantization (if enabled) applies AbsMax scaling to produce an 8-bit tensor
+   fed into integer or mixed-precision GEMM kernels.
+
+Key advantages:
+
+* 4x compression vs float16; often ~8x vs float32.
+* Reduced memory bandwidth → higher tokens/sec for LLM inference.
+* Channel-wise scaling preserves accuracy better than per-tensor scaling.
+
+Trade-offs:
+
+* Extra decode overhead; mitigated by fused kernels.
+* Some accuracy regression vs int8/float8, especially for very small models.
+"""
+
+"""
+## Where to go next
+
+* Quickstart with runnable LLM examples (coming soon)  
+* int4 internals & advanced configuration (link forthcoming)  
+* Exporting / deployment guide (ONNX / TFLite / Serving)  
+* GPTQ overview & best practices  
+* Quantization recipes & troubleshooting  
+
+(Links will be activated as the companion guides are published.)
+"""
+
+"""
+## Summary
+
+You can adopt Keras quantization either declaratively (policy-first `dtype` strings) or
+imperatively (`.quantize(mode)`). Start with `int8` / `float8` for balanced accuracy &
+performance, and explore `int4` / GPTQ for maximum compression of large language models.
+"""

guides/quantization_overview.py

@@ -123,6 +123,10 @@
             "path": "orbax_checkpoint",
             "title": "Orbax Checkpointing in Keras",
         },
+        {
+            "path": "quantization/overview",
+            "title": "Quantization in Keras",
+        },
         # {
         #     "path": "preprocessing_layers",
         #     "title": "Working with preprocessing layers",