Skip to content

Latest commit

 

History

History
188 lines (150 loc) · 8.59 KB

README.md

File metadata and controls

188 lines (150 loc) · 8.59 KB

quantization

Overview

The quantization module provides a unified, modular interface for quantizing neural networks, including large language models (LLMs) and vision-language models (VLMs). Its primary goal is to simplify the process of preparing, calibrating, and converting models into their quantized versions while abstracting the underlying complexities of different quantization algorithms. At the top level, the module exposes two primary public interface functions: prepare and convert.

  • Public Interface

    • prepare(model, quant_config, args, kwargs, inplace)
      • Prepares the given model for quantization based on the provided algorithm-specific configuration. This involves setting up necessary observers or hooks, and may optionally use example inputs (provided as separate positional and keyword arguments), which is particularly useful for activation quantization.
    • convert(model, inplace=True)
      • Converts the model that has been prepared and calibrated into its quantized form. It retrieves the quantizer attached by prepare() and uses the existing configuration and calibration state. It does not accept quant_config or select or instantiate a new quantizer. PTQ wrappers retain their resolved quantization policies.
  • Algorithm-Specific Configuration

    • Users supply algorithm-specific configuration objects (e.g., PTQConfig for activation quantization or GPTQConfig for weight quantization) to prepare(). These configuration objects encapsulate the parameters for each quantization method.
  • Internal Dispatch to Quantizer Implementations

    • prepare() selects the algorithm-specific quantizer through get_quantizer(quant_config) and attaches it to the returned prepared model as tico_quantizer. That quantizer owns its configuration; convert() reuses the attached instance.

This design ensures that users interact only with the high-level functions and configuration objects, while the internal details of each quantization algorithm are encapsulated within dedicated quantizer classes.

After successful conversion, the returned model no longer carries tico_quantizer. If conversion fails, the original prepared model retains the handle for retry. Retaining the handle does not roll back arbitrary in-place algorithm mutations. For algorithms supporting convert(prepared_model, inplace=False), the converted copy has no handle and the original prepared model keeps its handle even after success. GPTQ continues to require inplace=True.

Basic Usage

from tico.quantization import prepare, convert
from tico.quantization.config.gptq import GPTQConfig
import torch
import torch.nn as nn

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(8, 8)

    def forward(self, x):
        return self.linear(x)

model = LinearModel().eval()

# 1. Prepare for quantization
quant_config = GPTQConfig()
prepared_model = prepare(model, quant_config)

# 2. Calibration
for d in dataset:
    prepared_model(d)

# 3. Apply GPTQ
quantized_model = convert(prepared_model)

Purpose

The module is designed to:

  • Simplify Model Quantization: Provide a simple, high-level API that abstracts the calibration and conversion complexities.
  • Support Multiple Algorithms: Enable users to choose quantization algorithms by specifying configuration objects (e.g., PTQConfig, GPTQConfig).
  • Promote Modularity and Extensibility: Offer a clear separation between the public API, configuration classes, and internal quantizer implementations, ensuring the module is easy to extend and maintain.

Directory Structure

quantization/
├── algorithm/        # Algorithm-specific quantizers
│   ├── cle/          #   Cross-Layer Equalization
│   ├── fpi_gptq/     #   Fixed-point-iteration GPTQ variant
│   ├── gptq/         #   GPTQ weight quantization
│   ├── qwen3_vl_gptq/#   GPTQ specialization for Qwen3-VL
│   ├── smoothquant/  #   SmoothQuant
│   └── spinquant/    #   SpinQuant (rotation-based)
├── analysis/         # Numerical ablation, clipping, and sensitivity
├── config/           # Configuration definitions (PTQConfig, GPTQConfig, ...)
├── evaluation/       # Evaluation utilities
├── examples/         # Config-driven CLI examples (quantize / evaluate / inspect)
├── passes/           # Graph-level passes and transformations
├── recipes/          # Reusable pipeline layer behind the example CLIs
└── wrapq/            # Wrapper-based PTQ infrastructure (observers, wrappers)

The module is composed of two main implementation layers:

  • Algorithm layer (quantization/algorithm/)
    • Implements algorithm-specific quantization logic (e.g., GPTQ, SmoothQuant, SpinQuant, CLE).
  • Infrastructure layer (quantization/wrapq/)
    • Provides a generic, wrapper-based quantization backend shared across algorithms.

On top of these, two layers make the toolkit usable end-to-end from the command line:

  • Recipes (quantization/recipes/)
    • The reusable pipeline layer: model-family adapters, algorithm stages, calibration data builders, evaluation, export, and debug helpers.
  • Examples (quantization/examples/)
    • Thin CLI entry points (quantize.py, evaluate.py, inspector.py) driven by YAML presets in examples/configs/.

Documentation Map

Document Description
Quantized Circle design Fake-quant markers, qparam folding, requantization, and the NPU tensor contract
analysis/README.md Numerical PTQ ablation, clipping, and sensitivity APIs
algorithm/README.md Algorithm module design and usage guidelines
algorithm/gptq/README.md GPTQ
algorithm/smoothquant/README.md SmoothQuant
algorithm/spinquant/README.md SpinQuant
algorithm/cle/README.md Cross-Layer Equalization
wrapq/README.md Wrapper-based PTQ infrastructure
recipes/README.md Recipes developer guide (adapters, stages, import rules)
examples/README.md CLI examples: quantize, evaluate, inspect
examples/configs/README.md YAML config preset reference

Contributing: Adding a New Algorithm

We welcome contributions that enhance the module by adding support for new quantization algorithms. If you are interested in contributing a new algorithm, please follow these steps:

  1. Implement New Configuration and Quantizer Classes
  • Configuration

    • Create a new configuration class (e.g., NewAlgoConfig) that inherits from BaseConfig. Include all necessary parameters and default values specific to your algorithm.
  • Quantizer

    • Implement a new quantizer class (e.g., NewAlgoQuantizer) that inherits from BaseQuantizer. Implement the prepare and convert methods.
  1. Register the Quantizer

Decorate the quantizer with register_quantizer(NewAlgoConfig) from tico/quantization/quantizer_registry.py. prepare() resolves the quantizer by exact config type; configs whose name matches tico.quantization.algorithm.<name>.quantizer are also imported lazily on first use. Registration rules:

  • One config type maps to exactly one quantizer class. Re-registering the same class is a no-op, so repeated activation of an extension is safe.
  • Registering a different quantizer for an already mapped config type raises ValueError. Do not reuse a built-in config type for a new algorithm; define your own config class, including for quantizers shipped in separately installed packages.
  • registered_quantizer(config_cls) returns the current mapping without triggering the lazy import.
  1. Write Tests

Ensure that you add unit tests and integration tests for your new algorithm. Tests should cover both the prepare and convert phases and validate that the quantized model meets performance expectations.

  1. Document Your Changes

Update the README and any related documentation to include details about the new algorithm, its configuration parameters, and usage examples.

To expose the new algorithm through the CLI examples as a pipeline stage, follow the recipes developer guide — most new workflows should be added as YAML presets, not new scripts.