diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000..69686f2 --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,5 @@ +# Commits to ignore in git blame (formatting-only changes). +# GitHub picks this file up automatically; locally, enable it with: +# git config blame.ignoreRevsFile .git-blame-ignore-revs + +eb65ffb809108b784e2051767f0c0f3f9e672fb3 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dbee399..3192c21 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,19 +1,20 @@ +exclude: (^examples\/) repos: -- repo: https://github.com/psf/black - rev: 25.1.0 +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.15.20 hooks: - - id: black - language_version: python3 - args: ['--line-length=125', - '--skip-string-normalization'] + - id: ruff + args: [--fix] + - id: ruff-format - repo: https://github.com/tox-dev/pyproject-fmt - rev: v2.5.1 + rev: v2.25.1 hooks: - id: pyproject-fmt + args: ["--max-supported-python", "3.13"] - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + rev: v6.0.0 hooks: - id: check-added-large-files - id: check-case-conflict @@ -27,29 +28,26 @@ repos: - id: requirements-txt-fixer - id: trailing-whitespace -- repo: https://github.com/PyCQA/isort - rev: 6.0.1 - hooks: - - id: isort - args: ["--profile=black"] # <-- this one - - repo: https://github.com/asottile/pyupgrade - rev: v3.19.1 + rev: v3.21.2 hooks: - id: pyupgrade - args: ["--py36-plus"] + args: ["--py310-plus"] - repo: https://github.com/pycqa/flake8 - rev: 7.1.2 + rev: 7.3.0 hooks: - id: flake8 - exclude: docs/conf.py + exclude: docs/source/conf.py additional_dependencies: [flake8-bugbear, flake8-print] args: ['--max-line-length=125', # github viewer width - '--extend-ignore=E203'] # E203 is not PEP8 compliant + '--extend-ignore=E203,T201,F401', + # E203 is not PEP8 compliant + # F401 included in ruff (behaves slightly differently for noqa flags) + ] - repo: https://github.com/mgedmin/check-manifest - rev: "0.50" + rev: "0.51" hooks: - id: check-manifest stages: [manual] diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 575f578..8053304 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -12,12 +12,11 @@ build: # Build documentation in the "docs/" directory with Sphinx sphinx: - configuration: docs/conf.py + configuration: docs/source/conf.py # Optionally, but recommended, # declare the Python requirements required to build your documentation # See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html -# python: -# install: -# - requirements: docs/requirements.txt - +python: + install: + - requirements: docs/requirements.txt diff --git a/README.md b/README.md index d5d245b..ed4f679 100644 --- a/README.md +++ b/README.md @@ -1,61 +1,118 @@ -![alt text](docs/source/_static/pquant_white_font.png) +

+ PQuantML logo +

-## Prune and Quantize ML models -PQuant is a library for training compressed machine learning models, developed at CERN as part of the [Next Generation Triggers](https://nextgentriggers.web.cern.ch/t13/) project. +# PQuantML -Installation via pip: ```pip install pquant-ml```. +**PQuantML** is an end-to-end library for training compressed machine learning models, developed at CERN as part of the [Next Generation Triggers](https://nextgentriggers.web.cern.ch/t13/) project. -With TensorFlow ```pip install pquant-ml[tensorflow]```. +It supports: -With PyTorch ```pip install pquant-ml[torch]```. +- Joint pruning + quantization +- Layer-wise precision configuration +- Flexible training pipelines +- PyTorch and TensorFlow backends +- Knowledge distillation +- HGQ library integration +- Integration with hardware-friendly toolchains (e.g., [hls4ml](https://fastmachinelearning.org/hls4ml/)) -PQuant replaces the layers and activations it finds with a Compressed (in the case of layers) or Quantized (in the case of activations) variant. These automatically handle the quantization of the weights, biases and activations, and the pruning of the weights. -Both PyTorch and TensorFlow models are supported. +PQuantML enables efficient deployment of compact neural networks on resource-constrained hardware such as FPGAs and embedded accelerators. -### Layers that can be compressed +

+ PQuantML overview +

-* **PQConv*D**: Convolutional layers -* **PQAvgPool*D**: Average pooling layers -* **PQBatchNorm*D**: BatchNorm layers -* **PQDense**: Linear layer -* **PQActivation**: Activation layers (ReLU, Tanh) +## Installation -The various pruning methods have different training steps, such as a pre-training step and fine-tuning step. PQuant provides a training function, where the user provides the functions to train and validate an epoch, and PQuant handles the training while triggering the different training steps. +Install the base package via pip: +```bash +pip install pquant-ml +``` -![alt text](docs/source/_static/overview_pquant.png) +Install with a specific backend: +```bash +pip install "pquant-ml[tensorflow]" # TensorFlow backend +pip install "pquant-ml[torch]" # PyTorch backend +``` +## Supported layers -### Example -Example notebook can be found [here](https://github.com/nroope/PQuant/tree/main/examples). It handles the - 1. Creation of a torch model and data loaders. - 2. Creation of the training and validation functions. - 3. Loading a default pruning configuration of a pruning method. - 4. Using the configuration, the model, and the training and validation functions, call the training function of PQuant to train and compress the model. - 5. Creating a custom quantization and pruning configuration for a given model (disable pruning for some layers, different quantization bitwidths for different layers). - 6. Direct layers usage and layers replacement approaches. - 7. Usage of fine-tuning platform. +| Layer | Description | +| --- | --- | +| `PQConv*D` | Convolutional layers | +| `PQAvgPool*D` | Average pooling layers | +| `PQBatchNorm*D` | Batch normalization layers | +| `PQDense` | Linear (fully connected) layer | +| `PQActivation` | Activation layers: ReLU, Tanh, Leaky ReLU, GELU, Hard Tanh, or a user-provided activation function (Torch only) | +| `MultiHeadAttention` | Multi-head attention layer | +| `LayerNorm` | Layer normalization layer (currently Torch only) | -### Pruning methods -A description of the pruning methods and their hyperparameters can be found [here](docs/pruning_methods.md). +## Training -### Quantization parameters -A description of the quantization parameters can be found [here](docs/quantization_parameters.md). +Different pruning methods involve different training stages, such as pre-training and fine-tuning. PQuantML provides a generic training function: you supply your own training and validation functions along with the number of epochs, and PQuant handles the training loop while automatically triggering the appropriate stages for the chosen pruning method. +

+ Pruning methods overview. +

+## Quantization -For detailed documentation check this page: [PQuantML documentation](https://pquantml.readthedocs.io/en/latest/) +PQuantML supports two quantization modes, each with several granularity options. +**Fixed-point quantization** (for weights): +- per-weight +- per-channel +- per-tensor -### Authors - - Roope Niemi (CERN) - - Anastasiia Petrovych (CERN) - - Arghya Das (Purdue University) - - Enrico Lupi (CERN) - - Chang Sun (Caltech) - - Dimitrios Danopoulos (CERN) - - Marlon Joshua Helbing - - Mia Liu (Purdue University) - - Michael Kagan (SLAC National Accelerator Laboratory) - - Vladimir Loncar (CERN) - - Maurizio Pierini (CERN) +**HGQ (High Granularity Quantization):** +- per-weight +- per-tensor + +## Example + +Example notebooks are available in the [`examples/`](https://github.com/cern-nextgen/PQuantML/tree/main/examples) directory. It shows how to: + +1. Create a Torch model and data loaders. +2. Create the training and validation functions. +3. Load a default configuration for a pruning method. +4. Train and compress the model by passing the configuration, model, and training/validation functions to PQuant's training function. +5. Build a custom quantization and pruning configuration for a given model (e.g. disabling pruning for some layers, or using different quantization bit-widths per layer). +6. Use the direct-layer and layer-replacement approaches. +7. Use the HPO platform. + +## Documentation + +Full documentation is available at [pquantml.readthedocs.io](https://pquantml.readthedocs.io/en/latest/). + +## Citation + +The framework is described in **PQuantML: A Tool for End-to-End Hardware-aware Model Compression** ([arXiv:2603.26595](https://arxiv.org/abs/2603.26595)). + +If you use PQuantML in your work, please cite: + +```bibtex +@article{niemi2026pquantml, + title = {PQuantML: A Tool for End-to-End Hardware-aware Model Compression}, + author = {Niemi, Roope and Petrovych, Anastasiia and Das, Arghya and + Lupi, Enrico and Sun, Chang and Danopoulos, Dimitrios and + Helbing, Marlon Joshua and Liu, Mia and Kagan, Michael and + Loncar, Vladimir and Pierini, Maurizio}, + journal = {arXiv preprint arXiv:2603.26595}, + year = {2026} +} +``` + +## Authors + +- Roope Niemi (CERN) +- Anastasiia Petrovych (CERN) +- Arghya Das (Purdue University) +- Enrico Lupi (CERN) +- Chang Sun (Caltech) +- Dimitrios Danopoulos (CERN) +- Marlon Joshua Helbing +- Mia Liu (Purdue University) +- Michael Kagan (SLAC National Accelerator Laboratory) +- Vladimir Loncar (CERN) +- Maurizio Pierini (CERN) diff --git a/docs/Makefile b/docs/Makefile index 5647f38..e88d665 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -18,4 +18,4 @@ help: # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile @sphinx-apidoc -f -o autodoc/ ../src/HGQ - @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) \ No newline at end of file + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/requirements.txt b/docs/requirements.txt index 726ada1..950c059 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1,5 +1,5 @@ -sphinx furo myst-parser -sphinx_rtd_theme +sphinx sphinx-autodoc-typehints +sphinx_rtd_theme diff --git a/docs/source/_static/custom.css b/docs/source/_static/custom.css index 3300cdf..ee944c1 100644 --- a/docs/source/_static/custom.css +++ b/docs/source/_static/custom.css @@ -9,13 +9,13 @@ html body nav.wy-nav-top, } .wy-nav-content { - max-width: 2000px !important; + max-width: none !important; width: 100% !important; } .rst-content, .rst-content .section { - max-width: 2000px !important; + max-width: none!important; width: 100% !important; } @@ -47,16 +47,16 @@ html body nav.wy-nav-top, .rst-content .literal-block, .rst-content .line-block, .rst-content .topic { - max-width: 2000px !important; + max-width: none !important; } .wy-body-for-nav { - background-color: #ffffff !important; + background-color: #ffffff !important; } .wy-nav-content { background-color: #ffffff !important; - max-width: 1200px !important; + max-width: none !important; } .wy-side-nav-search { @@ -71,19 +71,19 @@ html body nav.wy-nav-top, } .wy-nav-side { - background-color: #b30000 !important; + background-color: #b30000 !important; } .wy-menu-vertical a { - color: #ffffff !important; + color: #ffffff !important; } .wy-menu-vertical a:hover { - background-color: #990000 !important; + background-color: #990000 !important; } .wy-menu-vertical li.current > a, .wy-menu-vertical li.toctree-l1.current > a { - background-color: #660000 !important; + background-color: #660000 !important; color: #ffffff !important; } @@ -110,7 +110,7 @@ h1, h2, h3, h4, h5, h6 { background-color: #ffffff !important; } -.rst-content > .document > .toctree-wrapper, +.rst-content > .document > .toctree-wrapper, .rst-content > .document > .section { background-color: #ffffff !important; } @@ -121,7 +121,7 @@ h1, h2, h3, h4, h5, h6 { } .rst-content table th { - background-color: #ffe6e6 !important; + background-color: #ffe6e6 !important; color: #b30000 !important; } @@ -158,7 +158,7 @@ h1, h2, h3, h4, h5, h6 { color: #000000 !important; padding: 8px !important; margin: 12px !important; - width: calc(100% - 24px) !important; + width: calc(100% - 24px) !important; box-sizing: border-box !important; } @@ -170,7 +170,7 @@ h1, h2, h3, h4, h5, h6 { .wy-side-nav-search { background-color: #ffffff !important; - padding: 20px 15px !important; + padding: 20px 15px !important; border-bottom: 2px solid #b30000 !important; } @@ -184,19 +184,19 @@ h1, h2, h3, h4, h5, h6 { } .wy-side-nav-search .version-switch { - color: #b30000 !important; + color: #b30000 !important; font-weight: 600 !important; } .wy-side-nav-search .version-switch :hover { - color: #990000 !important; + color: #990000 !important; } .wy-side-nav-search .fa-caret-down { color: #b30000 !important; } .wy-side-nav-search select.version-switch { - color: #b30000 !important; + color: #b30000 !important; background-color: #ffffff !important; border: 2px solid #b30000 !important; font-weight: 600 !important; @@ -205,8 +205,8 @@ h1, h2, h3, h4, h5, h6 { } .wy-side-nav-search select.rtd-version-select { - color: #b30000 !important; - background-color: #ffffff !important; + color: #b30000 !important; + background-color: #ffffff !important; border: 2px solid #b30000 !important; font-weight: 600 !important; padding: 6px !important; @@ -223,16 +223,16 @@ h1, h2, h3, h4, h5, h6 { } .wy-side-nav-search input[type="search"] { - height: 100% !important; - width: 100% !important; - font-size: 18px !important; - padding: 10px 14px !important; + height: 100% !important; + width: 100% !important; + font-size: 18px !important; + padding: 10px 14px !important; margin: 0 !important; - border-radius: 6px !important; - border: none !important; - box-shadow: none !important; - background: #ffffff !important; - color: #000000 !important; + border-radius: 6px !important; + border: none !important; + box-shadow: none !important; + background: #ffffff !important; + color: #000000 !important; box-sizing: border-box !important; } diff --git a/docs/source/_static/overview_pquant_updated.png b/docs/source/_static/overview_pquant_updated.png new file mode 100644 index 0000000..e3a08d3 Binary files /dev/null and b/docs/source/_static/overview_pquant_updated.png differ diff --git a/docs/source/_static/pruning_methods_overview.png b/docs/source/_static/pruning_methods_overview.png new file mode 100644 index 0000000..9539e70 Binary files /dev/null and b/docs/source/_static/pruning_methods_overview.png differ diff --git a/docs/source/conf.py b/docs/source/conf.py index 64ac78d..971c162 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -9,16 +9,14 @@ import os import sys -sys.path.insert(0, os.path.abspath('../')) +sys.path.insert(0, os.path.abspath("../")) -project = 'PQuantML' -copyright = '2025, Roope Niemi' -author = 'Roope Niemi, Anastasiia Petrovych' -release = "1.0.0" +project = "PQuantML" +copyright = "2025, Roope Niemi" +author = "Roope Niemi, Anastasiia Petrovych" +release = "0.0.6" version = release -# -- General configuration --------------------------------------------------- -# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration myst_enable_extensions = [ "amsmath", @@ -37,34 +35,32 @@ autosummary_generate = True -extensions = ['myst_parser', 'sphinx.ext.autodoc', 'sphinx.ext.autosummary', 'sphinx.ext.napoleon', 'sphinx_rtd_theme'] +extensions = ["myst_parser", "sphinx.ext.autodoc", "sphinx.ext.autosummary", "sphinx.ext.napoleon", "sphinx_rtd_theme"] -source_suffix = ['.rst', '.md'] +source_suffix = [".rst", ".md"] -templates_path = ['_templates'] -exclude_patterns = ['_build'] +templates_path = ["_templates"] +exclude_patterns = ["_build"] html_logo = "_static/pquant.png" html_theme_options = { - 'logo_only': True, - 'display_version': True, + "logo_only": True, + "display_version": True, } html_context = { - 'display_github': True, # Integrate GitHub - 'github_user': 'nroope', # Username - 'github_repo': "PQuant", # Repo name - 'github_version': 'master', # Version - 'conf_py_path': '/docs/', # Path in the checkout to the docs root + "display_github": True, # Integrate GitHub + "github_user": "nroope", # Username + "github_repo": "PQuant", # Repo name + "github_version": "master", # Version + "conf_py_path": "/docs/", # Path in the checkout to the docs root } -# -- Options for HTML output ------------------------------------------------- -# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output html_theme = "sphinx_rtd_theme" -html_static_path = ['_static'] -html_favicon = '_static/pquant.png' +html_static_path = ["_static"] +html_favicon = "_static/pquant.png" html_css_files = [ - 'custom.css', -] \ No newline at end of file + "custom.css", +] diff --git a/docs/source/faq.md b/docs/source/faq.md index 5a3f9f1..5af24cf 100644 --- a/docs/source/faq.md +++ b/docs/source/faq.md @@ -1,28 +1,38 @@ # FAQs ## What models formats does PQuantML currently support? -PQuantML primarily supports PyTorch and TensorFlow/Keras models and supports both direct construction and automatic layer replacement using `add_compression_layers(...)`. +PQuantML primarily supports PyTorch and TensorFlow/Keras models and supports both direct construction and automatic layer replacement using `add_compression_layers(...)` method. -## What are requirements to use PQuantML? -Install PyTorch with the correct CUDA version that matches your system and other frameworks, like TensorFlow. This prevents version mismatches and GPU compatibility issues. + +## What are the requirements for using PQuantML? +PQuantML supports two backends. If you are using the PyTorch backend, make sure to install a version of PyTorch built for the CUDA version installed on your system. If you are also using frameworks such as TensorFlow, ensure that all frameworks are compatible with the same CUDA version to avoid version conflicts and GPU compatibility issues. An example to install PyTorch with CUDA 13.0: ```python pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu130 ``` + +## Can I export models to ONNX? + +Yes. PQuantML supports exporting compatible models to the **ONNX** format, making it easy to deploy quantized models across different inference runtimes and hardware platforms. + + ## Can I use MLflow locally? -Yes. +Yes. PQuantML integrates with MLflow for experiment tracking and model logging and local usage is fully supported. -### Start local MLFlow UI: +### Start local MLFlow UI: ```python mlflow ui --host 0.0.0.0 --port 5000 ``` +By default, the MLflow UI is available at `http://localhost:5000`. + +### Use a local or remote Optuna database: +PQuantML also supports storing Optuna studies in either a local SQLite database or a remote database server. -### Use a local or remote database for Optuna tuning: ```python from pquant.core.finetuning import TuningTask tuner = TuningTask(config) diff --git a/docs/source/getting_started.md b/docs/source/getting_started.md index 63b5f0b..0ff3ba1 100644 --- a/docs/source/getting_started.md +++ b/docs/source/getting_started.md @@ -1,19 +1,20 @@ # Quick User Guide ```{note} -This section provides an overview of how to use the PQuantML library: defining models with pruning and quantization, running fine-tuning, and optionally converting the final model to hls4ml. +This section provides an overview of how to use the PQuantML library: defining models with pruning and quantization, running hyperparameters optimization, and optionally converting the final model to hls4ml. ``` -## Model definition & training - -To compress a model with PQuantML, all layers must be replaced with their PQuantML equivalents. For example, replace `Dense` by `PQDense`, `ReLU` by `PQActivation`, etc. +## Model definition & training +To enable pruning and quantization, a model must use PQuantML layers. This can be done in one of two ways: +- Direct layer definition, by building the model with PQuantML layers such as PQDense and PQActivation. +- Automatic layer replacement, by converting an existing PyTorch model using add_compression_layers(...). -Model compression behaviour such as pruning strength, quantization bit-widths, training parameters, etc. is controlled through the configuration object, which is a Pydantic model. +Model compression behaviour such as pruning strength, quantization bit-widths, training parameters, etc. is controlled through the configuration object, which is a Pydantic model to provide an automatic type checking. -### Load a default configuration -``` python +### Load the default DST configuration +```python from pquant import dst_config # Upload a default DST config @@ -39,13 +40,13 @@ def build_model(config): class Model(torch.nn.Module): def __init__(self): super().__init__() - self.dense1 = PQDense(config, 16, 64, + self.dense1 = PQDense(config, 16, 64, in_quant_bits = (1, 3, 3)) self.relu = PQActivation(config, "relu") self.dense2 = PQDense(config, 64, 32) self.dense3 = PQDense(config, 32, 32) - self.dense4 = PQDense(config, 32, 5, - quantize_output=True, + self.dense4 = PQDense(config, 32, 5, + quantize_output=True, out_quant_bits=(1, 3, 3)) def forward(self, x): @@ -57,7 +58,7 @@ def build_model(config): return Model(config) ``` - +This approach is recommended when developing a new architecture from scratch. ### Layer-replacement usage ```python @@ -78,7 +79,7 @@ def build_model(): x = self.relu(self.dense3(x)) x = self.dense4(x) return x - + return Model() @@ -86,10 +87,12 @@ def build_model(): model = add_compression_layers(model, config) ``` -### Fine-Tuning with PQuantML -PQuantML provides an automated fine-tuning and hyperparameter-optimization workflow through the `TuningTask API`. This allows you to search for optimal pruning, quantization, and training parameters using your own training, validation, and objective functions. +If you already have a model, it can be converted automatically by replacing supported layers with their PQuantML equivalents. + +### Hyperparameters optimization with PQuantML +PQuantML provides an automated hyperparameter-optimization workflow through the TuningTask API. This allows you to search for optimal pruning, quantization, and training parameters using your own training, validation, and objective functions. -```python +```python from pquant.core.finetuning import TuningTask, TuningConfig # Convert defined yaml file into the object @@ -114,7 +117,7 @@ tuner.set_optimizer_function(get_optimizer) tuner.set_scheduler_function(get_scheduler) ``` -To run optimization: +Run optimization: ```python device = "cuda" if torch.cuda.is_available() else "cpu" model = model.to(device) @@ -124,10 +127,10 @@ best_params = tuner.run_optimization(model, testloader=..., loss_func=...) ``` + ```{note} -`tuner.run_optimization()` automatically runs multiple compression–fine-tuning cycles, evaluates each trial using your objective function, and returns the best hyperparameters. +`tuner.run_optimization()` automatically runs multiple compression cycles, evaluates each trial using your objective function, and returns the best hyperparameter configuration. ``` - All other training code remains unchanged. ### Train a model @@ -142,37 +145,38 @@ Training is handled through the `train_model(...)` wrapper: ```python from pquant import train_model -trained_model = train_model(model = model, - config = config, - train_func = ..., - valid_func = ..., - trainloader = ..., +trained_model = train_model(model = model, + config = config, + train_func = ..., + valid_func = ..., + trainloader = ..., device="cuda", - testloader = ..., + testloader = ..., loss_func = loss_func, optimizer = optimizer, scheduler=scheduler ) -``` +``` ### Using different quantization settings per layer ```{note} -For FITCompress, HGQ, or architectures, where activations require different quantization bit-widths, each activation layer must be instantiated separately. +If different activation layers require different quantization settings (for example when using FITCompress or HGQ), instantiate each `PQActivation` layer separately instead of reusing a single activation module. ``` + ```python def build_model(config): class Model(torch.nn.Module): def __init__(self): super().__init__() - self.dense1 = PQDense(config, 16, 64, + self.dense1 = PQDense(config, 16, 64, in_quant_bits = (1, 3, 3)) self.relu1 = PQActivation(config, "relu") self.relu2 = PQActivation(config, "relu") self.relu3 = PQActivation(config, "relu") self.dense2 = PQDense(config, 64, 32) self.dense3 = PQDense(config, 32, 32) - self.dense4 = PQDense(config, 32, 5, - quantize_output=True, + self.dense4 = PQDense(config, 32, 5, + quantize_output=True, out_quant_bits=(1, 3, 3)) def forward(self, x): diff --git a/docs/source/index.rst b/docs/source/index.rst index 344ba5d..5348005 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,7 +1,7 @@ .. PQuantMLdocumentation master file, created by sphinx-quickstart on Mon Dec 8 16:28:11 2023. You can adapt this file completely to your liking, but it should at least - contain the root `toctree` directive. + contain the root toctree directive. =========================== PQuantML @@ -14,19 +14,23 @@ PQuantML .. image:: https://badge.fury.io/py/hgq.svg :target: https://pypi.org/project/pquant-ml/ -Welcome to the official documentation for **PQuantML**, a hardware-aware model compression framework supporting: +Welcome to the documentation for PQuantML, a hardware-aware model compression framework supporting: -- Joint pruning + quantization -- Layer-wise precision configuration -- Flexible training pipelines -- PyTorch and TensorFlow backends +- Joint pruning + quantization +- Layer-wise precision configuration +- Flexible training pipelines +- PyTorch and TensorFlow backends +- Knowledge distillation +- HGQ library intergration - Integration with hardware-friendly toolchains (e.g., hls4ml) PQuantML enables efficient deployment of compact neural networks on resource-constrained hardware such as FPGAs and embedded accelerators. +The paper describing the framework is available at: `PQuantML: A Tool for End-to-End Hardware-aware Model Compression `_. + .. rst-class:: light -.. image:: _static/overview_pquant.png +.. image:: _static/overview_pquant_updated.png :alt: PQuantML-overview :width: 100% :align: center @@ -37,19 +41,19 @@ PQuantML enables efficient deployment of compact neural networks on resource-con Key Features ------------ -- **Joint Quantization + Pruning:** Combine bit-width reduction with structured pruning. -- **Flexible Precision Control:** Per-layer and mixed-precision configuration. -- **Hardware-Aware Objective:** Include resource constraints (DSP, LUT, BRAM) in training. -- **Simple API:** Configure compression through a single YAML or Python object. -- **PyTorch Integration:** Works with custom training/validation loops. -- **Export Support:** Model conversion towards hardware toolchains. +- Joint Quantization + Pruning: Combine bit-width reduction with structured pruning. +- Flexible Precision Control: Per-layer and mixed-precision configuration. +- Hardware-Aware Objective: Include resource constraints (DSP, LUT, BRAM) in training. +- Simple API: Configure compression through a single YAML or Python object. +- PyTorch Integration: Works with custom training/validation loops. +- Export Support: Model conversion towards hardware toolchains. Contents ========================= .. toctree:: :maxdepth: 2 - + status install getting_started diff --git a/docs/source/install.md b/docs/source/install.md index ef039ab..8fed3e6 100644 --- a/docs/source/install.md +++ b/docs/source/install.md @@ -1,8 +1,10 @@ # Installation +Installation via pip: `pip install pquant-ml`. -Use `pip install pquant-ml` to install the latest version from PyPI. You will need an environment with `python>=3.10,<=3.12` installed. +With TensorFlow backend: `pip install pquant-ml[tensorflow]`. +With PyTorch backend: `pip install pquant-ml[torch]`. ```{warning} -PQuantML v1.0 requires `tensorflow>=2.17`, `mlflow>=2.0,<3.0`, and `python>=3.10,<=3.12`. +PQuantML v0.0.6 requires `tensorflow>=2.17`, `mlflow>=2.0,<3.0`, and `python>=3.10,<=3.12`. ``` diff --git a/docs/source/reference.md b/docs/source/reference.md index 10802c7..53aafef 100644 --- a/docs/source/reference.md +++ b/docs/source/reference.md @@ -2,19 +2,20 @@ ## Config file -The most important part of the library is a user-defined config yaml file. It has five separate sections: training, pruning, quantization, finetuning, and fitcompress section, `currently maintained by TensorFlow only`, parameters. By default, the parameters in the config are the following: +The most important part of the library is a user-defined `config.yaml` file. It has five separate sections: **training, pruning, quantization, hpo, and fitcompress**. By default, the parameters in the config are the following: ### Training parameters -The following table outlines the primary parameters used to configure the training process: +The following table outlines the description of default parameters used to configure the training process: | **Field** | **Type** | **Default** | **Description** | |------------------------|---------------------------------------------|---------------|------------------------------------------------------------| -| `epochs` | int | `200` | Total number of training epochs. | -| `fine_tuning_epochs` | int | `0` | Additional epochs for fine-tuning. | -| `pretraining_epochs` | int | `50` | Pretraining / warm-up epochs. | -| `rewind` | str | `"never"` | Weight rewinding policy. | -| `rounds` | int | `1` | Number of prune–fine-tune cycles. | -| `save_weights_epoch` | int | `-1` | Save checkpoint at this epoch (`-1` disables). | +| `epochs` | int | `200` | Number of epochs during the main training stage. | +| `fine_tuning_epochs` | int | `0` | Number of epochs during the fine-tuning stage. | +| `pretraining_epochs` | int | `50` | Number of epochs during the pretraining stage. | +| `rewind` | str | `"never"` | When to rewind the weights: never, every-round, or posttraining-stage. | +| `rounds` | int | `1` | Number of pruning/quantization rounds during training. | +| `save_weights_epoch` | int | `-1` | Epoch at which weights are saved for rewinding during the first round. | +| `pruning_first` | bool | `False` | Whether to prune before quantization. If false, pruning occurs after quantization.| ```{note} If you require additional parameters for the training or optimization loops, please define them directly in the config.yaml file. @@ -33,10 +34,11 @@ If you require additional parameters for the training or optimization loops, ple | `quantize_input` | bool | `true` | Whether inputs to layers are quantized by default. | | `quantize_output` | bool | `true` | Whether outputs of layers are quantized by default. | | `enable_quantization` | bool | `true` | Global switch to enable or disable quantization. | +| `granularity` | str | `"per_tensor"` | Whether bitwidths are shared across the whole tensor, per-channel, or per-weight. | | `hgq_gamma` | float | `0.0` | HGQ regularization coefficient for bitwidth stability. | | `hgq_beta` | float | `0.0` | HGQ loss coefficient scaling EBOPs. | | `layer_specific` | dict | `{}` | Dictionary for per-layer quantization overrides. | -| `use_hgq` | bool | `false` | Enable or disable High Granularity Quantization (HGQ). | +| `use_high_granularity_quantization` | bool | `false` | Enable or disable High Granularity Quantization (HGQ). | | `use_real_tanh` | bool | `false` | Use a real `tanh` instead of hard/approximate `tanh`. | | `overflow_mode_data` | str | `"SAT"` | Overflow handling mode for input and output quantizers(`SAT`, `SAT_SYM`, `WRAP`, `WRAP_SM`). | | `overflow_mode_parameters` | str | `"SAT"` | Overflow handling mode for weight and biases quantizers(`SAT`, `SAT_SYM`, `WRAP`, `WRAP_SM`). | @@ -44,7 +46,7 @@ If you require additional parameters for the training or optimization loops, ple | `use_relu_multiplier` | bool | `true` | Enable a learned bit-shift multiplier inside ReLU layers. | -### Fine-tuning parameters +### Hyperparameters optimization parameters | **Field** | **Type** | **Default** | **Description** | |-------------------------|--------------------------|--------------------|-----------------------------------------------| @@ -163,24 +165,24 @@ There are more details about every pruning method: | `t_start_collecting_batch` | int | `50` | Steps to skip before collecting statistics. | #### MDMM Pruning - -| **Field** | **Type** | **Default** | **Description** | -|--------------------|-----------------------|----------------------------|--------------------------------------------------------------| -| `pruning_method` | str | `mdmm` | Selects this pruning schema. | -| `constraint_type` | ConstraintType | `"Equality"` | Constraint form: equality / ≤ / ≥. | -| `target_value` | float | `0.0` | Target value for the chosen metric. | -| `metric_type` | MetricType | `"UnstructuredSparsity"` | Specifies which metric is constrained. | -| `target_sparsity` | float | `0.9` | Target sparsity when constraining sparsity. | -| `rf` | int | `1` | Regularization / frequency parameter. | -| `epsilon` | float | `1.0e-03` | Feasibility tolerance. | -| `scale` | float | `10.0` | Penalty scaling for constraint violation. | -| `damping` | float | `1.0` | Damping term for numerical stability. | -| `use_grad` | bool | `false` | Use gradient information during updates. | -| `l0_mode` | `"coarse"` \| `"smooth"` | `"coarse"` | L0 approximation mode. | -| `scale_mode` | `"mean"` \| `"sum"` | `"mean"` | Aggregation mode for penalties. | - - -Optionally, there is also FITCompress method implemented for PyTorch: +| **Field** | **Type** | **Default** | **Description** | +|--------------------|-----------------------|----------------------------|---------------------------------------------------------------------------------| +| `pruning_method` | str | `"mdmm"` | Selects this pruning schema. | +| `constraint_type` | ConstraintType | `"Equality"` | Constraint form: `Equality`, `LessThanOrEqual`, or `GreaterThanOrEqual`. | +| `target_value` | float | `0.0` | Target value for the chosen metric (used for the general constraint). | +| `metric_type` | MetricType | `"UnstructuredSparsity"` | Metric being constrained: `UnstructuredSparsity` or `StructuredSparsity`. | +| `target_sparsity` | float | `0.9` | Target sparsity when constraining sparsity. | +| `rf` | int | `1` | Regularization / frequency parameter. | +| `epsilon` | float | `1.0e-03` | Feasibility tolerance. | +| `scale` | float | `10.0` | Penalty scaling for constraint violation. | +| `damping` | float | `1.0` | Damping term for numerical stability. | +| `use_grad` | bool | `false` | Use gradient information during updates. | +| `l0_mode` | `"coarse"` \| `"smooth"` | `"coarse"` | L0 approximation mode. | +| `scale_mode` | `"mean"` \| `"sum"` | `"mean"` | Aggregation mode for penalties. | +| `constraint_lr` | float | `1.0e-03` | Learning rate for the Lagrange multiplier (dual variable). | + + +Optionally, there is also FITCompress method implemented for PyTorch-only: ### FitCompress method | **Field** | **Type** | **Default** | **Description** | |---------------------------|----------|-------------|---------------------------------------------------------------------------------| @@ -201,9 +203,19 @@ Optionally, there is also FITCompress method implemented for PyTorch: - `PQAvgPool*D`: Average pooling layers. - `PQBatchNorm*D`: BatchNorm layers. - `PQDense`: Linear layer. -- `PQActivation`: Activation layers (ReLU, Tanh) +- `PQActivation`: Activation layers (ReLU, Tanh, Leaky Relu, Gelu, Hard Tanh, or a user-provided activation function (Torch only) ). +- `MultiHeadAttention`: Multi-head attention layer. +- `LayerNorm`: Layer normalization layer (Currently Torch only). ```{note} -Currently, PQuantML supports two quantization modes: layer-wise fixed-point quantization, where each tensor uses a single -bit-width configuration, and High-Granularity Quantization (HGQ). +PQuantML supports two quantization modes, each with several granularity options: + +**Fixed-point quantization** (for weights): +- per-weight +- per-channel +- per-tensor + +**HGQ (High Granularity Quantization)**: +- per-weight (learned bit-widths per weight) +- per-tensor (learned bit-widths per tensor) ``` diff --git a/docs/source/status.md b/docs/source/status.md index 15f1aa9..276021b 100644 --- a/docs/source/status.md +++ b/docs/source/status.md @@ -1,17 +1,23 @@ # PQuantML Status -This page tracks the development status of PQuantML features +This page tracks the development status of PQuantML features. -## Release: v1.0.0 +## Release: v0.0.6 | Feature | Status | Notes | |---------------------------------|-----------------|-------| -| Compression pipeline | ✅ Complete | Included in v1.0.0 | +| Compression pipeline | ✅ Complete | Included in v0.0.6 | | Pruning methods (7 variants) | ✅ Complete | All documented | | Quantization (fixed-point) | ✅ Complete | Supports per-layer overrides | | HGQ support |✅ Complete | Supports HGQ quantization | -| hls4ml integration | ✅ Complete | Works in v1.0.0 | -| FITCompress | 🚧 Partially implemented | Works through PyTorch only | -| Documentation | 🚧 Improving | Expanded daily | - - +| hls4ml integration | ✅ Complete | Works in v0.0.6 | +| FITCompress | ✅ Complete | Supported in PyTorch only | +| Model fit support | ✅ Complete | Works in v0.0.6 | +| Alkaid converter support | ⏳ Coming in v0.0.7 | Implemented in dev | +| Onnx converter support | ⏳ Coming in v0.0.7 | Implemented in dev | +| Knowledge distillation | ⏳ Coming in v0.0.7 | Implemented in dev | +| Implementation of HGQ and pruning layers in Torch | ⏳ Coming in v0.0.7 | Implemented in dev | +| Additional test coverage for pruning methods | ⏳ Will be in the next release | Implemented in dev | +| CI/CD pipeline | 🚧 Work in progress | Due to the next release | +| MDMM pruning algorithm metrics extension | 🚧 Work in progress | Due to the next release | +| Documentation | 🚧 Improving | Expanded monthly | diff --git a/pyproject.toml b/pyproject.toml index c6e0c9e..58af8da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,5 @@ [build-system] build-backend = "setuptools.build_meta" - requires = [ "setuptools>=46.1", "setuptools-scm[toml]>=5" ] [project] @@ -28,12 +27,19 @@ optional-dependencies.all = [ "pytest>=8.4", "tensorflow>=2.17,<=2.20", "torch>= optional-dependencies.tensorflow = [ "tensorflow>=2.17,<=2.20" ] optional-dependencies.test = [ "pytest>=8.4" ] optional-dependencies.torch = [ "torch>=2.1" ] -urls.repository = "https://github.com/nroope/PQuantML" +urls.repository = "https://github.com/cern-nextgen/PQuantML" +entry-points.alkaid_keras.pquant = "pquant._alkaid_plugin._alkaid_keras_plugin:register" +entry-points.alkaid_torch.pquant = "pquant._alkaid_plugin._alkaid_torch_plugin:register" [tool.setuptools] packages = [ "pquant" ] -include-package-data = true package-dir = { "" = "src" } +include-package-data = true [tool.setuptools_scm] write_to = "src/pquant/_version.py" + +[tool.ruff] +line-length = 125 +src = [ "src" ] +lint.extend-select = [ "I" ] diff --git a/src/pquant/__init__.py b/src/pquant/__init__.py index 299fa45..ecfd1f8 100644 --- a/src/pquant/__init__.py +++ b/src/pquant/__init__.py @@ -5,7 +5,7 @@ # flake8: noqa backend = os.getenv("KERAS_BACKEND", "tensorflow") if backend == "torch": - from . import configs, pruning_methods + from . import configs from .core.hyperparameter_optimization import ( PQConfig, ap_config, @@ -19,19 +19,26 @@ pdp_config, wanda_config, ) - from .core.torch import activations, layers, optimizers, quantizer + from .core.torch import ( + activations, + layers, + optimizers, + pruning_methods, + quantizer, + tracing, + ) from .core.torch.layers import ( add_compression_layers, apply_final_compression, get_ebops, get_layer_keep_ratio, get_model_losses, - load_torch_hgq_model, post_training_prune, ) + from .core.torch.tracing import check_quantization, print_quantization_check from .core.torch.train import train_model - _forwards = ["activations", "layers", "quantizer", "optimizers"] + _forwards = ["activations", "layers", "quantizer", "optimizers", "tracing"] for name in _forwards: mod = importlib.import_module(f".core.torch.{name}", package="pquant") @@ -56,12 +63,13 @@ _forwards.append("load_from_file") _forwards.append("load_from_dictionary") _forwards.append("get_ebops") - _forwards.append("load_torch_hgq_model") + _forwards.append("check_quantization") + _forwards.append("print_quantization_check") _forwards.append("PQConfig") __all__ = _forwards else: - from . import configs, pruning_methods + from . import configs from .core.hyperparameter_optimization import ( PQConfig, ap_config, @@ -74,7 +82,7 @@ pdp_config, wanda_config, ) - from .core.keras import activations, layers, quantizer + from .core.keras import activations, layers, pruning_methods, quantizer from .core.keras.layers import ( add_compression_layers, apply_final_compression, diff --git a/src/pquant/pruning_methods/__init__.py b/src/pquant/_alkaid_plugin/__init__.py similarity index 100% rename from src/pquant/pruning_methods/__init__.py rename to src/pquant/_alkaid_plugin/__init__.py diff --git a/src/pquant/_alkaid_plugin/_alkaid_common.py b/src/pquant/_alkaid_plugin/_alkaid_common.py new file mode 100644 index 0000000..9280b8a --- /dev/null +++ b/src/pquant/_alkaid_plugin/_alkaid_common.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +from alkaid.trace.ops import quantize as alkaid_quantize + + +class PQuantAlkaidError(ValueError): + """Raised for PQuant states that cannot be replayed by Alkaid.""" + + +def _to_numpy(value: Any) -> np.ndarray: + if value is None: + return np.array(0.0) + if isinstance(value, np.ndarray): + return value + if hasattr(value, "detach"): + value = value.detach() + if hasattr(value, "cpu"): + value = value.cpu() + return value.numpy() + try: + import keras + + return np.asarray(keras.ops.convert_to_numpy(value)) + except Exception: + return np.asarray(value) + + +def _to_bool(value: Any, default: bool = False) -> bool: + if value is None: + return default + try: + arr = _to_numpy(value) + except Exception: + return bool(value) + if arr.shape == (): + return bool(arr.item()) + return bool(np.all(arr)) + + +def _to_int_bits(value: Any) -> np.ndarray: + return np.rint(_to_numpy(value)).astype(np.int64) + + +def _raw_module_attr(obj: Any, name: str, default: Any = None) -> Any: + for storage_name in ("_parameters", "_buffers", "_modules"): + storage = getattr(obj, storage_name, None) + if isinstance(storage, dict) and name in storage: + return storage[name] + try: + return object.__getattribute__(obj, name) + except AttributeError: + return getattr(obj, name, default) + + +def _quantizer_kif(quantizer: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + if hasattr(quantizer, "_parameters"): + if not bool(_raw_module_attr(quantizer, "use_hgq", False)): + return ( + _to_int_bits(_raw_module_attr(quantizer, "k")), + _to_int_bits(_raw_module_attr(quantizer, "i")), + _to_int_bits(_raw_module_attr(quantizer, "f")), + ) + inner = _raw_module_attr(quantizer, "quantizer") + if hasattr(inner, "_parameters") or hasattr(inner, "_buffers"): + k = _raw_module_attr(inner, "_k") + i = _raw_module_attr(inner, "_i_raw", None) + if i is None: + i = _raw_module_attr(inner, "_i") + f = _raw_module_attr(inner, "_f") + return _to_int_bits(k), _to_int_bits(i), _to_int_bits(f) + k, i, f = quantizer.get_quantization_bits() + return _to_int_bits(k), _to_int_bits(i), _to_int_bits(f) + + +def _replay_quantizer(quantizer: Any, x: Any) -> Any: + k, i, f = _quantizer_kif(quantizer) + inner = _raw_module_attr(quantizer, "quantizer", None) + overflow = _raw_module_attr(quantizer, "overflow", _raw_module_attr(inner, "overflow_mode", "WRAP")) + round_mode = _raw_module_attr(quantizer, "round_mode", _raw_module_attr(inner, "round_mode", "TRN")) + return alkaid_quantize(x, k=k, i=i, f=f, overflow_mode=str(overflow).upper(), round_mode=str(round_mode).upper()) + + +def _replay_quantizer_if_enabled(layer: Any, quantizer_name: str, x: Any, flag_name: str) -> Any: + if not bool(getattr(layer, "enable_quantization", True)): + return x + if not bool(getattr(layer, flag_name, True)): + return x + quantizer = getattr(layer, quantizer_name, None) + return _replay_quantizer(quantizer, x) + + +def _assert_final_compression(layer: Any) -> None: + if not _to_bool(_raw_module_attr(layer, "final_compression_done", False)): + raise PQuantAlkaidError( + f"{type(layer).__name__} must have apply_final_compression() applied before Alkaid conversion." + ) + + +def _final_bias(layer: Any) -> np.ndarray: + """The layer's final (compressed) bias as numpy, or a scalar zero when absent.""" + _assert_final_compression(layer) + bias = getattr(layer, "_bias", None) + if bias is None: + return np.array(0.0) + return _to_numpy(bias) + + +def _scale_by_relu_multiplier(layer: Any, x: Any) -> Any: + """Apply PQActivation's power-of-two ReLU multiplier, which is used only without HGQ.""" + applies = ( + not _to_bool(getattr(layer, "use_hgq", False)) + and _to_bool(getattr(layer, "use_multiplier", False)) + and layer.activation_name == "relu" + and hasattr(layer, "multiplier") + ) + if not applies: + return x + return x * (2.0 ** np.rint(_to_numpy(layer.multiplier))) + + +def _replay_table(table: Any, x: Any, table_fn: Any) -> Any: + """Replay a lookup-table activation: quantize input, apply the table, quantize output.""" + if not (table.quantize_output and table.enable_quantization): + name = getattr(table, "name", None) or type(table).__name__ + raise PQuantAlkaidError(f"PQSoftmax table {name!r} must have an enabled output quantizer for Alkaid conversion.") + x = _replay_quantizer_if_enabled(table, "input_quantizer", x, "quantize_input") + out = x.apply(table_fn(table)) + return _replay_quantizer(table.output_quantizer, out) + + +def _replay_softmax(layer: Any, inputs: Any, table_fn: Any) -> Any: + """Replay PQSoftmax through its exp and inverse-sum lookup tables.""" + inputs = _replay_quantizer_if_enabled(layer, "input_quantizer", inputs, "quantize_input") + if layer.stable: + inputs = np.max(inputs, axis=layer.axes, keepdims=True) - inputs # type: ignore + exponents = _replay_table(layer.exp_table, inputs, table_fn) + sums = np.sum(exponents, axis=layer.axes, keepdims=True) + inverse_sums = _replay_table(layer.inv_table, sums, table_fn) + out = exponents * inverse_sums + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +def _mark_plugin_loaded(framework: str) -> None: + """Record the pquant plugin as loaded in Alkaid's plugin loader, if that loader exists.""" + try: + from alkaid.converter import _plugin_loader + + _plugin_loader._LOADED.add(("pquant", framework)) + except Exception: + pass diff --git a/src/pquant/_alkaid_plugin/_alkaid_keras_plugin.py b/src/pquant/_alkaid_plugin/_alkaid_keras_plugin.py new file mode 100644 index 0000000..f848222 --- /dev/null +++ b/src/pquant/_alkaid_plugin/_alkaid_keras_plugin.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +from math import prod + +import keras +import numpy as np +from alkaid.converter.builtin.keras.layers._base import ReplayOperationBase +from alkaid.converter.builtin.keras.layers.activation import keras_numpy_unary_map +from alkaid.converter.builtin.keras.layers.batchnorm import ReplayBatchNormalization +from alkaid.converter.builtin.keras.layers.conv import _conv +from alkaid.converter.builtin.keras.layers.pool import ReplayPool +from alkaid.trace import FVArray +from alkaid.trace.ops import einsum, extract_patches +from keras.layers import DepthwiseConv1D, DepthwiseConv2D + +from pquant._alkaid_plugin._alkaid_common import ( + PQuantAlkaidError, + _assert_final_compression, + _final_bias, + _mark_plugin_loaded, + _replay_quantizer, + _replay_quantizer_if_enabled, + _replay_softmax, + _scale_by_relu_multiplier, + _to_numpy, +) +from pquant.core.keras.activations import PQActivation +from pquant.core.keras.layers import ( + PQAvgPool1d, + PQAvgPool2d, + PQBatchNormalization, + PQConv1d, + PQConv2d, + PQDense, + PQDepthwiseConv2d, + PQMultiheadAttention, + PQSeparableConv2d, + PQSoftmax, +) +from pquant.core.keras.quantizer import Quantizer + + +def _final_kernel(layer) -> np.ndarray: + """The layer's final (compressed) kernel as numpy. Asserts compression was applied.""" + _assert_final_compression(layer) + return _to_numpy(layer._kernel) + + +def _table_fn(table): + """Numpy-callable for a PQActivation lookup table, evaluated in float32 like the keras runtime.""" + fn = table.activation_function + + def apply_fn(v: np.ndarray) -> np.ndarray: + t = keras.ops.cast(keras.ops.convert_to_tensor(v), "float32") + return np.asarray(keras.ops.convert_to_numpy(fn(t)), dtype=np.float64) + + return apply_fn + + +def _unpack_query_key_value(inputs): + if not isinstance(inputs, (list, tuple)): + return inputs, inputs, inputs + if len(inputs) == 3: + return inputs[0], inputs[1], inputs[2] + if len(inputs) == 2: + return inputs[0], inputs[1], inputs[1] + return inputs[0], inputs[0], inputs[0] + + +class ReplayPQuantQuantizer(ReplayOperationBase): + __activation_handled__ = True + handles = (Quantizer,) + + def call(self, x: FVArray) -> FVArray: + return _replay_quantizer(self.op, x) + + +class ReplayPQuantDense(ReplayOperationBase): + handles = (PQDense,) + + def call(self, inputs: FVArray) -> FVArray: + layer = self.op + inputs = _replay_quantizer_if_enabled(layer, "input_quantizer", inputs, "quantize_input") + out = np.einsum("...c,cC->...C", inputs, _final_kernel(layer)) + _final_bias(layer) + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +class ReplayPQuantConv(ReplayOperationBase): + handles = (PQConv1d, PQConv2d, PQDepthwiseConv2d) + + def call(self, inputs: FVArray) -> FVArray: + layer = self.op + inputs = _replay_quantizer_if_enabled(layer, "input_quantizer", inputs, "quantize_input") + kernel = _final_kernel(layer) + bias = _final_bias(layer) + + if isinstance(layer, (DepthwiseConv1D, DepthwiseConv2D)): + ch_in, dm = kernel.shape[-2:] + kernel = kernel.reshape(*kernel.shape[:-2], 1, ch_in * dm) + groups = ch_in + else: + groups = layer.groups + + x = extract_patches( + inputs, + size=layer.kernel_size, + strides=layer.strides, + dilation_rate=layer.dilation_rate, + padding=layer.padding, + data_format=layer.data_format, + ) + ch_out = kernel.shape[-1] + ch_in_per_g = kernel.shape[-2] + k_vol = int(prod(layer.kernel_size)) + out = _conv( + x, + kernel, + k_vol=k_vol, + groups=groups, + ch_in_per_g=ch_in_per_g, + out_per_g=ch_out // groups, + ) + if bias.shape != (): + out = out + bias + if layer.data_format == "channels_first": + out = np.moveaxis(out, -1, 1) # type: ignore + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +class ReplayPQuantSeparableConv(ReplayOperationBase): + handles = (PQSeparableConv2d,) + + def call(self, inputs: FVArray) -> FVArray: + layer = self.op + x = ReplayPQuantConv(layer.depthwise_conv).call(inputs) + return ReplayPQuantConv(layer.pointwise_conv).call(x) + + +class ReplayPQuantBatchNormalization(ReplayBatchNormalization): + handles = (PQBatchNormalization,) + + def fused_scale_offset(self) -> tuple[np.ndarray, np.ndarray]: + layer = self.op + _assert_final_compression(layer) + mean = _to_numpy(keras.ops.cast(layer.moving_mean, layer.dtype)) + variance = _to_numpy(keras.ops.cast(layer.moving_variance, layer.dtype)) + if layer.scale: + gamma = _to_numpy(keras.ops.cast(layer.gamma, layer.dtype)) + else: + gamma = np.ones_like(mean) + if layer.center: + beta = _to_numpy(keras.ops.cast(layer.beta, layer.dtype)) + else: + beta = np.zeros_like(mean) + scale = gamma / np.sqrt(variance + layer.epsilon) + offset = beta - mean * scale + return scale, offset + + def call(self, inputs: FVArray, mask=None) -> FVArray: + layer = self.op + inputs = _replay_quantizer_if_enabled(layer, "input_quantizer", inputs, "quantize_input") + scale, offset = self.fused_scale_offset() + shape = [1] * inputs.ndim + axis = layer.axis if isinstance(layer.axis, (list, tuple)) else [layer.axis] + for a in axis: + aa = a if a >= 0 else inputs.ndim + a + shape[aa] = inputs.shape[aa] + out = inputs + if not np.all(scale == 1): + out = out * scale.reshape(shape) # type: ignore + if not np.all(offset == 0): + out = out + offset.reshape(shape) # type: ignore + return out + + +class ReplayPQuantAvgPool(ReplayPool): + __activation_handled__ = True + handles = (PQAvgPool1d, PQAvgPool2d) + + def call(self, inputs: FVArray, mask: None = None) -> FVArray: + layer = self.op + inputs = _replay_quantizer_if_enabled(layer, "input_quantizer", inputs, "quantize_input") + out = super().call(inputs, mask=mask) + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +class ReplayPQuantActivation(ReplayOperationBase): + __activation_handled__ = True + handles = (PQActivation,) + + def call(self, inputs: FVArray) -> FVArray: + layer = self.op + inputs = _scale_by_relu_multiplier(layer, inputs) + inputs = _replay_quantizer_if_enabled(layer, "input_quantizer", inputs, "quantize_input") + if layer.activation_name not in keras_numpy_unary_map: + raise PQuantAlkaidError(f"Unsupported PQuant activation for Alkaid conversion: {layer.activation_name!r}") + out = keras_numpy_unary_map[layer.activation_name](inputs) + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +class ReplayPQuantSoftmax(ReplayOperationBase): + __activation_handled__ = True + handles = (PQSoftmax,) + + def call(self, inputs: FVArray, mask=None) -> FVArray: + if mask is not None: + raise PQuantAlkaidError("PQSoftmax masks are not supported in Alkaid conversion.") + return _replay_softmax(self.op, inputs, _table_fn) + + +class ReplayPQuantMultiheadAttention(ReplayOperationBase): + __activation_handled__ = True + handles = (PQMultiheadAttention,) + + def call(self, inputs, key_padding_mask=None, attn_mask=None, need_weights=True): + layer = self.op + if key_padding_mask is not None or attn_mask is not None: + raise PQuantAlkaidError("Attention masks are not supported in Alkaid conversion.") + + query, key, value = _unpack_query_key_value(inputs) + batch_size, query_len = query.shape[0], query.shape[1] + key_len = key.shape[1] + num_heads, head_dim = layer.num_heads, layer.head_dim + + q = ReplayPQuantDense(layer.q_proj).call(query) # (B, T, E) + k = ReplayPQuantDense(layer.k_proj).call(key) # (B, S, E) + v = ReplayPQuantDense(layer.v_proj).call(value) # (B, S, E) + + # Reshape to (B, H, T/S, head_dim) + q = q.reshape(batch_size, query_len, num_heads, head_dim).transpose(0, 2, 1, 3) + k = k.reshape(batch_size, key_len, num_heads, head_dim).transpose(0, 2, 1, 3) + v = v.reshape(batch_size, key_len, num_heads, head_dim).transpose(0, 2, 1, 3) + + scale = float(np.float32(layer.scale)) + attn_scores = einsum("bhtd,bhsd->bhts", q, k) * scale + + # The softmax's own input/output quantizers handle the scores and the attention weights + attn_weights = ReplayPQuantSoftmax(layer.softmax).call(attn_scores) + + # Weighted sum of values (dropout is an inference no-op): (B, H, T, head_dim) + out = einsum("bhts,bhsd->bhtd", attn_weights, v) + + # Merge heads: (B, T, E) + out = out.transpose(0, 2, 1, 3).reshape(batch_size, query_len, layer.embed_dim) + out = ReplayPQuantDense(layer.out_proj).call(out) + + if need_weights: + # Average attention weights over heads: (B, T, S) + return out, np.mean(attn_weights, axis=1) + return (out,) + + +def register() -> None: + """Entry point for Alkaid's ``alkaid_keras`` second-level plugin group.""" + _mark_plugin_loaded("keras") diff --git a/src/pquant/_alkaid_plugin/_alkaid_torch_plugin.py b/src/pquant/_alkaid_plugin/_alkaid_torch_plugin.py new file mode 100644 index 0000000..57f11e3 --- /dev/null +++ b/src/pquant/_alkaid_plugin/_alkaid_torch_plugin.py @@ -0,0 +1,259 @@ +from __future__ import annotations + +import builtins +import operator +from functools import wraps +from typing import Any + +import numpy as np +import torch +from alkaid.converter.builtin.torch.layers.direct import torch_numpy_unary_map +from alkaid.converter.builtin.torch.layers.functional import ( + _functional_map, + conv_nd_replay, + replay_avg_pool, +) +from alkaid.converter.builtin.torch.layers.methods import _method_map +from alkaid.converter.builtin.torch.layers.modules import ( + ReplayBatchNorm, + ReplayModuleBase, +) +from alkaid.trace import FVArray + +from pquant._alkaid_plugin._alkaid_common import ( + PQuantAlkaidError, + _assert_final_compression, + _final_bias, + _mark_plugin_loaded, + _replay_quantizer, + _replay_quantizer_if_enabled, + _replay_softmax, + _scale_by_relu_multiplier, + _to_numpy, +) +from pquant.core.torch.activations import PQActivation, PQSoftmax +from pquant.core.torch.layers import ( + PQAvgPool1d, + PQAvgPool2d, + PQBatchNorm1d, + PQBatchNorm2d, + PQConv1d, + PQConv2d, + PQDense, +) +from pquant.core.torch.quantizer import Quantizer + + +def _final_weight(layer: torch.nn.Module) -> np.ndarray: + """The layer's final (compressed) weight as numpy. Asserts compression was applied.""" + _assert_final_compression(layer) + return _to_numpy(layer._weight) + + +def _activation_numpy_fn(layer: PQActivation): + """The numpy elementwise function for a named PQActivation (relu/tanh/gelu/...).""" + name = layer.activation_name + fn = torch_numpy_unary_map.get(name) or torch_numpy_unary_map.get(name.replace("_", "")) + if fn is not None: + return fn + if name == "leaky_relu": + slope = float(getattr(layer.activation_function, "negative_slope", 0.1015625)) + return lambda x: np.where(x < 0, x * slope, x) # type: ignore + raise PQuantAlkaidError(f"Unsupported PQuant activation for Alkaid conversion: {name!r}") + + +def _table_fn(table): + """Numpy-callable for a PQActivation lookup table, evaluated in float32 like the torch runtime.""" + fn = table.activation_function + + def apply_fn(v: np.ndarray) -> np.ndarray: + with torch.no_grad(): + t = torch.as_tensor(v, dtype=torch.float32, device="cpu") + return fn(t).detach().cpu().numpy().astype(np.float64) + + return apply_fn + + +class ReplayPQuantQuantizer(ReplayModuleBase): + handles = (Quantizer,) + + def call(self, input: FVArray) -> FVArray: + return _replay_quantizer(self.module, input) + + +class ReplayPQuantDense(ReplayModuleBase): + handles = (PQDense,) + + def call(self, input: FVArray) -> FVArray: + layer = self.module + input = _replay_quantizer_if_enabled(layer, "input_quantizer", input, "quantize_input") + out = input @ _final_weight(layer).T + bias = _final_bias(layer) + if bias.shape != (): + out = out + bias + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +class ReplayPQuantConv(ReplayModuleBase): + handles = (PQConv1d, PQConv2d) + + def call(self, input: FVArray) -> FVArray: + layer = self.module + input = _replay_quantizer_if_enabled(layer, "input_quantizer", input, "quantize_input") + out = conv_nd_replay( + input, + _final_weight(layer), + _final_bias(layer), + stride=layer.stride, + padding=layer.padding, + dilation=layer.dilation, + groups=layer.groups, + ) + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +class ReplayPQuantBatchNorm(ReplayBatchNorm): + handles = (PQBatchNorm1d, PQBatchNorm2d) + + def fused_scale_offset(self) -> tuple[np.ndarray, np.ndarray]: + layer = self.module + _assert_final_compression(layer) + mean = _to_numpy(layer.running_mean) + variance = _to_numpy(layer.running_var) + gamma = _to_numpy(layer._weight) if layer._weight is not None else np.ones_like(mean) + beta = _to_numpy(layer._bias) if layer._bias is not None else np.zeros_like(mean) + scale = gamma / np.sqrt(variance + layer.eps) + offset = beta - mean * scale + return scale, offset + + def call(self, input: FVArray) -> FVArray: + layer = self.module + input = _replay_quantizer_if_enabled(layer, "input_quantizer", input, "quantize_input") + return super().call(input) + + +class ReplayPQuantAvgPool(ReplayModuleBase): + handles = (PQAvgPool1d, PQAvgPool2d) + + def call(self, input: FVArray) -> FVArray: + layer = self.module + input = _replay_quantizer_if_enabled(layer, "input_quantizer", input, "quantize_input") + out = replay_avg_pool( + input, + layer.kernel_size, + layer.stride, + layer.padding, + layer.ceil_mode, + layer.count_include_pad, + ) + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +class ReplayPQuantActivation(ReplayModuleBase): + handles = (PQActivation,) + + def call(self, input: FVArray) -> FVArray: + layer = self.module + input = _scale_by_relu_multiplier(layer, input) + input = _replay_quantizer_if_enabled(layer, "input_quantizer", input, "quantize_input") + out = _activation_numpy_fn(layer)(input) + return _replay_quantizer_if_enabled(layer, "output_quantizer", out, "quantize_output") + + +class ReplayPQuantSoftmax(ReplayModuleBase): + """Replay PQSoftmax as a single fx-leaf module""" + + handles = (PQSoftmax,) + + def call(self, inputs: FVArray, mask=None) -> FVArray: + module = self.module + if mask is not None: + raise PQuantAlkaidError("PQSoftmax masks are not supported in Alkaid conversion.") + if not module.built: + raise PQuantAlkaidError("PQSoftmax must be built (one real forward) before Alkaid conversion.") + return _replay_softmax(module, inputs, _table_fn) + + +def _patch_root_quantizer_trace() -> None: + """Let a bare Quantizer be traced as the model root, which Alkaid's fx tracer cannot handle.""" + import alkaid.converter.builtin.torch.main as torch_main + + tracer_cls = torch_main.TorchALIRTracer + marker = "__alkaid_pquant_patched_root_quantizer__" + if getattr(tracer_cls, marker, False): + return + original = tracer_cls.apply_model + + @wraps(original) + def wrapped(self, verbose: bool, inputs: tuple[FVArray, ...]): + if isinstance(self.model, Quantizer): + if isinstance(inputs, FVArray): + inputs = (inputs,) + replay = ReplayPQuantQuantizer(self.model) + dump = replay(*inputs) + return {"inputs": tuple(inputs), "quantizer/final": dump["final"], "final": dump["final"]}, ["final"] + return original(self, verbose, inputs) + + tracer_cls.apply_model = wrapped + setattr(tracer_cls, marker, True) + + +def _replay_getattr(obj: Any, name: str, *default: Any) -> Any: + if default: + return getattr(obj, name, default[0]) + return getattr(obj, name) + + +def _replay_tensor(data: Any, *args: Any, **kwargs: Any) -> Any: + if isinstance(data, FVArray): + return data + if isinstance(data, torch.Tensor): + data = data.detach().cpu().numpy() + return np.asarray(data) + + +def _normalize_shape(args: tuple[Any, ...]) -> tuple[int, ...]: + if len(args) == 1 and isinstance(args[0], (tuple, list, torch.Size)): + return tuple(int(v) for v in args[0]) + return tuple(int(v) for v in args) + + +def _replay_zeros(*size: Any, **kwargs: Any) -> np.ndarray: + return np.zeros(_normalize_shape(size), dtype=np.float32) + + +def _replay_ones(*size: Any, **kwargs: Any) -> np.ndarray: + return np.ones(_normalize_shape(size), dtype=np.float32) + + +def _replay_full(size: Any, fill_value: Any, **kwargs: Any) -> np.ndarray: + return np.full(_normalize_shape((size,)), fill_value, dtype=np.float32) + + +def _replay_zeros_like(x: Any, **kwargs: Any) -> np.ndarray: + return np.zeros(tuple(x.shape), dtype=np.float32) + + +def _replay_ones_like(x: Any, **kwargs: Any) -> np.ndarray: + return np.ones(tuple(x.shape), dtype=np.float32) + + +def _register_functional_helpers() -> None: + _functional_map.setdefault(operator.pow, lambda a, b: a**b) + _functional_map.setdefault(torch.pow, lambda a, b: a**b) + _functional_map.setdefault(builtins.getattr, _replay_getattr) + _functional_map.setdefault(torch.tensor, _replay_tensor) + _functional_map.setdefault(torch.as_tensor, _replay_tensor) + _functional_map.setdefault(torch.zeros, _replay_zeros) + _functional_map.setdefault(torch.ones, _replay_ones) + _functional_map.setdefault(torch.full, _replay_full) + _functional_map.setdefault(torch.zeros_like, _replay_zeros_like) + _functional_map.setdefault(torch.ones_like, _replay_ones_like) + _method_map.setdefault("pow", lambda receiver, exponent, **_kwargs: receiver**exponent) + + +def register() -> None: + """Entry point for Alkaid's ``alkaid_torch`` second-level plugin group.""" + _patch_root_quantizer_trace() + _register_functional_helpers() + _mark_plugin_loaded("torch") diff --git a/src/pquant/configs/config_ap.yaml b/src/pquant/configs/config_ap.yaml index 0c3cb22..729ed2d 100644 --- a/src/pquant/configs/config_ap.yaml +++ b/src/pquant/configs/config_ap.yaml @@ -13,6 +13,7 @@ quantization_parameters: default_data_keep_negatives: 0. default_data_integer_bits: 0. default_data_fractional_bits: 7. + dynamic_data_quantization: false granularity: "per_tensor" quantize_input: true quantize_output: false diff --git a/src/pquant/configs/config_autosparse.yaml b/src/pquant/configs/config_autosparse.yaml index e5ce8ed..7ac7ffd 100644 --- a/src/pquant/configs/config_autosparse.yaml +++ b/src/pquant/configs/config_autosparse.yaml @@ -16,6 +16,7 @@ quantization_parameters: default_data_keep_negatives: 0. default_data_integer_bits: 0. default_data_fractional_bits: 7. + dynamic_data_quantization: false granularity: "per_tensor" quantize_input: true quantize_output: false diff --git a/src/pquant/configs/config_cs.yaml b/src/pquant/configs/config_cs.yaml index 88594c8..58bbdad 100644 --- a/src/pquant/configs/config_cs.yaml +++ b/src/pquant/configs/config_cs.yaml @@ -12,6 +12,7 @@ quantization_parameters: default_data_keep_negatives: 0. default_data_integer_bits: 0. default_data_fractional_bits: 7. + dynamic_data_quantization: false granularity: "per_tensor" quantize_input: true quantize_output: false diff --git a/src/pquant/configs/config_dst.yaml b/src/pquant/configs/config_dst.yaml index 9a54073..86b958a 100644 --- a/src/pquant/configs/config_dst.yaml +++ b/src/pquant/configs/config_dst.yaml @@ -14,6 +14,7 @@ quantization_parameters: default_data_keep_negatives: 0. default_data_integer_bits: 0. default_data_fractional_bits: 7. + dynamic_data_quantization: false granularity: "per_tensor" quantize_input: true quantize_output: false diff --git a/src/pquant/configs/config_fitcompress.yaml b/src/pquant/configs/config_fitcompress.yaml index cd8e502..e18b8e8 100644 --- a/src/pquant/configs/config_fitcompress.yaml +++ b/src/pquant/configs/config_fitcompress.yaml @@ -11,6 +11,7 @@ quantization_parameters: default_data_keep_negatives: 0. default_data_integer_bits: 0. default_data_fractional_bits: 7. + dynamic_data_quantization: false granularity: "per_tensor" quantize_input: true quantize_output: false diff --git a/src/pquant/configs/config_mdmm.yaml b/src/pquant/configs/config_mdmm.yaml index abbae3c..7586882 100644 --- a/src/pquant/configs/config_mdmm.yaml +++ b/src/pquant/configs/config_mdmm.yaml @@ -24,6 +24,7 @@ quantization_parameters: default_data_keep_negatives: 0. default_data_integer_bits: 0. default_data_fractional_bits: 7. + dynamic_data_quantization: false granularity: "per_tensor" quantize_input: true quantize_output: false diff --git a/src/pquant/configs/config_pdp.yaml b/src/pquant/configs/config_pdp.yaml index 937b5b9..1461d76 100644 --- a/src/pquant/configs/config_pdp.yaml +++ b/src/pquant/configs/config_pdp.yaml @@ -14,6 +14,7 @@ quantization_parameters: default_data_keep_negatives: 0. default_data_integer_bits: 0. default_data_fractional_bits: 7. + dynamic_data_quantization: false granularity: "per_tensor" quantize_input: true quantize_output: false diff --git a/src/pquant/configs/config_wanda.yaml b/src/pquant/configs/config_wanda.yaml index f128994..ab9cdae 100644 --- a/src/pquant/configs/config_wanda.yaml +++ b/src/pquant/configs/config_wanda.yaml @@ -16,6 +16,7 @@ quantization_parameters: default_data_keep_negatives: 0. default_data_integer_bits: 0. default_data_fractional_bits: 7. + dynamic_data_quantization: false granularity: "per_tensor" quantize_input: true quantize_output: false diff --git a/src/pquant/configs/finetuning.yaml b/src/pquant/configs/finetuning.yaml index 70a6237..165aef5 100644 --- a/src/pquant/configs/finetuning.yaml +++ b/src/pquant/configs/finetuning.yaml @@ -13,6 +13,7 @@ quantization_parameters: integer_bits: 4 fractional_bits: 6 use_high_granularity_quantization: false + dynamic_data_quantization: false granularity: "per_tensor" use_real_tanh: false use_symmetric_quantization: false diff --git a/src/pquant/core/constants.py b/src/pquant/core/constants.py index 993042b..830d062 100644 --- a/src/pquant/core/constants.py +++ b/src/pquant/core/constants.py @@ -1,3 +1,5 @@ +from enum import Enum + import optuna from pquant.data_models.pruning_model import ( @@ -10,15 +12,13 @@ PDPPruningModel, WandaPruningModel, ) -from pquant.pruning_methods.constraint_functions import ( - EqualityConstraint, - GreaterThanOrEqualConstraint, - LessThanOrEqualConstraint, -) -from pquant.pruning_methods.metric_functions import ( - StructuredSparsityMetric, - UnstructuredSparsityMetric, -) + + +class QuantizationGranularity(str, Enum): + PER_TENSOR = "per_tensor" + PER_CHANNEL = "per_channel" + PER_WEIGHT = "per_weight" + PRUNING_MODEL_REGISTRY = { "cs": CSPruningModel, @@ -47,21 +47,9 @@ DB_STORAGE = "sqlite:///optuna_study.db" TORCH_BACKEND = "torch" -TF_BACKEND = 'tensorflow' +TF_BACKEND = "tensorflow" FINETUNING_DIRECTION = {"maximize", "minimize"} CONFIG_FILE = "config.yaml" N_JOBS = 1 - - -METRIC_REGISTRY = { - "UnstructuredSparsity": UnstructuredSparsityMetric, - "StructuredSparsity": StructuredSparsityMetric, -} - -CONSTRAINT_REGISTRY = { - "Equality": EqualityConstraint, - "LessThanOrEqual": LessThanOrEqualConstraint, - "GreaterThanOrEqual": GreaterThanOrEqualConstraint, -} diff --git a/src/pquant/core/hyperparameter_optimization.py b/src/pquant/core/hyperparameter_optimization.py index 55a6077..b7c9d36 100644 --- a/src/pquant/core/hyperparameter_optimization.py +++ b/src/pquant/core/hyperparameter_optimization.py @@ -2,11 +2,11 @@ import json import logging import os -from typing import Annotated, Callable, Dict, Optional, Union +from collections.abc import Callable +from typing import Annotated import keras import optuna -import torch import yaml from pydantic import BaseModel, Field, field_validator @@ -37,12 +37,11 @@ def get_sampler(sampler_type, **kwargs): raise ValueError(f"Unknown sampler type: {sampler_type}") -def log_model_by_backend(model, name, signature=None, registered_model_name=None): +def log_model_by_backend(model, name, backend, signature=None, registered_model_name=None): import mlflow - backend = keras.backend.backend() kwargs = { - "artifact_path": name, + "name": name, "signature": signature, "registered_model_name": registered_model_name, } @@ -58,26 +57,26 @@ class MetricFunction(BaseModel): function_name: Callable direction: str - @field_validator('direction') + @field_validator("direction") def validate_direction(cls, direction): if direction not in constants.FINETUNING_DIRECTION: - raise ValueError("direction must be 'maximize' or 'minimize'") + raise ValueError("Direction must be 'maximize' or 'minimize'") return direction class PQConfig(BaseModel): hpo_parameters: BaseHyperparameterOptimizationModel pruning_parameters: Annotated[ - Union[ - CSPruningModel, - DSTPruningModel, - FITCompressPruningModel, - PDPPruningModel, - WandaPruningModel, - AutoSparsePruningModel, - ActivationPruningModel, - MDMMPruningModel, - ], + ( + CSPruningModel + | DSTPruningModel + | FITCompressPruningModel + | PDPPruningModel + | WandaPruningModel + | AutoSparsePruningModel + | ActivationPruningModel + | MDMMPruningModel + ), Field(discriminator="pruning_method"), ] quantization_parameters: BaseQuantizationModel @@ -86,10 +85,10 @@ class PQConfig(BaseModel): @classmethod def load_from_file(cls, path_to_config_file): - if path_to_config_file.endswith(('.yaml', '.yml')): + if path_to_config_file.endswith((".yaml", ".yml")): with open(path_to_config_file) as f: config_data = yaml.safe_load(f) - elif path_to_config_file.endswith('.json'): + elif path_to_config_file.endswith(".json"): with open(path_to_config_file) as f: config_data = json.load(f) else: @@ -115,16 +114,67 @@ def get_dict(self): return self.model_dump(mode="json") +class BackendAdapter: + def __init__(self, model): + self.backend = self._detect_backend(model) + self.device = None + + def clone_model(self, model): + if self.backend == constants.TORCH_BACKEND: + return copy.deepcopy(model) + elif self.backend == constants.TF_BACKEND: + new_model = keras.models.clone_model(model) + new_model.set_weights(model.get_weights()) + return new_model + + def get_backend(self): + return self.backend + + def get_device(self): + return self.device + + def _detect_backend(self, model): + if hasattr(model, "parameters"): + return constants.TORCH_BACKEND + elif isinstance(model, keras.Model): + return constants.TF_BACKEND + else: + raise ValueError("Unsupported model type") + + def move_to_device(self, model): + if self.backend == constants.TORCH_BACKEND: + self.device = next(model.parameters()).device + return model.to(self.device) + return model + + def eval(self, model): + if self.backend == constants.TORCH_BACKEND: + model.eval() + return model + + def tensor_to_numpy(self, tensor): + if self.backend == constants.TORCH_BACKEND: + return tensor.detach().cpu().numpy() + elif self.backend == constants.TF_BACKEND: + return tensor.numpy() + + def forward(self, model, x): + if self.backend == constants.TORCH_BACKEND: + x = x.to(self.device) + return model(x) + elif self.backend == constants.TF_BACKEND: + return model(x, training=False) + + class TuningTask: def __init__(self, config: PQConfig): self.config = config self.hyperparameters = {} - self.objectives: Dict[str, MetricFunction] = {} - self._training_function: Optional[Callable] = None - self._validation_function: Optional[Callable] = None - self._optimizer_function: Optional[Callable] = None - self._scheduler_function: Optional[Callable] = None - self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.objectives: dict[str, MetricFunction] = {} + self._training_function: Callable | None = None + self._validation_function: Callable | None = None + self._optimizer_function: Callable | None = None + self._scheduler_function: Callable | None = None self.enable_mlflow = False self.tracking_uri = None self.storage_db = None @@ -200,7 +250,7 @@ def set_hyperparameters(self): if numerical_params: self.set_numerical_params(numerical_params) - elif categorical_params: + if categorical_params: self.set_categorical_params(categorical_params) def set_numerical_params(self, numerical_params): @@ -245,50 +295,57 @@ def register_hyperparameter(self, name, optuna_func, *args, **kwargs): def objective(self, trial, model, train_func, valid_func, **kwargs): from pquant import add_compression_layers, train_model + config_copy = copy.deepcopy(self.config) + applied_parameters = {} for param_name, (optuna_func, func_args, func_kwargs) in self.hyperparameters.items(): new_value = optuna_func(trial, *func_args, **func_kwargs) logging.info(f"Suggested {param_name} = {new_value}") applied = False for sub_config in [ - self.config.training_parameters, - self.config.pruning_parameters, - self.config.quantization_parameters, - self.config.fitcompress_parameters, + config_copy.training_parameters, + config_copy.pruning_parameters, + config_copy.quantization_parameters, + config_copy.fitcompress_parameters, ]: if hasattr(sub_config, param_name): setattr(sub_config, param_name, new_value) + applied_parameters[param_name] = new_value applied = True break if not applied: logging.error(f"'{param_name}' not found in config: value not applied.") - trainloader = kwargs['trainloader'] + trainloader = kwargs["trainloader"] raw_input_batch = next(iter(trainloader)) - sample_input = raw_input_batch[0] - sample_output = model(sample_input.to(next(model.parameters()).device)) + sample_input = raw_input_batch[0] + model_copy = self.adapter.clone_model(model) + model_copy = self.adapter.move_to_device(model_copy) + sample_output = self.adapter.forward(model_copy, sample_input) input_shape = sample_input.shape - compressed_model = add_compression_layers(model, self.config, input_shape) + + compressed_model = add_compression_layers(model_copy, config_copy, input_shape) optimizer_func = self.get_optimizer_function() - optimizer = optimizer_func(self.config, compressed_model) + optimizer = optimizer_func(config_copy, compressed_model) scheduler_func = self.get_scheduler_function() - scheduler = scheduler_func(optimizer, self.config) + scheduler = scheduler_func(optimizer, config_copy) + device = self.adapter.get_device() trained_model = train_model( compressed_model, - self.config, + config_copy, train_func, valid_func, optimizer=optimizer, scheduler=scheduler, - device=self.device, + device=device, writer=None, **kwargs, ) - trained_model.eval() + self.adapter.eval(trained_model) objectives = [ - metric_object.function_name(trained_model, device=self.device, **kwargs) + metric_object.function_name(trained_model, device=device, **kwargs) for _, metric_object in self.objectives.items() ] @@ -297,23 +354,27 @@ def objective(self, trial, model, train_func, valid_func, **kwargs): from mlflow.models import infer_signature with mlflow.start_run(nested=True): - mlflow.log_params({param_name: getattr(self.config, param_name) for param_name in self.config.model_fields}) + mlflow.log_params(applied_parameters) mlflow.log_metrics({key: val for key, val in zip(self.objectives.keys(), objectives)}) - signature = infer_signature(sample_input.cpu().numpy(), sample_output.detach().cpu().numpy()) + signature = infer_signature( + self.adapter.tensor_to_numpy(sample_input), self.adapter.tensor_to_numpy(sample_output) + ) mlflow.log_text(yaml.safe_dump(self.get_dict()), "config.yaml") - model_name = self.config.hpo_parameters.model_name + model_name = config_copy.hpo_parameters.model_name log_model_by_backend( model=trained_model, name=model_name, signature=signature, registered_model_name=model_name, + backend=self.adapter.get_backend(), ) return objectives if len(objectives) > 1 else objectives[0] def run_optimization(self, model, **kwargs): hpo_parameters = self.config.hpo_parameters + num_trials = hpo_parameters.num_trials if self.enable_mlflow: import mlflow @@ -330,12 +391,11 @@ def run_optimization(self, model, **kwargs): load_if_exists=True, directions=[metric_object.direction for _, metric_object in self.objectives.items()], ) - - num_trials = hpo_parameters.num_trials + self.adapter = BackendAdapter(model) study.optimize( lambda trial: self.objective( trial, - copy.deepcopy(model.cpu()).to(self.device), + model, self.get_training_function(), self.get_validation_function(), **kwargs, diff --git a/src/pquant/core/keras/activations.py b/src/pquant/core/keras/activations.py index dcf5f6b..fccbe8b 100644 --- a/src/pquant/core/keras/activations.py +++ b/src/pquant/core/keras/activations.py @@ -1,7 +1,8 @@ -from typing import Tuple +from math import prod from typing import TypeVar as T import keras +from keras import ops from keras.ops import maximum, minimum, relu, tanh from pquant.core.keras.quantizer import Quantizer @@ -29,10 +30,11 @@ def __init__( self, config, activation="relu", - in_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, quantize_input=True, quantize_output=False, + enable_ebops=True, **kwargs, ): super().__init__(**kwargs) @@ -55,8 +57,14 @@ def __init__( self.k_output, self.i_output, self.f_output = out_quant_bits self.in_quant_bits = in_quant_bits self.out_quant_bits = out_quant_bits - self.activation_name = activation.lower() - self.activation_function = activation_registry.get(self.activation_name) + if isinstance(activation, str): + self.activation_name = activation.lower() + self.activation_function = activation_registry.get(self.activation_name) + else: + # A callable was passed directly instead of a registry key (e.g. the exp/inv + # lookup-table functions used by QSoftmax). + self.activation_function = activation + self.activation_name = getattr(activation, "__name__", activation.__class__.__name__).lower() self.config = config self.enable_quantization = config.quantization_parameters.enable_quantization self.use_hgq = config.quantization_parameters.use_high_granularity_quantization @@ -68,11 +76,13 @@ def __init__( self.hgq_gamma = config.quantization_parameters.hgq_gamma self.hgq_heterogeneous = config.quantization_parameters.hgq_heterogeneous self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress + self.dynamic_data = config.quantization_parameters.dynamic_data_quantization self.post_fitcompress_calibration = False self.saved_inputs = [] self.quantize_input = quantize_input self.quantize_output = quantize_output + self.enable_ebops = enable_ebops self.built = False def build(self, input_shape): @@ -89,6 +99,8 @@ def build(self, input_shape): is_heterogeneous=self.use_hgq, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, + granularity=self.config.quantization_parameters.granularity, ) if self.quantize_output: self.output_quantizer = Quantizer( @@ -101,6 +113,8 @@ def build(self, input_shape): is_heterogeneous=self.use_hgq, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, + granularity=self.config.quantization_parameters.granularity, ) if self.use_multiplier: @@ -127,9 +141,11 @@ def post_pre_train_function(self): self.output_quantizer.post_pre_train_function() def ebops(self): + if not self.enable_ebops: + return 0.0 if self.quantize_input and self.quantize_output: - bw_inp = self.input_quantizer.quantizer.bits_(self.input_shape) - bw_out = self.output_quantizer.quantizer.bits_(self.input_shape) + bw_inp = self.input_quantizer.get_total_bits(self.input_shape) + bw_out = self.output_quantizer.get_total_bits(self.input_shape) return keras.ops.sum((2.0**bw_inp) * bw_out) * 1e-4 # type: ignore return 0.0 @@ -175,6 +191,7 @@ def get_config(self): "config": self.config.get_dict(), "quantize_input": self.quantize_input, "quantize_output": self.quantize_output, + "enable_ebops": self.enable_ebops, "activation": self.activation_name, "in_quant_bits": self.in_quant_bits, "out_quant_bits": self.out_quant_bits, @@ -184,3 +201,254 @@ def get_config(self): def extra_repr(self): return f"quantize_input = {self.quantize_input}, quantize_output = {self.quantize_output}" + + +@keras.saving.register_keras_serializable(package="PQuantML") +class PQSoftmax(keras.layers.Layer): + """Quantized softmax that mirrors HGQ's ``QSoftmax``. + + Args: + config: PQuant configuration object (or a serialized dict). + axis: Axis (or axes) the softmax normalizes over. + stable: If True, subtract the max before exponentiating for numerical stability. + input_scaler: Scalar multiplied with the logits before exponentiating. + parallelization_factor: hls4ml parallelization factor used in the ebops cost model. + quantize_input: Whether to quantize the softmax logits before the exp table. + quantize_output: Whether to quantize the exp * inv product (the softmax output). + in_quant_bits: (k, i, f) bits for the softmax input quantizer. + out_quant_bits: (k, i, f) bits for the softmax output quantizer. + exp_in_quant_bits / exp_out_quant_bits: (k, i, f) bits for the exp table. + inv_in_quant_bits / inv_out_quant_bits: (k, i, f) bits for the inv table. + """ + + def __init__( + self, + config, + axis: int | tuple[int, ...] = -1, + stable: bool = True, + input_scaler: float = 1.0, + parallelization_factor: int = -1, + quantize_input: bool = True, + quantize_output: bool = False, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + exp_in_quant_bits: tuple[T, T, T] = None, + exp_out_quant_bits: tuple[T, T, T] = None, + inv_in_quant_bits: tuple[T, T, T] = None, + inv_out_quant_bits: tuple[T, T, T] = None, + **kwargs, + ): + super().__init__(**kwargs) + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + + self.supports_masking = True + self._axis = tuple(axis) if isinstance(axis, (tuple, list)) else (axis,) + self.axes = self._axis + self.stable = stable + self.input_scaler = input_scaler + self.parallelization_factor = parallelization_factor + self.quantize_input = quantize_input + self.quantize_output = quantize_output + self.epsilon = keras.backend.epsilon() + + if in_quant_bits is not None: + self.k_input, self.i_input, self.f_input = in_quant_bits + else: + self.k_input = config.quantization_parameters.default_data_keep_negatives + self.i_input = config.quantization_parameters.default_data_integer_bits + self.f_input = config.quantization_parameters.default_data_fractional_bits + + if out_quant_bits is not None: + self.k_output, self.i_output, self.f_output = out_quant_bits + else: + self.k_output = 0 + self.i_output = config.quantization_parameters.default_data_integer_bits + self.f_output = config.quantization_parameters.default_data_fractional_bits + + self.in_quant_bits = in_quant_bits + self.out_quant_bits = out_quant_bits + self.exp_in_quant_bits = exp_in_quant_bits + self.exp_out_quant_bits = exp_out_quant_bits + self.inv_in_quant_bits = inv_in_quant_bits + self.inv_out_quant_bits = inv_out_quant_bits + + self.overflow_mode_data = config.quantization_parameters.overflow_mode_data + self.round_mode = config.quantization_parameters.round_mode + self.use_hgq = config.quantization_parameters.use_high_granularity_quantization + self.hgq_gamma = config.quantization_parameters.hgq_gamma + self.hgq_beta = config.quantization_parameters.hgq_beta + self.enable_quantization = config.quantization_parameters.enable_quantization + self.dynamic_data = config.quantization_parameters.dynamic_data_quantization + self.is_pretraining = True + + i_data = config.quantization_parameters.default_data_integer_bits + f_data = config.quantization_parameters.default_data_fractional_bits + k_data = config.quantization_parameters.default_data_keep_negatives + exp_in_quant_bits = exp_in_quant_bits if exp_in_quant_bits is not None else (k_data, i_data, f_data) + exp_out_quant_bits = exp_out_quant_bits if exp_out_quant_bits is not None else (0, i_data, f_data) + inv_in_quant_bits = inv_in_quant_bits if inv_in_quant_bits is not None else (k_data, i_data, f_data) + inv_out_quant_bits = inv_out_quant_bits if inv_out_quant_bits is not None else (0, i_data, f_data) + + def _exp(x): + if self.stable: + return ops.exp(-x * self.input_scaler) + return ops.exp(x * self.input_scaler) + + def _inv(x): + return 1.0 / (x + self.epsilon) + + self.exp_table = PQActivation( + config, + _exp, + in_quant_bits=exp_in_quant_bits, + out_quant_bits=exp_out_quant_bits, + quantize_input=stable, + quantize_output=True, + enable_ebops=stable, + ) + self.inv_table = PQActivation( + config, + _inv, + in_quant_bits=inv_in_quant_bits, + out_quant_bits=inv_out_quant_bits, + quantize_input=True, + quantize_output=True, + ) + + def build(self, input_shape): + ndim = len(input_shape) + self.axes = tuple(sorted(a if a >= 0 else a + ndim for a in self._axis)) + self.input_shape = (1,) + tuple(input_shape[1:]) + + def _data_quantizer(k, i, f): + return Quantizer( + k=k, + i=i, + f=f, + overflow=self.overflow_mode_data, + round_mode=self.round_mode, + is_data=True, + is_heterogeneous=self.use_hgq, + hgq_gamma=self.hgq_gamma, + place="datalane", + dynamic_data=self.dynamic_data, + ) + + self.input_quantizer = _data_quantizer(self.k_input, self.i_input, self.f_input) + self.output_quantizer = _data_quantizer(self.k_output, self.i_output, self.f_output) + accum_shape = tuple(1 if i in self.axes else s for i, s in enumerate(input_shape)) + self.input_quantizer.build(input_shape) + self.output_quantizer.build(input_shape) + if not self.exp_table.built: + self.exp_table.build(input_shape) + if not self.inv_table.built: + self.inv_table.build(accum_shape) + for table, shape in ((self.exp_table, input_shape), (self.inv_table, accum_shape)): + if table.quantize_input and not table.input_quantizer.built: + table.input_quantizer.build(shape) + if table.quantize_output and not table.output_quantizer.built: + table.output_quantizer.build(shape) + super().build(input_shape) + + def get_input_quantization_bits(self): + return self.input_quantizer.get_quantization_bits() + + def get_output_quantization_bits(self): + return self.output_quantizer.get_quantization_bits() + + def post_pre_train_function(self): + self.is_pretraining = False + if self.quantize_input: + self.input_quantizer.post_pre_train_function() + if self.quantize_output: + self.output_quantizer.post_pre_train_function() + self.exp_table.post_pre_train_function() + self.inv_table.post_pre_train_function() + + def ebops(self): + shape = self.input_shape + accum_shape = tuple(1 if i in self.axes else s for i, s in enumerate(shape)) + max_instance = prod(accum_shape) + n_instance = self.parallelization_factor if self.parallelization_factor > 0 else max_instance + factor = n_instance / max_instance + + inp_bits = self.input_quantizer.get_total_bits(shape) + exp_bits = self.exp_table.output_quantizer.get_total_bits(shape) + inv_bits = self.inv_table.output_quantizer.get_total_bits(accum_shape) + + substract_ebops = ops.sum(inp_bits) if self.stable else 0.0 + accum_ebops = ops.sum(exp_bits) - ops.sum(ops.min(exp_bits, axis=self.axes)) + mult_ebops = ops.sum(exp_bits * inv_bits) + + ebops = substract_ebops + accum_ebops + mult_ebops + if not self.stable: + ebops = ebops + ops.sum((2.0**inp_bits) * exp_bits) * 1e-4 + ebops = ebops * factor + return ebops + self.exp_table.ebops() + self.inv_table.ebops() + + def hgq_loss(self): + if self.is_pretraining or not self.use_hgq: + return 0.0 + loss = self.hgq_beta * self.ebops() + if self.quantize_input: + loss += self.input_quantizer.hgq_loss() + if self.quantize_output: + loss += self.output_quantizer.hgq_loss() + return loss + + def call(self, inputs, training=None, mask=None): + if self.quantize_input and self.enable_quantization: + inputs = self.input_quantizer(inputs, training=training) + + if self.stable: + inputs = ops.max(inputs, axis=self.axes, keepdims=True) - inputs + + exp_inp = self.exp_table(inputs) + + if mask is not None: + exp_inp = ops.cast(mask, exp_inp.dtype) * exp_inp + + sums = ops.sum(exp_inp, axis=self.axes, keepdims=True) + divisor = self.inv_table(sums) + + out = exp_inp * divisor + if self.quantize_output and self.enable_quantization: + out = self.output_quantizer(out, training=training) + return out + + def compute_output_shape(self, input_shape): + return input_shape + + def get_config(self): + config = super().get_config() + config.update( + { + "config": self.config.get_dict(), + "axis": self.axes, + "stable": self.stable, + "input_scaler": self.input_scaler, + "parallelization_factor": self.parallelization_factor, + "quantize_input": self.quantize_input, + "quantize_output": self.quantize_output, + "in_quant_bits": self.in_quant_bits, + "out_quant_bits": self.out_quant_bits, + "exp_in_quant_bits": self.exp_in_quant_bits, + "exp_out_quant_bits": self.exp_out_quant_bits, + "inv_in_quant_bits": self.inv_in_quant_bits, + "inv_out_quant_bits": self.inv_out_quant_bits, + } + ) + return config + + @classmethod + def from_config(cls, config): + config = config.copy() + config.pop("exp_table", None) + config.pop("inv_table", None) + config.pop("input_quantizer", None) + config.pop("output_quantizer", None) + return cls(**config) diff --git a/src/pquant/core/keras/convert_to_onnx.py b/src/pquant/core/keras/convert_to_onnx.py new file mode 100644 index 0000000..781b8fa --- /dev/null +++ b/src/pquant/core/keras/convert_to_onnx.py @@ -0,0 +1,1476 @@ +""" +Convert a PQuant Keras functional model to ONNX or QONNX format. + +Pass ``use_qonnx=True`` to emit QONNX ``Quant`` custom nodes (requires the +qonnx runtime). Pass ``use_qonnx=False`` (default) to emit standard +``Clip + QuantizeLinear + DequantizeLinear`` nodes runnable with plain +onnxruntime. + +Fixed-point (k, i, f) mapping +------------------------------ +QONNX: + scale = 2^(-f) + zero_point = 0 + bit_width = k + i + f + signed = int(k) + +Standard ONNX (QDQ): + scale = 2^(-f) + zero_point = 0 (int8 signed, uint8 unsigned) + clip range = [-2^i, 2^i - 2^(-f)] signed + = [0, 2^i - 2^(-f)] unsigned + +Keras weight layout (kernel always stored as HWIO regardless of data_format): + Conv2D kernel: [kH, kW, in/g, out] → [out, in/g, kH, kW] for ONNX + Conv1D kernel: [kL, in/g, out] → [out, in/g, kL] for ONNX + DepthwiseConv2D kernel: [kH, kW, in, dm] → [in*dm, 1, kH, kW] for ONNX + Dense kernel: [in, out] → stored as [out, in] for Gemm (transB=1) + +Data format: + channels_first Data flows as NCHW; Conv/Pool ONNX ops work naturally. + channels_last Transpose(NHWC→NCHW) is inserted before each Conv/Pool/BN + node and Transpose(NCHW→NHWC) is inserted after. The + logical data format in the ONNX graph therefore stays NHWC + at every inter-layer edge; only the PQ ops run internally in + NCHW. ONNX-aware optimisers (e.g. onnxsim) can fold the + redundant back-to-back transposes away. +""" + +import functools +import logging + +import keras +import numpy as np +import onnx +import onnx.helper as oh +import onnx.numpy_helper as onh +from keras import ops +from onnx import TensorProto + +from pquant.core.keras.activations import PQActivation +from pquant.core.keras.layers import ( + PQBatchNormalization, + PQConv1d, + PQConv2d, + PQDense, + PQDepthwiseConv2d, + PQMultiheadAttention, +) + +# --------------------------------------------------------------------------- +# QONNX Quant node +# --------------------------------------------------------------------------- + +ROUND_MODE_MAP = { + "TRN": "FLOOR", + "RND": "ROUND", + "RND_CONV": "ROUND", + "TRN_ZERO": "TRUNCATE", + "RND_ZERO": "ROUND", + "RND_MIN_INF": "FLOOR", + "RND_INF": "ROUND", +} + + +def _quant_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT"): + """Build a QONNX Quant node. k/i/f are numpy arrays. Returns ([node], output_name).""" + k_val = int(float(np.array(k).ravel()[0])) + f_arr = np.array(f, dtype=np.float32) + i_arr = np.array(i, dtype=np.float32) + if f_arr.size > 1: + i_arr = i_arr.ravel().max() + f_arr = f_arr.ravel().min() + i_val = float(i_arr) + f_val = float(f_arr) + scale = float(2.0 ** (-f_val)) + bit_width = float(k_val + i_val + f_val) + qonnx_rnd = ROUND_MODE_MAP.get(rounding_mode, "ROUND") + # SAT_SYM excludes the most-negative value → QONNX narrow=1 + narrow = 1 if (k_val == 1 and overflow_mode == "SAT_SYM") else 0 + + scale_name = f"{name_prefix}_scale" + zp_name = f"{name_prefix}_zero_point" + bw_name = f"{name_prefix}_bit_width" + out_name = f"{name_prefix}_quantized" + + initializers.append(onh.from_array(np.array(scale, dtype=np.float32), name=scale_name)) + initializers.append(onh.from_array(np.array(0.0, dtype=np.float32), name=zp_name)) + initializers.append(onh.from_array(np.array(bit_width, dtype=np.float32), name=bw_name)) + + node = oh.make_node( + op_type="Quant", + inputs=[input_name, scale_name, zp_name, bw_name], + outputs=[out_name], + domain="qonnx.custom_op.general", + signed=k_val, + narrow=narrow, + rounding_mode=qonnx_rnd, + ) + return [node], out_name + + +# --------------------------------------------------------------------------- +# Standard ONNX QDQ triple +# --------------------------------------------------------------------------- + + +def _qdq_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT", include_clip=True): + """Build QuantizeLinear+DequantizeLinear nodes, optionally preceded by a Clip. + + Returns ([nodes], output_name). Set include_clip=False to skip the Clip node + (safe when values are guaranteed to be in-range at inference time, since + QuantizeLinear saturates naturally). + """ + k_val = int(float(np.array(k).ravel()[0])) + i_val = float(np.array(i, dtype=np.float32).ravel()[0]) + f_val = float(np.array(f, dtype=np.float32).ravel()[0]) + scale = float(2.0 ** (-f_val)) + signed = k_val == 1 + + clip_max = float(2.0**i_val - 2.0 ** (-f_val)) + if not signed: + clip_min = 0.0 + elif overflow_mode == "SAT_SYM": + clip_min = -clip_max # symmetric: -(2^i - 2^(-f)) + else: + clip_min = float(-(2.0**i_val)) # SAT: -2^i + zp_val = np.int8(0) if signed else np.uint8(0) + + scale_name = f"{name_prefix}_scale" + zp_name = f"{name_prefix}_zero_point" + quantized_name = f"{name_prefix}_quantized" + out_name = f"{name_prefix}_dequantized" + + initializers += [ + onh.from_array(np.array(scale, dtype=np.float32), name=scale_name), + onh.from_array(np.array(zp_val), name=zp_name), + ] + + if include_clip: + clip_min_name = f"{name_prefix}_clip_min" + clip_max_name = f"{name_prefix}_clip_max" + clipped_name = f"{name_prefix}_clipped" + initializers += [ + onh.from_array(np.array(clip_min, dtype=np.float32), name=clip_min_name), + onh.from_array(np.array(clip_max, dtype=np.float32), name=clip_max_name), + ] + nodes = [ + oh.make_node("Clip", inputs=[input_name, clip_min_name, clip_max_name], outputs=[clipped_name]), + oh.make_node("QuantizeLinear", inputs=[clipped_name, scale_name, zp_name], outputs=[quantized_name]), + ] + else: + nodes = [ + oh.make_node("QuantizeLinear", inputs=[input_name, scale_name, zp_name], outputs=[quantized_name]), + ] + + nodes.append(oh.make_node("DequantizeLinear", inputs=[quantized_name, scale_name, zp_name], outputs=[out_name])) + return nodes, out_name + + +# --------------------------------------------------------------------------- +# integer weight storage helper +# --------------------------------------------------------------------------- + + +def _int_weight_node(name_prefix, weight_np, k, i, f, initializers): # noqa: ARG001 (i unused) + """ + Store a weight tensor as int8/uint8 + DequantizeLinear. + + weight_np must already be in ONNX layout (transposed from Keras) and on the + fixed-point grid after apply_final_compression(). + + k/i/f are numpy arrays (may be per-tensor scalar or per-channel 1-D after + caller has already squeezed/reshaped appropriately). + + Granularity: + - per-tensor (f is scalar): single scale. + - per-channel (f is 1-D of length out_channels): axis=0 on weight tensor. + - per-weight (fully per-element): falls back to float32 storage. + + Returns ([node], output_name). + """ + k_np = np.array(k, dtype=np.float32) + f_np = np.array(f, dtype=np.float32) + k_val = int(float(k_np.ravel()[0])) + dtype = np.int8 if k_val == 1 else np.uint8 + out_channels = weight_np.shape[0] + out_name = f"{name_prefix}_dequantized" + + if f_np.size == 1: + # per-tensor + scale_np = np.array(float(2.0 ** (-float(f_np.ravel()[0]))), dtype=np.float32) + int_w = np.round(weight_np / float(scale_np)).astype(dtype) + per_ch = False + else: + f_1d = f_np.ravel() + if f_1d.size == out_channels: + # per-channel: one f value per output channel + scale_1d = (2.0 ** (-f_1d)).astype(np.float32) + bcast = scale_1d.reshape((out_channels,) + (1,) * (weight_np.ndim - 1)) + int_w = np.round(weight_np / bcast).astype(dtype) + scale_np = scale_1d + per_ch = True + else: + # per-weight: ONNX cannot represent; fall back to float32 + float_name = f"{name_prefix}_float" + initializers.append(onh.from_array(weight_np, name=float_name)) + return [], float_name + + int_name = f"{name_prefix}_int" + scale_name = f"{name_prefix}_dq_scale" + zp_name = f"{name_prefix}_dq_zp" + + zp_np = np.zeros(out_channels if per_ch else 1, dtype=dtype) + initializers += [ + onh.from_array(int_w, name=int_name), + onh.from_array(scale_np, name=scale_name), + onh.from_array(zp_np if per_ch else np.array(dtype(0)), name=zp_name), + ] + node_kwargs = {"axis": 0} if per_ch else {} + node = oh.make_node("DequantizeLinear", inputs=[int_name, scale_name, zp_name], outputs=[out_name], **node_kwargs) + return [node], out_name + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _np(tensor): + """Convert a Keras tensor / variable / scalar to a float32 numpy array.""" + return np.array(tensor, dtype=np.float32) + + +def _bn_transpose_info(layer): + """ + Return (need_transpose, perm_fwd, perm_bwd) for a BatchNormalization layer. + + ONNX BN (opset < 14) always normalises on axis 1 (NCHW). If the Keras + layer uses axis=-1 (channels_last), we must insert Transpose nodes around + the BN op. We infer ndim from the layer's stored input_shape. + """ + axis = getattr(layer, "axis", 1) + stored = getattr(layer, "input_shape", None) + ndim = len(stored) if stored is not None else 4 + eff_axis = axis if axis >= 0 else (ndim + axis) + + if eff_axis == 1 or ndim <= 2: + # channels already at position 1, or 2-D input — no transpose needed + return False, None, None + + if ndim == 4 and eff_axis == 3: + return True, [0, 3, 1, 2], [0, 2, 3, 1] + + if ndim == 3 and eff_axis == 2: + return True, [0, 2, 1], [0, 2, 1] + + # Fallback: general permutation that moves eff_axis to position 1 + perm_fwd = [0, eff_axis] + [i for i in range(1, ndim) if i != eff_axis] + # Inverse permutation + perm_bwd = [0] * ndim + for i, p in enumerate(perm_fwd): + perm_bwd[p] = i + return True, perm_fwd, perm_bwd + + +def _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn): + if ( + getattr(layer, "input_quantizer", None) is not None + and getattr(layer, "quantize_input", True) + and getattr(layer, "enable_quantization", True) + ): + q = layer.input_quantizer + k, i, f = q.get_quantization_bits() + new_nodes, current = quant_fn( + f"{prefix}_in", + current, + q.round_mode, + _np(k), + _np(i), + _np(f), + initializers, + overflow_mode=getattr(q, "overflow", "SAT"), + ) + nodes.extend(new_nodes) + return current + + +def _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn): + if ( + getattr(layer, "output_quantizer", None) is not None + and getattr(layer, "quantize_output", False) + and getattr(layer, "enable_quantization", True) + ): + q = layer.output_quantizer + k, i, f = q.get_quantization_bits() + new_nodes, current = quant_fn( + f"{prefix}_out", + current, + q.round_mode, + _np(k), + _np(i), + _np(f), + initializers, + overflow_mode=getattr(q, "overflow", "SAT"), + ) + nodes.extend(new_nodes) + return current + + +def _add_transpose(name, input_name, perm, nodes): + """Emit a Transpose node and return the output name.""" + out = f"{name}_transpose_{''.join(str(p) for p in perm)}" + nodes.append(oh.make_node("Transpose", inputs=[input_name], outputs=[out], perm=list(perm))) + return out + + +def _channels_last(layer): + return getattr(layer, "data_format", "channels_first") == "channels_last" + + +def _weight_f_for_onnx(f_np, out_channels): + """Squeeze/ravel a Keras per-channel f array to shape (out_channels,) for ONNX.""" + f_flat = f_np.ravel() + if f_flat.size == 1: + return f_flat # scalar, return as-is + if f_flat.size == out_channels: + return f_flat + # Per-element or mismatched: take the minimum to avoid overflow + return np.array([f_flat.min()], dtype=np.float32) + + +# --------------------------------------------------------------------------- +# per-layer graph builders +# --------------------------------------------------------------------------- + + +def _add_dense(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """Dense / PQDense. Keras kernel [in, out] stored as [out, in]; Gemm uses transB=1. + + Storing the weight pre-transposed means axis=0 is always the output dimension, + which is required for per-channel DequantizeLinear and avoids a runtime Transpose. + """ + current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + # Transpose kernel to [out, in]; Gemm will use transB=1 so Y = X @ W^T = X @ kernel. + kernel_np = _np(layer._kernel).T # [out, in] + out_units = kernel_np.shape[0] + + if use_qonnx: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + wfp_name = f"{prefix}_weight_fp" + initializers.append(onh.from_array(kernel_np, name=wfp_name)) + w_nodes, q_weight = _quant_node( + f"{prefix}_weight", + wfp_name, + layer.weight_quantizer.round_mode, + _np(k_w), + _np(i_w), + _np(f_w), + initializers, + overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(w_nodes) + elif store_integer_weights: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + f_np_w = _np(f_w) + f_for_onnx = _weight_f_for_onnx(f_np_w, out_units) + k_for_onnx = _weight_f_for_onnx(_np(k_w), out_units) + i_for_onnx = _weight_f_for_onnx(_np(i_w), out_units) + w_nodes, q_weight = _int_weight_node(f"{prefix}_weight", kernel_np, k_for_onnx, i_for_onnx, f_for_onnx, initializers) + nodes.extend(w_nodes) + else: + q_weight = f"{prefix}_weight" + initializers.append(onh.from_array(kernel_np, name=q_weight)) + + gemm_inputs = [current, q_weight] + + if layer._bias is not None: + bias_np = _np(layer._bias) + if use_qonnx: + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + bfp_name = f"{prefix}_bias_fp" + initializers.append(onh.from_array(bias_np, name=bfp_name)) + b_nodes, q_bias = _quant_node( + f"{prefix}_bias", + bfp_name, + layer.bias_quantizer.round_mode, + _np(k_b), + _np(i_b), + _np(f_b), + initializers, + overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights: + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, _np(k_b), _np(i_b), _np(f_b), initializers) + nodes.extend(b_nodes) + else: + q_bias = f"{prefix}_bias" + initializers.append(onh.from_array(bias_np, name=q_bias)) + gemm_inputs.append(q_bias) + + gemm_out = f"{prefix}_gemm" + nodes.append(oh.make_node("Gemm", inputs=gemm_inputs, outputs=[gemm_out], transB=1)) + current = gemm_out + + current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + return current + + +def _add_conv(layer, prefix, current, nodes, initializers, ndim, quant_fn, use_qonnx, store_integer_weights): + """PQConv2d / PQConv1d. Keras kernel: [*kernel, in/g, out] → ONNX [out, in/g, *kernel].""" + cl = _channels_last(layer) + + if cl: + perm_to_nchw = [0, 3, 1, 2] if ndim == 2 else [0, 2, 1] + perm_to_nhwx = [0, 2, 3, 1] if ndim == 2 else [0, 2, 1] + current = _add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) + + current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + kernel_np = _np(layer._kernel) + # Transpose kernel from Keras HWIO to ONNX OIHW + if ndim == 2: + kernel_onnx = np.transpose(kernel_np, (3, 2, 0, 1)) # [kH,kW,in,out] → [out,in,kH,kW] + else: + kernel_onnx = np.transpose(kernel_np, (2, 1, 0)) # [kL,in,out] → [out,in,kL] + + out_channels = kernel_onnx.shape[0] + + if use_qonnx: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + wfp_name = f"{prefix}_weight_fp" + initializers.append(onh.from_array(kernel_onnx, name=wfp_name)) + w_nodes, q_weight = _quant_node( + f"{prefix}_weight", + wfp_name, + layer.weight_quantizer.round_mode, + _np(k_w), + _np(i_w), + _np(f_w), + initializers, + overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(w_nodes) + elif store_integer_weights: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + f_for_onnx = _weight_f_for_onnx(_np(f_w), out_channels) + k_for_onnx = _weight_f_for_onnx(_np(k_w), out_channels) + i_for_onnx = _weight_f_for_onnx(_np(i_w), out_channels) + w_nodes, q_weight = _int_weight_node( + f"{prefix}_weight", kernel_onnx, k_for_onnx, i_for_onnx, f_for_onnx, initializers + ) + nodes.extend(w_nodes) + else: + q_weight = f"{prefix}_weight" + initializers.append(onh.from_array(kernel_onnx, name=q_weight)) + + conv_inputs = [current, q_weight] + + if layer._bias is not None: + bias_np = _np(layer._bias) + if use_qonnx: + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + bfp_name = f"{prefix}_bias_fp" + initializers.append(onh.from_array(bias_np, name=bfp_name)) + b_nodes, q_bias = _quant_node( + f"{prefix}_bias", + bfp_name, + layer.bias_quantizer.round_mode, + _np(k_b), + _np(i_b), + _np(f_b), + initializers, + overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights: + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, _np(k_b), _np(i_b), _np(f_b), initializers) + nodes.extend(b_nodes) + else: + q_bias = f"{prefix}_bias" + initializers.append(onh.from_array(bias_np, name=q_bias)) + conv_inputs.append(q_bias) + + # Padding + padding = layer.padding + if isinstance(padding, str): + auto_pad = "SAME_UPPER" if padding == "same" else "VALID" + pads = None + else: + p = list(padding) if hasattr(padding, "__iter__") else [padding] * ndim + pads = p + p # ONNX format: [begin_0, begin_1, ..., end_0, end_1, ...] + auto_pad = "NOTSET" + + to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 + conv_attrs = dict( + kernel_shape=to_list(layer.kernel_size, ndim), + strides=to_list(layer.strides, ndim), + dilations=to_list(layer.dilation_rate, ndim), + group=getattr(layer, "groups", 1), + auto_pad=auto_pad, + ) + if pads is not None: + conv_attrs["pads"] = pads + + conv_out = f"{prefix}_conv" + nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) + current = conv_out + + current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + + if cl: + current = _add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) + return current + + +def _add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """PQDepthwiseConv2d. + + Keras kernel: [kH, kW, in, depth_mult] + ONNX Conv with groups=in: weight [in*depth_mult, 1, kH, kW] + """ + cl = _channels_last(layer) + + if cl: + current = _add_transpose(f"{prefix}_pre", current, [0, 3, 1, 2], nodes) + + current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + kernel_np = _np(layer._kernel) # [kH, kW, in, depth_mult] + in_ch, depth_mult = kernel_np.shape[2], kernel_np.shape[3] + # Rearrange to [in*depth_mult, 1, kH, kW] + kernel_onnx = np.transpose(kernel_np, (2, 3, 0, 1)).reshape(in_ch * depth_mult, 1, *kernel_np.shape[:2]) + + out_channels = kernel_onnx.shape[0] + + if use_qonnx: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + wfp_name = f"{prefix}_weight_fp" + initializers.append(onh.from_array(kernel_onnx, name=wfp_name)) + w_nodes, q_weight = _quant_node( + f"{prefix}_weight", + wfp_name, + layer.weight_quantizer.round_mode, + _np(k_w), + _np(i_w), + _np(f_w), + initializers, + overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(w_nodes) + elif store_integer_weights: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + f_for_onnx = _weight_f_for_onnx(_np(f_w), out_channels) + k_for_onnx = _weight_f_for_onnx(_np(k_w), out_channels) + i_for_onnx = _weight_f_for_onnx(_np(i_w), out_channels) + w_nodes, q_weight = _int_weight_node( + f"{prefix}_weight", kernel_onnx, k_for_onnx, i_for_onnx, f_for_onnx, initializers + ) + nodes.extend(w_nodes) + else: + q_weight = f"{prefix}_weight" + initializers.append(onh.from_array(kernel_onnx, name=q_weight)) + + conv_inputs = [current, q_weight] + + if layer._bias is not None: + bias_np = _np(layer._bias) + if use_qonnx: + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + bfp_name = f"{prefix}_bias_fp" + initializers.append(onh.from_array(bias_np, name=bfp_name)) + b_nodes, q_bias = _quant_node( + f"{prefix}_bias", + bfp_name, + layer.bias_quantizer.round_mode, + _np(k_b), + _np(i_b), + _np(f_b), + initializers, + overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights: + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, _np(k_b), _np(i_b), _np(f_b), initializers) + nodes.extend(b_nodes) + else: + q_bias = f"{prefix}_bias" + initializers.append(onh.from_array(bias_np, name=q_bias)) + conv_inputs.append(q_bias) + + padding = layer.padding + if isinstance(padding, str): + auto_pad = "SAME_UPPER" if padding == "same" else "VALID" + pads = None + else: + p = list(padding) if hasattr(padding, "__iter__") else [padding, padding] + pads = p + p + auto_pad = "NOTSET" + + to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 + conv_attrs = dict( + kernel_shape=to_list(layer.kernel_size, 2), + strides=to_list(layer.strides, 2), + dilations=to_list(layer.dilation_rate, 2), + group=in_ch, + auto_pad=auto_pad, + ) + if pads is not None: + conv_attrs["pads"] = pads + + conv_out = f"{prefix}_conv" + nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) + current = conv_out + + current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + + if cl: + current = _add_transpose(f"{prefix}_post", current, [0, 2, 3, 1], nodes) + return current + + +def _add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """PQBatchNormalization / standard BatchNormalization.""" + need_tr, perm_to_nchw, perm_to_nhwx = _bn_transpose_info(layer) + + if need_tr: + current = _add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) + + current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + is_pq = isinstance(layer, PQBatchNormalization) + + gamma_np = _np(layer.gamma) if layer.gamma is not None else None + beta_np = _np(layer.beta) if layer.beta is not None else None + + if gamma_np is None: + # scale=False: use ones + n_ch = _np(layer.moving_mean).shape[0] + gamma_np = np.ones(n_ch, dtype=np.float32) + if beta_np is None: + # center=False: use zeros + n_ch = _np(layer.moving_mean).shape[0] + beta_np = np.zeros(n_ch, dtype=np.float32) + + if is_pq and use_qonnx: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + gfp = f"{prefix}_gamma_fp" + initializers.append(onh.from_array(gamma_np, name=gfp)) + g_nodes, q_gamma = _quant_node( + f"{prefix}_gamma", + gfp, + layer.weight_quantizer.round_mode, + _np(k_w), + _np(i_w), + _np(f_w), + initializers, + overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(g_nodes) + + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + bfp = f"{prefix}_beta_fp" + initializers.append(onh.from_array(beta_np, name=bfp)) + b_nodes, q_beta = _quant_node( + f"{prefix}_beta", + bfp, + layer.bias_quantizer.round_mode, + _np(k_b), + _np(i_b), + _np(f_b), + initializers, + overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif is_pq and store_integer_weights: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + g_nodes, q_gamma = _int_weight_node(f"{prefix}_gamma", gamma_np, _np(k_w), _np(i_w), _np(f_w), initializers) + nodes.extend(g_nodes) + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + b_nodes, q_beta = _int_weight_node(f"{prefix}_beta", beta_np, _np(k_b), _np(i_b), _np(f_b), initializers) + nodes.extend(b_nodes) + else: + q_gamma = f"{prefix}_gamma" + q_beta = f"{prefix}_beta" + initializers.append(onh.from_array(gamma_np, name=q_gamma)) + initializers.append(onh.from_array(beta_np, name=q_beta)) + + mean_name = f"{prefix}_running_mean" + var_name = f"{prefix}_running_var" + initializers.append(onh.from_array(_np(layer.moving_mean), name=mean_name)) + initializers.append(onh.from_array(_np(layer.moving_variance), name=var_name)) + + bn_out = f"{prefix}_bn" + nodes.append( + oh.make_node( + "BatchNormalization", + inputs=[current, q_gamma, q_beta, mean_name, var_name], + outputs=[bn_out], + epsilon=float(layer.epsilon), + ) + ) + current = bn_out + + if need_tr: + current = _add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) + return current + + +def _add_dense_nd(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """PQDense for rank-3 inputs (B, T, E). + + Uses MatMul + Add instead of Gemm so the op works for any rank ≥ 2. + The kernel is stored as [out, in] (same layout as _add_dense / _int_weight_node), + then transposed to [in, out] at runtime via a Transpose node so that + MatMul(input, kernel_t) broadcasts correctly over the sequence dimension. + """ + current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + kernel_np = _np(layer._kernel).T # [out, in] + out_units = kernel_np.shape[0] + + if use_qonnx: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + wfp_name = f"{prefix}_weight_fp" + initializers.append(onh.from_array(kernel_np, name=wfp_name)) + w_nodes, q_weight = _quant_node( + f"{prefix}_weight", + wfp_name, + layer.weight_quantizer.round_mode, + _np(k_w), + _np(i_w), + _np(f_w), + initializers, + overflow_mode=getattr(layer.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(w_nodes) + elif store_integer_weights: + k_w, i_w, f_w = layer.weight_quantizer.get_quantization_bits() + f_for_onnx = _weight_f_for_onnx(_np(f_w), out_units) + k_for_onnx = _weight_f_for_onnx(_np(k_w), out_units) + i_for_onnx = _weight_f_for_onnx(_np(i_w), out_units) + w_nodes, q_weight = _int_weight_node(f"{prefix}_weight", kernel_np, k_for_onnx, i_for_onnx, f_for_onnx, initializers) + nodes.extend(w_nodes) + else: + q_weight = f"{prefix}_weight" + initializers.append(onh.from_array(kernel_np, name=q_weight)) + + # Transpose [out, in] → [in, out] so MatMul(input[..., in], kernel_t[in, out]) works + kernel_t_name = f"{prefix}_weight_t" + nodes.append(oh.make_node("Transpose", inputs=[q_weight], outputs=[kernel_t_name], perm=[1, 0])) + + mm_out = f"{prefix}_mm" + nodes.append(oh.make_node("MatMul", inputs=[current, kernel_t_name], outputs=[mm_out])) + current = mm_out + + if layer._bias is not None: + bias_np = _np(layer._bias) + if use_qonnx: + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + bfp_name = f"{prefix}_bias_fp" + initializers.append(onh.from_array(bias_np, name=bfp_name)) + b_nodes, q_bias = _quant_node( + f"{prefix}_bias", + bfp_name, + layer.bias_quantizer.round_mode, + _np(k_b), + _np(i_b), + _np(f_b), + initializers, + overflow_mode=getattr(layer.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights: + k_b, i_b, f_b = layer.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, _np(k_b), _np(i_b), _np(f_b), initializers) + nodes.extend(b_nodes) + else: + q_bias = f"{prefix}_bias" + initializers.append(onh.from_array(bias_np, name=q_bias)) + add_out = f"{prefix}_bias_add" + nodes.append(oh.make_node("Add", inputs=[current, q_bias], outputs=[add_out])) + current = add_out + + current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + return current + + +def _add_mha(layer, prefix, q_input, k_input, v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """Build ONNX nodes for PQMultiheadAttention (Keras version, always batch-first). + + Decomposes multi-head attention into primitive ONNX ops: + + Q/K/V MatMul projections (rank-3 MatMul via _add_dense_nd) + Reshape (B, L, E) → (B, H, L, head_dim) + Transpose + MatMul(Q, K^T) * scale → optional Quant + Softmax → optional Quant + MatMul(attn_weights, V) → optional context Quant + Transpose + Reshape → (B, T, E) + out_proj MatMul + + Returns (out_name, avg_attn_weights_name). + """ + H = layer.num_heads + head_dim = layer.head_dim + E = layer.embed_dim + scale_val = float(layer.scale) + + # --- Q / K / V projections: (B, L, E) → (B, L, E) --- + q_proj_out = _add_dense_nd( + layer.q_proj, f"{prefix}_q_proj", q_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + k_proj_out = _add_dense_nd( + layer.k_proj, f"{prefix}_k_proj", k_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + v_proj_out = _add_dense_nd( + layer.v_proj, f"{prefix}_v_proj", v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + # --- Helper: (B, L, E) → (B, H, L, head_dim) using dynamic shapes --- + def _split_heads(x_name, pfx): + shape_out = f"{pfx}_shape" + b_scalar = f"{pfx}_b_sc" + l_scalar = f"{pfx}_l_sc" + b_1d = f"{pfx}_b_1d" + l_1d = f"{pfx}_l_1d" + h_1d_const = f"{pfx}_H_1d" + hd_1d_const = f"{pfx}_hd_1d" + shape_4d = f"{pfx}_shape4d" + reshaped = f"{pfx}_reshaped" + transposed = f"{pfx}_transposed" + idx0 = f"{pfx}_gi0" + idx1 = f"{pfx}_gi1" + ax0 = f"{pfx}_ax0" + + nodes.append(oh.make_node("Shape", inputs=[x_name], outputs=[shape_out])) + initializers.extend( + [ + onh.from_array(np.array(0, dtype=np.int64), name=idx0), + onh.from_array(np.array(1, dtype=np.int64), name=idx1), + onh.from_array(np.array([0], dtype=np.int64), name=ax0), + onh.from_array(np.array([H], dtype=np.int64), name=h_1d_const), + onh.from_array(np.array([head_dim], dtype=np.int64), name=hd_1d_const), + ] + ) + nodes.append(oh.make_node("Gather", inputs=[shape_out, idx0], outputs=[b_scalar])) + nodes.append(oh.make_node("Gather", inputs=[shape_out, idx1], outputs=[l_scalar])) + nodes.append(oh.make_node("Unsqueeze", inputs=[b_scalar, ax0], outputs=[b_1d])) + nodes.append(oh.make_node("Unsqueeze", inputs=[l_scalar, ax0], outputs=[l_1d])) + nodes.append(oh.make_node("Concat", inputs=[b_1d, l_1d, h_1d_const, hd_1d_const], outputs=[shape_4d], axis=0)) + nodes.append(oh.make_node("Reshape", inputs=[x_name, shape_4d], outputs=[reshaped])) + # (B, L, H, head_dim) → (B, H, L, head_dim) + nodes.append(oh.make_node("Transpose", inputs=[reshaped], outputs=[transposed], perm=[0, 2, 1, 3])) + return transposed + + q_h = _split_heads(q_proj_out, f"{prefix}_q") + k_h = _split_heads(k_proj_out, f"{prefix}_k") + v_h = _split_heads(v_proj_out, f"{prefix}_v") + + k_t_name = f"{prefix}_k_T" + nodes.append(oh.make_node("Transpose", inputs=[k_h], outputs=[k_t_name], perm=[0, 1, 3, 2])) + + raw_scores = f"{prefix}_scores_raw" + scaled_scores = f"{prefix}_scores_scaled" + scale_cst = f"{prefix}_attn_scale" + nodes.append(oh.make_node("MatMul", inputs=[q_h, k_t_name], outputs=[raw_scores])) + initializers.append(onh.from_array(np.array(scale_val, dtype=np.float32), name=scale_cst)) + nodes.append(oh.make_node("Mul", inputs=[raw_scores, scale_cst], outputs=[scaled_scores])) + current = scaled_scores + + if layer.softmax.quantize_input and getattr(layer, "enable_quantization", True): + q = layer.softmax.input_quantizer + k_q, i_q, f_q = q.get_quantization_bits() + q_nodes, current = quant_fn( + f"{prefix}_attn_score_q", + current, + q.round_mode, + _np(k_q), + _np(i_q), + _np(f_q), + initializers, + overflow_mode=getattr(q, "overflow", "SAT"), + ) + nodes.extend(q_nodes) + + # --- Softmax (axis=-1); approximate_softmax falls back to standard Softmax in ONNX --- + attn_w_name = f"{prefix}_attn_weights" + nodes.append(oh.make_node("Softmax", inputs=[current], outputs=[attn_w_name], axis=-1)) + current = attn_w_name + + # --- Softmax output quantization (the MHA enables the softmax's output quantizer) --- + if layer.softmax.quantize_output and getattr(layer, "enable_quantization", True): + q = layer.softmax.output_quantizer + k_q, i_q, f_q = q.get_quantization_bits() + q_nodes, current = quant_fn( + f"{prefix}_attn_weight_q", + current, + q.round_mode, + _np(k_q), + _np(i_q), + _np(f_q), + initializers, + overflow_mode=getattr(q, "overflow", "SAT"), + ) + nodes.extend(q_nodes) + + ctx_raw = f"{prefix}_ctx_raw" + nodes.append(oh.make_node("MatMul", inputs=[current, v_h], outputs=[ctx_raw])) + current_ctx = ctx_raw + + ctx_t = f"{prefix}_ctx_t" + ctx_shape = f"{prefix}_ctx_shape" + ctx_b_sc = f"{prefix}_ctx_b_sc" + ctx_t_sc = f"{prefix}_ctx_t_sc" + ctx_b_1d = f"{prefix}_ctx_b_1d" + ctx_t_1d = f"{prefix}_ctx_t_1d" + ctx_E_1d = f"{prefix}_ctx_E_1d" + ctx_ax0 = f"{prefix}_ctx_ax0" + ctx_gi0 = f"{prefix}_ctx_gi0" + ctx_gi1 = f"{prefix}_ctx_gi1" + ctx_3d = f"{prefix}_ctx_shape3d" + ctx_merged = f"{prefix}_ctx_merged" + + nodes.append(oh.make_node("Transpose", inputs=[current_ctx], outputs=[ctx_t], perm=[0, 2, 1, 3])) + nodes.append(oh.make_node("Shape", inputs=[ctx_t], outputs=[ctx_shape])) + initializers += [ + onh.from_array(np.array(0, dtype=np.int64), name=ctx_gi0), + onh.from_array(np.array(1, dtype=np.int64), name=ctx_gi1), + onh.from_array(np.array([0], dtype=np.int64), name=ctx_ax0), + onh.from_array(np.array([E], dtype=np.int64), name=ctx_E_1d), + ] + nodes.append(oh.make_node("Gather", inputs=[ctx_shape, ctx_gi0], outputs=[ctx_b_sc])) + nodes.append(oh.make_node("Gather", inputs=[ctx_shape, ctx_gi1], outputs=[ctx_t_sc])) + nodes.append(oh.make_node("Unsqueeze", inputs=[ctx_b_sc, ctx_ax0], outputs=[ctx_b_1d])) + nodes.append(oh.make_node("Unsqueeze", inputs=[ctx_t_sc, ctx_ax0], outputs=[ctx_t_1d])) + nodes.append(oh.make_node("Concat", inputs=[ctx_b_1d, ctx_t_1d, ctx_E_1d], outputs=[ctx_3d], axis=0)) + nodes.append(oh.make_node("Reshape", inputs=[ctx_t, ctx_3d], outputs=[ctx_merged])) + + # --- Output projection: (B, T, E) → (B, T, E) --- + out = _add_dense_nd( + layer.out_proj, f"{prefix}_out_proj", ctx_merged, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + # --- Average attention weights over heads: (B, H, T, S) → (B, T, S) --- + avg_attn = f"{prefix}_avg_attn_weights" + nodes.append(oh.make_node("ReduceMean", inputs=[attn_w_name], outputs=[avg_attn], axes=[1], keepdims=0)) + + return out, avg_attn + + +def _add_avgpool(layer, prefix, current, nodes, initializers, ndim, quant_fn): + cl = _channels_last(layer) + + if cl: + perm_to_nchw = [0, 3, 1, 2] if ndim == 2 else [0, 2, 1] + perm_to_nhwx = [0, 2, 3, 1] if ndim == 2 else [0, 2, 1] + current = _add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) + + current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 + + pool_out = f"{prefix}_pool" + nodes.append( + oh.make_node( + "AveragePool", + inputs=[current], + outputs=[pool_out], + kernel_shape=to_list(layer.pool_size, ndim), + strides=to_list(layer.strides, ndim), + pads=[0] * (ndim * 2), + count_include_pad=0, + ) + ) + current = pool_out + + current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + + if cl: + current = _add_transpose(f"{prefix}_post", current, perm_to_nhwx, nodes) + return current + + +def _add_global_avgpool(layer, prefix, current, nodes, ndim): + cl = _channels_last(layer) + + if cl: + perm_to_nchw = [0, 3, 1, 2] if ndim == 2 else [0, 2, 1] + current = _add_transpose(f"{prefix}_pre", current, perm_to_nchw, nodes) + + pool_out = f"{prefix}_global_pool" + nodes.append(oh.make_node("GlobalAveragePool", inputs=[current], outputs=[pool_out])) + current = pool_out + + if cl: + # GlobalAveragePool returns [N, C, 1, 1]; emit Flatten to [N, C]. + # Actually after GlobalAveragePool output is [N, C, 1, 1]; transpose back would give + # [N, 1, 1, C] which then needs squeezing — that's the same as just squeezing [N, C]. + # Emit Flatten to [N, C] instead of bothering with transpose. + flatten_name = f"{prefix}_flatten" + nodes.append(oh.make_node("Flatten", inputs=[pool_out], outputs=[flatten_name], axis=1)) + current = flatten_name + + return current + + +def _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn): + """PQActivation: [input QDQ] → [multiplier scale] → activation → [output QDQ]. + + Supported activations: relu, tanh, hard_tanh (= Clip(-1, 1)). + + The optional relu multiplier is baked to a constant: 2^round(m). + """ + # --- optional input quantization --- + current = _maybe_quant_input(layer, prefix, current, nodes, initializers, quant_fn) + + # --- optional learnable multiplier (relu only) --- + if ( + getattr(layer, "use_multiplier", False) + and getattr(layer, "activation_name", "") == "relu" + and hasattr(layer, "multiplier") + ): + m_val = float(np.array(layer.multiplier).ravel()[0]) + scale = float(2.0 ** round(m_val)) + scale_name = f"{prefix}_mul_scale" + scaled_out = f"{prefix}_scaled" + initializers.append(onh.from_array(np.array(scale, dtype=np.float32), name=scale_name)) + nodes.append(oh.make_node("Mul", inputs=[current, scale_name], outputs=[scaled_out])) + current = scaled_out + + # --- activation --- + act = getattr(layer, "activation_name", "relu") + act_out = f"{prefix}_act" + if act == "relu": + nodes.append(oh.make_node("Relu", inputs=[current], outputs=[act_out])) + elif act == "tanh": + nodes.append(oh.make_node("Tanh", inputs=[current], outputs=[act_out])) + elif act == "hard_tanh": + # hard_tanh(x) = clip(x, -1, 1) + cmin_name = f"{prefix}_htanh_min" + cmax_name = f"{prefix}_htanh_max" + initializers += [ + onh.from_array(np.array(-1.0, dtype=np.float32), name=cmin_name), + onh.from_array(np.array(1.0, dtype=np.float32), name=cmax_name), + ] + nodes.append(oh.make_node("Clip", inputs=[current, cmin_name, cmax_name], outputs=[act_out])) + else: + raise TypeError(f"PQActivation: unsupported activation {act!r} for ONNX export") + current = act_out + + # --- optional output quantization --- + current = _maybe_quant_output(layer, prefix, current, nodes, initializers, quant_fn) + return current + + +# --------------------------------------------------------------------------- +# shared layer dispatcher +# --------------------------------------------------------------------------- + + +def _emit_layer( + layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, input_onnx_names=None +): + """Emit ONNX nodes for a single Keras layer. Returns the ONNX output name.""" + + # --- PQuant layers --- + if isinstance(layer, PQMultiheadAttention): + # input_onnx_names = [query, key, value] or [single_input] for self-attention + if len(input_onnx_names) >= 3: + q_in, k_in, v_in = input_onnx_names[0], input_onnx_names[1], input_onnx_names[2] + elif len(input_onnx_names) == 2: + q_in, k_in, v_in = input_onnx_names[0], input_onnx_names[1], input_onnx_names[1] + else: + q_in = k_in = v_in = input_onnx_names[0] + return _add_mha(layer, prefix, q_in, k_in, v_in, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + + if isinstance(layer, PQActivation): + return _add_pq_activation(layer, prefix, current, nodes, initializers, quant_fn) + + if isinstance(layer, PQDense): + return _add_dense(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + + if isinstance(layer, PQDepthwiseConv2d): + return _add_depthwise_conv(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + + if isinstance(layer, PQConv2d): + return _add_conv( + layer, + prefix, + current, + nodes, + initializers, + ndim=2, + quant_fn=quant_fn, + use_qonnx=use_qonnx, + store_integer_weights=store_integer_weights, + ) + + if isinstance(layer, PQConv1d): + return _add_conv( + layer, + prefix, + current, + nodes, + initializers, + ndim=1, + quant_fn=quant_fn, + use_qonnx=use_qonnx, + store_integer_weights=store_integer_weights, + ) + + if isinstance(layer, PQBatchNormalization): + return _add_batchnorm(layer, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + + # --- Standard Keras layers --- + if isinstance(layer, keras.layers.BatchNormalization): + return _add_batchnorm( + layer, prefix, current, nodes, initializers, quant_fn=quant_fn, use_qonnx=False, store_integer_weights=False + ) + + if isinstance(layer, keras.layers.Conv2D): + # Plain Conv2D (non-PQ): wrap in a minimal shim + _layer = layer + _layer._bias = layer.bias + return _add_conv_plain(layer, prefix, current, nodes, initializers) + + if isinstance(layer, keras.layers.Dense): + out = f"{prefix}_gemm" + w_name = f"{prefix}_weight" + initializers.append(onh.from_array(_np(layer.kernel).T, name=w_name)) # store [out, in] + gemm_inputs = [current, w_name] + if layer.bias is not None: + b_name = f"{prefix}_bias" + initializers.append(onh.from_array(_np(layer.bias), name=b_name)) + gemm_inputs.append(b_name) + nodes.append(oh.make_node("Gemm", inputs=gemm_inputs, outputs=[out], transB=1)) + return out + + if isinstance(layer, (keras.layers.ReLU, keras.layers.Activation)): + activation = ( + layer.activation.__name__ + if isinstance(layer, keras.layers.Activation) and callable(layer.activation) + else getattr(layer, "activation", "relu") + ) + act_name = activation if isinstance(activation, str) else "relu" + out = f"{prefix}_act" + if "relu" in act_name.lower(): + nodes.append(oh.make_node("Relu", inputs=[current], outputs=[out])) + elif "sigmoid" in act_name.lower(): + nodes.append(oh.make_node("Sigmoid", inputs=[current], outputs=[out])) + elif "tanh" in act_name.lower(): + nodes.append(oh.make_node("Tanh", inputs=[current], outputs=[out])) + else: + raise TypeError(f"Unsupported Activation for ONNX export: {act_name!r}") + return out + + if isinstance(layer, keras.layers.Flatten): + out = f"{prefix}_flatten" + nodes.append(oh.make_node("Flatten", inputs=[current], outputs=[out], axis=1)) + return out + + if isinstance(layer, keras.layers.Reshape): + target_shape = list(layer.target_shape) + # Prepend batch dim (-1 means dynamic) + full_shape = [-1] + target_shape + shape_name = f"{prefix}_shape" + out = f"{prefix}_reshape" + initializers.append(onh.from_array(np.array(full_shape, dtype=np.int64), name=shape_name)) + nodes.append(oh.make_node("Reshape", inputs=[current, shape_name], outputs=[out])) + return out + + if isinstance(layer, keras.layers.Add): + assert input_onnx_names is not None and len(input_onnx_names) == 2 + out = f"{prefix}_add" + nodes.append(oh.make_node("Add", inputs=input_onnx_names, outputs=[out])) + return out + + if isinstance(layer, keras.layers.Concatenate): + assert input_onnx_names is not None + axis = layer.axis + # Negative axis: leave as-is; onnx Concat supports negative axes + out = f"{prefix}_concat" + nodes.append(oh.make_node("Concat", inputs=input_onnx_names, outputs=[out], axis=axis)) + return out + + if isinstance(layer, keras.layers.Multiply): + assert input_onnx_names is not None and len(input_onnx_names) == 2 + out = f"{prefix}_mul" + nodes.append(oh.make_node("Mul", inputs=input_onnx_names, outputs=[out])) + return out + + if isinstance(layer, keras.layers.AveragePooling2D): + return _add_avgpool(layer, prefix, current, nodes, initializers, ndim=2, quant_fn=quant_fn) + + if isinstance(layer, keras.layers.AveragePooling1D): + return _add_avgpool(layer, prefix, current, nodes, initializers, ndim=1, quant_fn=quant_fn) + + if isinstance(layer, keras.layers.GlobalAveragePooling2D): + return _add_global_avgpool(layer, prefix, current, nodes, ndim=2) + + if isinstance(layer, keras.layers.GlobalAveragePooling1D): + return _add_global_avgpool(layer, prefix, current, nodes, ndim=1) + + if isinstance(layer, (keras.layers.Dropout,)): + return current # identity at inference + + raise TypeError(f"Unsupported Keras layer type for ONNX export: {type(layer).__name__!r}") + + +def _add_conv_plain(layer, prefix, current, nodes, initializers): + """Emit a plain (non-PQ) Conv2D layer.""" + cl = _channels_last(layer) + if cl: + current = _add_transpose(f"{prefix}_pre", current, [0, 3, 1, 2], nodes) + + kernel_np = _np(layer.kernel) + kernel_onnx = np.transpose(kernel_np, (3, 2, 0, 1)) + w_name = f"{prefix}_weight" + initializers.append(onh.from_array(kernel_onnx, name=w_name)) + conv_inputs = [current, w_name] + + if layer.bias is not None: + b_name = f"{prefix}_bias" + initializers.append(onh.from_array(_np(layer.bias), name=b_name)) + conv_inputs.append(b_name) + + padding = layer.padding + auto_pad = "SAME_UPPER" if padding == "same" else "VALID" + to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 + conv_attrs = dict( + kernel_shape=to_list(layer.kernel_size, 2), + strides=to_list(layer.strides, 2), + dilations=to_list(layer.dilation_rate, 2), + group=layer.groups, + auto_pad=auto_pad, + ) + conv_out = f"{prefix}_conv" + nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) + current = conv_out + + if cl: + current = _add_transpose(f"{prefix}_post", current, [0, 2, 3, 1], nodes) + return current + + +# --------------------------------------------------------------------------- +# Keras functional model graph traversal +# --------------------------------------------------------------------------- + + +def _build_tensor_onnx_map(model): + """ + Return a dict mapping id(KerasTensor) → ONNX tensor name for model.inputs. + Multi-input models are supported; inputs are named "input_0", "input_1", etc. + (or just "input" for single-input models). + """ + tensor_to_onnx = {} + for i, inp in enumerate(model.inputs): + name = "input" if len(model.inputs) == 1 else f"input_{i}" + tensor_to_onnx[id(inp)] = name + return tensor_to_onnx + + +def _inbound_input_names(layer, tensor_to_onnx): + """Return the list of ONNX input names for this layer based on its inbound node.""" + if not layer._inbound_nodes: + return [] + node = layer._inbound_nodes[0] + input_tensors = node.input_tensors + if not isinstance(input_tensors, (list, tuple)): + input_tensors = [input_tensors] + result = [] + for t in input_tensors: + key = id(t) + if key not in tensor_to_onnx: + raise RuntimeError( + f"Layer {layer.name!r}: input tensor not found in tensor_to_onnx map. " + "Ensure model.layers is in topological order." + ) + result.append(tensor_to_onnx[key]) + return result + + +def _register_layer_output(layer, onnx_name, tensor_to_onnx): + """Register the ONNX output name for a layer's output tensor(s). + + onnx_name may be a plain string (single-output layer) or a tuple of strings + (multi-output layer, e.g. PQMultiheadAttention returns (out, avg_attn_weights)). + """ + if not layer._inbound_nodes: + return + node = layer._inbound_nodes[0] + out_tensors = node.output_tensors + if not isinstance(out_tensors, (list, tuple)): + out_tensors = [out_tensors] + if isinstance(onnx_name, (list, tuple)): + for tensor, name in zip(out_tensors, onnx_name): + tensor_to_onnx[id(tensor)] = name + else: + tensor_to_onnx[id(out_tensors[0])] = onnx_name + + +# --------------------------------------------------------------------------- +# main conversion +# --------------------------------------------------------------------------- + + +def convert_to_onnx( + model: keras.Model, + input_shape: tuple, + output_path: str = "model.onnx", + opset: int = 13, + use_qonnx: bool = False, + store_integer_weights: bool = False, + include_clip: bool = True, + batch_size: int | None = None, +) -> onnx.ModelProto: + """ + Convert a Keras functional model of PQuant layers to ONNX or QONNX. + + The model must have apply_final_compression() called on all PQ layers + before calling this function. Only inference-mode semantics are exported. + + Args: + model: Trained keras.Model. Must be a functional model + (built with the Keras functional API or subclassed + models whose layers are accessible via model.layers). + input_shape: Shape of a single sample excluding batch, e.g. (3, 32, 32). + For channels_last Conv models use e.g. (32, 32, 3). + output_path: Where to save the .onnx file. + opset: ONNX opset version (≥13 required for per-channel + DequantizeLinear). + use_qonnx: Emit QONNX Quant custom nodes if True. + store_integer_weights: Store weight initializers as int8/uint8 + + DequantizeLinear instead of float32 (ignored when + use_qonnx=True). + include_clip: Prepend a Clip node before each QuantizeLinear when + True (default). Set to False to emit bare + QuantizeLinear+DequantizeLinear pairs — safe when + values are guaranteed in-range at inference time since + QuantizeLinear saturates naturally. Ignored when + use_qonnx=True. + batch_size: If not None, fix the batch dimension of all graph + inputs and outputs to this value. If None (default), + the batch dimension is left dynamic. + + Returns: + The constructed onnx.ModelProto. + """ + quant_fn = _quant_node if use_qonnx else functools.partial(_qdq_node, include_clip=include_clip) + + onnx_nodes: list[onnx.NodeProto] = [] + initializers: list[onnx.TensorProto] = [] + + tensor_to_onnx = _build_tensor_onnx_map(model) + last_output_name: str = "" + + for layer in model.layers: + # Skip InputLayer — already seeded in tensor_to_onnx + if isinstance(layer, keras.layers.InputLayer): + continue + + input_onnx_names = _inbound_input_names(layer, tensor_to_onnx) + if not input_onnx_names: + continue + + current = input_onnx_names[0] # primary input (used by single-input layers) + prefix = layer.name.replace("/", "_").replace(":", "_") + + output_name = _emit_layer( + layer, + prefix, + current, + onnx_nodes, + initializers, + quant_fn, + use_qonnx, + store_integer_weights, + input_onnx_names=input_onnx_names, + ) + + _register_layer_output(layer, output_name, tensor_to_onnx) + # For multi-output layers (e.g. MHA returns (out, avg_attn)), track only the + # primary output as the graph's last output name. + last_output_name = output_name[0] if isinstance(output_name, tuple) else output_name + + # Determine output shape via a forward pass + dummy = np.zeros((1, *input_shape), dtype=np.float32) + dummy_out = model(dummy, training=False) + dummy_out_np = np.array(ops.convert_to_numpy(dummy_out)) + batch_dim = batch_size # None → dynamic, int → fixed + output_shape = [batch_dim] + list(dummy_out_np.shape[1:]) + + # Build ONNX graph + if len(model.inputs) == 1: + input_vis = [oh.make_tensor_value_info("input", TensorProto.FLOAT, [batch_dim, *input_shape])] + else: + input_vis = [ + oh.make_tensor_value_info(f"input_{i}", TensorProto.FLOAT, [batch_dim, *input_shape]) + for i in range(len(model.inputs)) + ] + output_vi = oh.make_tensor_value_info(last_output_name, TensorProto.FLOAT, output_shape) + + graph = oh.make_graph( + nodes=onnx_nodes, + name="pquant_keras_onnx", + inputs=input_vis, + outputs=[output_vi], + initializer=initializers, + ) + + opset_imports = [oh.make_opsetid("", opset)] + if use_qonnx: + opset_imports.append(oh.make_opsetid("qonnx.custom_op.general", 1)) + model_proto = oh.make_model(graph, opset_imports=opset_imports) + model_proto.ir_version = 6 + + # ONNX opset >= 9: initializers are implicit constants and must NOT appear in + # graph.input — otherwise tools treat weight tensors as runtime inputs. + # Some onnx library versions add them automatically for backward compatibility; + # strip them here so only the actual data inputs remain. + _init_names = {t.name for t in model_proto.graph.initializer} + _data_inputs = [vi for vi in model_proto.graph.input if vi.name not in _init_names] + del model_proto.graph.input[:] + model_proto.graph.input.extend(_data_inputs) + + onnx.checker.check_model(model_proto) + onnx.save(model_proto, output_path) + fmt = "QONNX" if use_qonnx else "ONNX (QDQ)" + logging.info("Saved %s Keras model → %s", fmt, output_path) + return model_proto + + +# --------------------------------------------------------------------------- +# usage example +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + import pquant + from pquant import apply_final_compression + + cfg = pquant.pdp_config() + + inp = keras.Input(shape=(3, 32, 32)) + x = PQConv2d(cfg, filters=16, kernel_size=3, padding="same")(inp) + x = PQBatchNormalization(cfg)(x) + x = keras.layers.ReLU()(x) + x = PQConv2d(cfg, filters=32, kernel_size=3, padding="same")(x) + x = keras.layers.ReLU()(x) + x = keras.layers.Flatten()(x) + x = PQDense(cfg, units=10)(x) + model = keras.Model(inp, x) + + apply_final_compression(model) + + convert_to_onnx(model, input_shape=(3, 32, 32), output_path="model_keras.onnx") + + import onnxruntime as ort + + sess = ort.InferenceSession("model_keras.onnx") + out = sess.run(None, {"input": np.random.randn(2, 3, 32, 32).astype(np.float32)}) + print("Output shape:", out[0].shape) # noqa: T201 diff --git a/src/pquant/core/keras/layers.py b/src/pquant/core/keras/layers.py index aabf69f..e602cb5 100644 --- a/src/pquant/core/keras/layers.py +++ b/src/pquant/core/keras/layers.py @@ -1,4 +1,5 @@ -from typing import Tuple, TypeVar +from math import prod +from typing import TypeVar import keras from keras import constraints, initializers, ops, regularizers @@ -23,13 +24,37 @@ ) from pquant.core.hyperparameter_optimization import PQConfig -from pquant.core.keras.activations import PQActivation +from pquant.core.keras.activations import PQActivation, PQSoftmax from pquant.core.keras.quantizer import Quantizer -from pquant.core.utils import get_pruning_layer +from pquant.core.keras.utils import get_pruning_layer T = TypeVar("T") +def resolve_data_quant_bits(quant_bits, config): + """Return (k, i, f) from an explicit tuple, or the config's data-lane defaults.""" + if quant_bits is not None: + return quant_bits + parameters = config.quantization_parameters + return ( + parameters.default_data_keep_negatives, + parameters.default_data_integer_bits, + parameters.default_data_fractional_bits, + ) + + +def resolve_weight_quant_bits(quant_bits, config): + """Return (k, i, f) from an explicit tuple, or the config's weight defaults.""" + if quant_bits is not None: + return quant_bits + parameters = config.quantization_parameters + return ( + parameters.default_weight_keep_negatives, + parameters.default_weight_integer_bits, + parameters.default_weight_fractional_bits, + ) + + @keras.saving.register_keras_serializable(package="PQuantML") class PQWeightBiasBase(keras.layers.Layer): def __init__( @@ -38,10 +63,14 @@ def __init__( layer_type, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, enable_pruning=None, *args, **kwargs, @@ -49,32 +78,10 @@ def __init__( super().__init__(**kwargs) if isinstance(config, dict): config = PQConfig.load_from_config(config) - if in_quant_bits is not None: - self.k_input, self.i_input, self.f_input = in_quant_bits - else: - self.k_input = config.quantization_parameters.default_data_keep_negatives - self.i_input = config.quantization_parameters.default_data_integer_bits - self.f_input = config.quantization_parameters.default_data_fractional_bits - - if weight_quant_bits is not None: - self.k_weight, self.i_weight, self.f_weight = weight_quant_bits - else: - self.k_weight = config.quantization_parameters.default_weight_keep_negatives - self.i_weight = config.quantization_parameters.default_weight_integer_bits - self.f_weight = config.quantization_parameters.default_weight_fractional_bits - if bias_quant_bits is not None: - self.k_bias, self.i_bias, self.f_bias = bias_quant_bits - else: - self.k_bias = config.quantization_parameters.default_weight_keep_negatives - self.i_bias = config.quantization_parameters.default_weight_integer_bits - self.f_bias = config.quantization_parameters.default_weight_fractional_bits - - if out_quant_bits is not None: - self.k_output, self.i_output, self.f_output = out_quant_bits - else: - self.k_output = config.quantization_parameters.default_data_keep_negatives - self.i_output = config.quantization_parameters.default_data_integer_bits - self.f_output = config.quantization_parameters.default_data_fractional_bits + self.k_input, self.i_input, self.f_input = resolve_data_quant_bits(in_quant_bits, config) + self.k_weight, self.i_weight, self.f_weight = resolve_weight_quant_bits(weight_quant_bits, config) + self.k_bias, self.i_bias, self.f_bias = resolve_weight_quant_bits(bias_quant_bits, config) + self.k_output, self.i_output, self.f_output = resolve_data_quant_bits(out_quant_bits, config) self.layer_type = layer_type self.pruning_layer = get_pruning_layer(config=config, layer_type=self.layer_type) @@ -86,6 +93,10 @@ def __init__( self.weight_quant_bits = weight_quant_bits self.bias_quant_bits = bias_quant_bits self.out_quant_bits = out_quant_bits + self.weight_quant_granularity = weight_quant_granularity + self.in_quant_granularity = in_quant_granularity + self.bias_quant_granularity = bias_quant_granularity + self.out_quant_granularity = out_quant_granularity self.pruning_first = config.training_parameters.pruning_first self.enable_quantization = config.quantization_parameters.enable_quantization self.round_mode = config.quantization_parameters.round_mode @@ -96,6 +107,7 @@ def __init__( self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress self.hgq_gamma = config.quantization_parameters.hgq_gamma self.granularity = config.quantization_parameters.granularity + self.dynamic_data = config.quantization_parameters.dynamic_data_quantization self.final_compression_done = False self.built = False self.parallelization_factor = -1 @@ -105,52 +117,60 @@ def __init__( self._is_finetuning = False self.config = config + weight_granularity = weight_quant_granularity if weight_quant_granularity is not None else self.granularity + bias_granularity = bias_quant_granularity if bias_quant_granularity is not None else self.granularity + in_granularity = in_quant_granularity if in_quant_granularity is not None else self.granularity + out_granularity = out_quant_granularity if out_quant_granularity is not None else self.granularity + self.weight_quantizer = Quantizer( - k=ops.convert_to_tensor(self.k_weight), - i=ops.convert_to_tensor(self.i_weight), - f=ops.convert_to_tensor(self.f_weight), + k=self.k_weight, + i=self.i_weight, + f=self.f_weight, overflow=self.overflow_mode_parameters, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=False, - granularity=self.granularity, + granularity=weight_granularity, hgq_gamma=self.hgq_gamma, place="weight", ) - - # if self.use_bias: self.bias_quantizer = Quantizer( - k=ops.convert_to_tensor(self.k_bias), - i=ops.convert_to_tensor(self.i_bias), - f=ops.convert_to_tensor(self.f_bias), + k=self.k_bias, + i=self.i_bias, + f=self.f_bias, overflow=self.overflow_mode_parameters, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=False, + granularity=bias_granularity, hgq_gamma=self.hgq_gamma, place="bias", ) self.input_quantizer = Quantizer( - k=ops.convert_to_tensor(self.k_input), - i=ops.convert_to_tensor(self.i_input), - f=ops.convert_to_tensor(self.f_input), + k=self.k_input, + i=self.i_input, + f=self.f_input, overflow=self.overflow_mode_data, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=in_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, ) self.output_quantizer = Quantizer( - k=ops.convert_to_tensor(self.k_output), - i=ops.convert_to_tensor(self.i_output), - f=ops.convert_to_tensor(self.f_output), + k=self.k_output, + i=self.i_output, + f=self.f_output, overflow=self.overflow_mode_data, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=out_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, ) def set_enable_pruning(self, enable_pruning): @@ -170,7 +190,7 @@ def get_output_quantization_bits(self): def build(self, input_shape): self.input_shape = (1,) + tuple(input_shape[1:]) - self.n_parallel = ops.prod(input_shape[1:-1]) + self.n_parallel = int(prod(input_shape[1:-1])) self.parallelization_factor = self.parallelization_factor if self.parallelization_factor > 0 else self.n_parallel self.is_pretraining = self.add_weight( shape=(), @@ -188,8 +208,60 @@ def build(self, input_shape): ) super().build(input_shape=input_shape) + def _build_quantizers(self, input_shape): + """Build the quantizer lanes in use that are not built yet. The output quantizer is + only created when the output is actually quantized (nothing reads it otherwise).""" + output_shape = self.compute_output_shape(input_shape) + if not self.input_quantizer.built: + self.input_quantizer.build(input_shape) + if not self.weight_quantizer.built: + self.weight_quantizer.build(self._kernel.shape) + if self.use_bias and not self.bias_quantizer.built: + self.bias_quantizer.build(self._bias.shape) + if self.quantize_output and not self.output_quantizer.built: + self.output_quantizer.build(output_shape) + + def _build_pruning_layer(self): + if self.enable_pruning and self.pruning_layer is not None and not self.pruning_layer.built: + pruning_shape = tuple(self._kernel.shape[i] for i in self.weight_transpose) + self.pruning_layer.build(pruning_shape) + + @property + def kernel(self): + if self.final_compression_done: + return self._kernel + if self.pruning_first: + weight = self._prune(self._kernel) + if self.enable_quantization: + weight = self.weight_quantizer(weight) + return weight + weight = self._kernel + if self.enable_quantization: + weight = self.weight_quantizer(weight) + return self._prune(weight) + + @kernel.setter + def kernel(self, kernel): + self._kernel = kernel + + @property + def bias(self): + if self.final_compression_done or self._bias is None: + return self._bias + bias = self._bias + if self.enable_quantization: + bias = self.bias_quantizer(self._bias) + return bias + + @bias.setter + def bias(self, bias): + self._bias = bias + def apply_final_compression(self): - pass + self._kernel.assign(self.kernel) + if self._bias is not None: + self._bias.assign(self.bias) + self.final_compression_done = True def save_own_variables(self, store): if not self.built: @@ -224,16 +296,67 @@ def pre_finetune_function(self): self._is_finetuning = True if hasattr(self, "is_finetuning"): self.is_finetuning.assign(1.0) + if self.pruning_layer is not None: + self.pruning_layer.pre_finetune_function() + + def pre_epoch_function(self, epoch, total_epochs): + if self.enable_pruning: + self.pruning_layer.pre_epoch_function(epoch, total_epochs) - def save_weights(self): + def post_epoch_function(self, epoch, total_epochs, **kwargs): + if self.enable_pruning: + self.pruning_layer.post_epoch_function(epoch, total_epochs, **kwargs) + self._update_pruning_mask() + + def post_round_function(self): + self.pruning_layer.post_round_function() + + def _update_pruning_mask(self): + if self.enable_pruning and hasattr(self.pruning_layer, "update_mask"): + kernel = self._handle_transpose(self._kernel, self.weight_transpose, True) + self.pruning_layer.update_mask(kernel) + + def _save_weights(self): self.init_weight = ops.copy(self._kernel) - def rewind_weights(self): + def _rewind_weights(self): self._kernel.assign(self.init_weight) def ebops(self): return 0.0 + def _masked_weight_bits(self, bw_ker): + """Zero the bit counts of weights that are pruned away or below the quantization step size.""" + mask = self._handle_transpose(self.pruning_layer.get_hard_mask(), self.weight_transpose_back, do_transpose=True) + _, _, f = self.get_weight_quantization_bits() + quantization_step_size = 2 ** (-f - 1) + step_size_mask = ops.cast(ops.abs(self._kernel) > quantization_step_size, self._kernel.dtype) + return bw_ker * mask * step_size_mask + + def _bias_ebops(self): + size = ops.cast(ops.prod(self.input_shape), self.dtype) + bw_bias = self.bias_quantizer.get_total_bits(ops.shape(self._bias)) + return ops.mean(bw_bias) * size + + def _conv_ebops(self, conv_bits, rank, include_mask): + bw_inp = self.input_quantizer.get_total_bits(self.input_shape) + bw_ker = self.weight_quantizer.get_total_bits(ops.shape(self._kernel)) + if include_mask: + bw_ker = self._masked_weight_bits(bw_ker) + if self.parallelization_factor < 0: + ebops = ops.sum(conv_bits(bw_inp, bw_ker)) + else: + if self.do_transpose_data: # channels_last + reduce_axis_input = tuple(range(rank + 1)) + else: + reduce_axis_input = (0,) + tuple(range(2, rank + 2)) + bw_inp = ops.max(bw_inp, axis=reduce_axis_input) + bw_ker = ops.sum(bw_ker, axis=tuple(range(rank))) + ebops = ops.sum(bw_inp[:, None] * bw_ker) + if self.use_bias: + ebops += self._bias_ebops() + return ebops + def hgq_loss(self): if not self.use_hgq: return ops.convert_to_tensor(0.0) @@ -248,39 +371,39 @@ def hgq_loss(self): loss += self.output_quantizer.hgq_loss() return ops.where(ops.cast(self.is_pretraining, "bool"), ops.zeros_like(loss), loss) - def handle_transpose(self, x, transpose, do_transpose=False): + def _handle_transpose(self, x, transpose, do_transpose=False): if do_transpose: x = ops.transpose(x, transpose) return x - def prune(self, weight): + def _prune(self, weight): if self.enable_pruning: - weight = self.handle_transpose(weight, self.weight_transpose, True) + weight = self._handle_transpose(weight, self.weight_transpose, True) weight = self.pruning_layer(weight) - weight = self.handle_transpose(weight, self.weight_transpose_back, True) + weight = self._handle_transpose(weight, self.weight_transpose_back, True) return weight def pre_forward(self, x, training): if self.quantize_input and self.enable_quantization: x = self.input_quantizer(x, training=training) if self.pruning_method == "wanda" and self.enable_pruning: - self.collect_input(x, self._kernel, training) + self._collect_input(x, self._kernel, training) return x - def post_forward(self, x, training): + def _post_forward(self, x, training): if self.quantize_output and self.enable_quantization: x = self.output_quantizer(x, training=training) if self.pruning_method == "activation_pruning" and self.enable_pruning: - self.collect_output(x, training) + self._collect_output(x, training) return x - def collect_input(self, x, weight, training): - collect_x = self.handle_transpose(x, self.data_transpose, self.do_transpose_data) - weight_channels_first = self.handle_transpose(weight, self.weight_transpose, True) + def _collect_input(self, x, weight, training): + collect_x = self._handle_transpose(x, self.data_transpose, self.do_transpose_data) + weight_channels_first = self._handle_transpose(weight, self.weight_transpose, True) self.pruning_layer.collect_input(collect_x, weight_channels_first, training) - def collect_output(self, x, training): - collect_x = self.handle_transpose(x, self.data_transpose, self.do_transpose_data) + def _collect_output(self, x, training): + collect_x = self._handle_transpose(x, self.data_transpose, self.do_transpose_data) self.pruning_layer.collect_output(collect_x, training) @classmethod @@ -312,6 +435,10 @@ def get_config(self): "weight_quant_bits": self.weight_quant_bits, "bias_quant_bits": self.bias_quant_bits, "out_quant_bits": self.out_quant_bits, + "weight_quant_granularity": self.weight_quant_granularity, + "in_quant_granularity": self.in_quant_granularity, + "bias_quant_granularity": self.bias_quant_granularity, + "out_quant_granularity": self.out_quant_granularity, "enable_pruning": self.enable_pruning, "final_compression_done": self.final_compression_done, } @@ -344,10 +471,14 @@ def __init__( bias: bool = True, device=None, dtype=None, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, enable_pruning=None, **kwargs, ): @@ -375,18 +506,19 @@ def __init__( weight_quant_bits=weight_quant_bits, bias_quant_bits=bias_quant_bits, out_quant_bits=out_quant_bits, + weight_quant_granularity=weight_quant_granularity, + in_quant_granularity=in_quant_granularity, + bias_quant_granularity=bias_quant_granularity, + out_quant_granularity=out_quant_granularity, enable_pruning=enable_pruning, **kwargs, ) self.depthwise_regularizer = depthwise_regularizer self.use_bias = use_bias - self.strides = strides - self.dilation_rate = dilation_rate self.weight_transpose = (2, 3, 0, 1) self.weight_transpose_back = (2, 3, 0, 1) self.data_transpose = (0, 3, 1, 2) self.do_transpose_data = self.data_format == "channels_last" - self._weight = None self._bias = None def build(self, input_shape): @@ -423,122 +555,30 @@ def build(self, input_shape): ) else: self._bias = None - if self.use_hgq: - self.input_quantizer.build(input_shape) - self.weight_quantizer.build(self._kernel.shape) - if self.use_bias: - self.bias_quantizer.build(self._bias.shape) - self.output_quantizer.build(self.compute_output_shape(input_shape)) - else: - if not self.input_quantizer.built: - self.input_quantizer.build(input_shape) - if not self.weight_quantizer.built: - self.weight_quantizer.build(self._kernel.shape) - if self.use_bias and not self.bias_quantizer.built: - self.bias_quantizer.build(self._bias.shape) - if self.quantize_output and not self.output_quantizer.built: - self.output_quantizer.build(self.compute_output_shape(input_shape)) - self.input_shape = (1,) + input_shape[1:] - if self.enable_pruning and self.pruning_layer is not None and not self.pruning_layer.built: - pruning_shape = tuple(self._kernel.shape[i] for i in self.weight_transpose) - self.pruning_layer.build(pruning_shape) - - @property - def kernel(self): - if self.final_compression_done: - return self._kernel - if self.pruning_first: - weight = self.prune(self._kernel) - if self.enable_quantization: - weight = self.weight_quantizer(weight) - return weight - else: - weight = self._kernel - if self.enable_quantization: - weight = self.weight_quantizer(weight) - return self.prune(weight) - - @kernel.setter - def kernel(self, kernel): - self._kernel = kernel - - @property - def bias(self): - if self.final_compression_done or self._bias is None: - return self._bias - bias = self._bias - if self.enable_quantization: - bias = self.bias_quantizer(self._bias) - return bias - - @bias.setter - def bias(self, bias): - self._bias = bias + self._build_quantizers(input_shape) + self._build_pruning_layer() def ebops(self, include_mask=False): - bw_inp = self.input_quantizer.get_total_bits(self.input_shape) - bw_ker = self.weight_quantizer.get_total_bits(ops.shape(self._kernel)) - if include_mask: - mask = self.handle_transpose(self.pruning_layer.get_hard_mask(), self.weight_transpose_back, do_transpose=True) - bw_ker = bw_ker * mask - _, _, f = self.get_weight_quantization_bits() - quantization_step_size = 2 ** (-f - 1) - step_size_mask = ops.cast((ops.abs(self._kernel) > quantization_step_size), self._kernel.dtype) - bw_ker = bw_ker * step_size_mask - if self.parallelization_factor < 0: - ebops = ops.sum( - ops.depthwise_conv( - bw_inp, - bw_ker, - strides=self.strides, - padding=self.padding, - data_format=None, - dilation_rate=self.dilation_rate, - ) + def conv_bits(bw_inp, bw_ker): + return ops.depthwise_conv( + bw_inp, + bw_ker, + strides=self.strides, + padding=self.padding, + data_format=None, + dilation_rate=self.dilation_rate, ) - else: - reduce_axis_kernel = tuple(range(0, 3)) - if self.data_format == "channels_last": # Is channels last - reduce_axis_input = reduce_axis_kernel - else: - reduce_axis_input = (0,) + tuple(range(2, 4)) - bw_inp = ops.max(bw_inp, axis=reduce_axis_input) - reduce_axis_kernel = tuple(range(0, 2)) - bw_ker = ops.sum(bw_ker, axis=reduce_axis_kernel) - ebops = ops.sum(bw_inp[:, None] * bw_ker) - if self.use_bias: - size = ops.cast(ops.prod(self.input_shape), self.dtype) - bw_bias = self.bias_quantizer.get_total_bits(ops.shape(self._bias)) - ebops += ops.mean(bw_bias) * size - return ebops + + return self._conv_ebops(conv_bits, rank=2, include_mask=include_mask) def call(self, x, training=None): x = self.pre_forward(x, training) x = super().call(x) - x = self.post_forward(x, training) + x = self._post_forward(x, training) if self.use_hgq and self.enable_quantization: self.add_loss(self.hgq_loss()) return x - # Is it supposed to be like this? - def apply_final_compression(self): - self._kernel.assign(self.kernel) - if self._bias is not None: - self._bias.assign(self.bias) - self.final_compression_done = True - - def extra_repr(self) -> str: - """ - Return the extra representation of the module. - """ - return ( - f"in_features={self.in_features} " - f"out_features={self.out_features} " - f"bias={self._bias is not None} " - f"quantize_input={self.quantize_input} " - f"quantize_output={self.quantize_output} " - ) - def _normalize_tuple(value, n): if isinstance(value, int): @@ -569,10 +609,14 @@ def __init__( activity_regularizer=None, kernel_constraint=None, bias_constraint=None, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, enable_pruning=None, **kwargs, ): @@ -585,6 +629,10 @@ def __init__( weight_quant_bits=weight_quant_bits, bias_quant_bits=bias_quant_bits, out_quant_bits=out_quant_bits, + weight_quant_granularity=weight_quant_granularity, + in_quant_granularity=in_quant_granularity, + bias_quant_granularity=bias_quant_granularity, + out_quant_granularity=out_quant_granularity, enable_pruning=enable_pruning, activity_regularizer=activity_regularizer, **kwargs, @@ -631,90 +679,21 @@ def build(self, input_shape): else: self._bias = None super().build(input_shape) - if self.use_hgq: - self.input_quantizer.build(input_shape) - self.weight_quantizer.build(self._kernel.shape) - if self.use_bias: - self.bias_quantizer.build(self._bias.shape) - self.output_quantizer.build(self.compute_output_shape(input_shape)) - else: - if not self.input_quantizer.built: - self.input_quantizer.build(input_shape) - if not self.weight_quantizer.built: - self.weight_quantizer.build(self._kernel.shape) - if self.use_bias and not self.bias_quantizer.built: - self.bias_quantizer.build(self._bias.shape) - if self.quantize_output and not self.output_quantizer.built: - self.output_quantizer.build(self.compute_output_shape(input_shape)) - if self.enable_pruning and self.pruning_layer is not None and not self.pruning_layer.built: - pruning_shape = tuple(self._kernel.shape[i] for i in self.weight_transpose) - self.pruning_layer.build(pruning_shape) - - @property - def kernel(self): - if self.final_compression_done: - return self._kernel - if self.pruning_first: - weight = self.prune(self._kernel) - if self.enable_quantization: - weight = self.weight_quantizer(weight) - return weight - else: - weight = self._kernel - if self.enable_quantization: - weight = self.weight_quantizer(weight) - return self.prune(weight) - - @property - def bias(self): - if self.final_compression_done or self._bias is None: - return self._bias - bias = self._bias - if self.enable_quantization: - bias = self.bias_quantizer(self._bias) - return bias - - @bias.setter - def bias(self, bias): - self._bias = bias + self._build_quantizers(input_shape) + self._build_pruning_layer() def ebops(self, include_mask=False): - bw_inp = self.input_quantizer.get_total_bits(self.input_shape) - bw_ker = self.weight_quantizer.get_total_bits(ops.shape(self._kernel)) - if include_mask: - mask = self.handle_transpose(self.pruning_layer.get_hard_mask(), self.weight_transpose_back, do_transpose=True) - bw_ker = bw_ker * mask - _, _, f = self.get_weight_quantization_bits() - quantization_step_size = 2 ** (-f - 1) - step_size_mask = ops.cast((ops.abs(self._kernel) > quantization_step_size), self._kernel.dtype) - bw_ker = bw_ker * step_size_mask - if self.parallelization_factor < 0: - ebops = ops.sum( - ops.conv( - bw_inp, - bw_ker, - strides=self.strides, - padding=self.padding, - data_format=None, - dilation_rate=self.dilation_rate, - ) + def conv_bits(bw_inp, bw_ker): + return ops.conv( + bw_inp, + bw_ker, + strides=self.strides, + padding=self.padding, + data_format=None, + dilation_rate=self.dilation_rate, ) - else: - reduce_axis_kernel = tuple(range(0, 3)) - if self.do_transpose_data: # Is channels last - reduce_axis_input = reduce_axis_kernel - else: - reduce_axis_input = (0,) + tuple(range(2, 4)) - bw_inp = ops.max(bw_inp, axis=reduce_axis_input) - reduce_axis_kernel = tuple(range(0, 2)) - bw_ker = ops.sum(bw_ker, axis=reduce_axis_kernel) - ebops = ops.sum(bw_inp[:, None] * bw_ker) - if self.use_bias: - size = ops.cast(ops.prod(self.input_shape), self.dtype) - bw_bias = self.bias_quantizer.get_total_bits(ops.shape(self._bias)) - ebops += ops.mean(bw_bias) * size - return ebops + return self._conv_ebops(conv_bits, rank=2, include_mask=include_mask) def compute_output_shape(self, input_shape): return compute_conv_output_shape( @@ -727,12 +706,6 @@ def compute_output_shape(self, input_shape): dilation_rate=self.dilation_rate, ) - def apply_final_compression(self): - self._kernel.assign(self.kernel) - if self._bias is not None: - self._bias.assign(self.bias) - self.final_compression_done = True - def call(self, x, training=None): x = self.pre_forward(x, training) x = ops.conv( @@ -746,7 +719,7 @@ def call(self, x, training=None): if self.use_bias: bias_shape = (1, 1, 1, self.filters) if self.data_format == "channels_last" else (1, self.filters, 1, 1) x = x + ops.reshape(self.bias, bias_shape) - x = self.post_forward(x, training) + x = self._post_forward(x, training) if self.use_hgq and self.enable_quantization: self.add_loss(self.hgq_loss()) return x @@ -841,18 +814,61 @@ def __init__( ) self.do_transpose_data = data_format == "channels_last" - def build(self, input_shape): - super().build(input_shape) - def apply_final_compression(self): self.depthwise_conv.apply_final_compression() self.pointwise_conv.apply_final_compression() + def post_pre_train_function(self): + self.depthwise_conv.post_pre_train_function() + self.pointwise_conv.post_pre_train_function() + + def pre_finetune_function(self): + self.depthwise_conv.pre_finetune_function() + self.pointwise_conv.pre_finetune_function() + + def pre_epoch_function(self, epoch, total_epochs): + self.depthwise_conv.pre_epoch_function(epoch, total_epochs) + self.pointwise_conv.pre_epoch_function(epoch, total_epochs) + + def post_epoch_function(self, epoch, total_epochs, **kwargs): + self.depthwise_conv.post_epoch_function(epoch, total_epochs, **kwargs) + self.pointwise_conv.post_epoch_function(epoch, total_epochs, **kwargs) + + def post_round_function(self): + self.depthwise_conv.post_round_function() + self.pointwise_conv.post_round_function() + + def _save_weights(self): + self.depthwise_conv._save_weights() + self.pointwise_conv._save_weights() + + def _rewind_weights(self): + self.depthwise_conv._rewind_weights() + self.pointwise_conv._rewind_weights() + + def build(self, input_shape): + self.depthwise_conv.build(input_shape) + intermediate_shape = self.depthwise_conv.compute_output_shape(input_shape) + self.pointwise_conv.build(intermediate_shape) + super().build(input_shape) + + def compute_output_shape(self, input_shape): + intermediate_shape = self.depthwise_conv.compute_output_shape(input_shape) + return self.pointwise_conv.compute_output_shape(intermediate_shape) + def call(self, x, training=None): x = self.depthwise_conv(x, training=training) x = self.pointwise_conv(x, training=training) return x + @classmethod + def from_config(cls, config): + final_compression_done = config.pop("final_compression_done", False) + instance = cls(**config) + instance.depthwise_conv.final_compression_done = final_compression_done + instance.pointwise_conv.final_compression_done = final_compression_done + return instance + def get_config(self): config = super().get_config() config.update( @@ -868,6 +884,7 @@ def get_config(self): "use_bias": self.pointwise_conv.use_bias, "quantize_input": self.depthwise_conv.quantize_input, "quantize_output": self.pointwise_conv.quantize_output, + "final_compression_done": self.depthwise_conv.final_compression_done, } ) return config @@ -882,10 +899,14 @@ def __init__( kernel_size, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, enable_pruning=None, strides=1, padding="valid", @@ -912,6 +933,10 @@ def __init__( weight_quant_bits=weight_quant_bits, bias_quant_bits=bias_quant_bits, out_quant_bits=out_quant_bits, + weight_quant_granularity=weight_quant_granularity, + in_quant_granularity=in_quant_granularity, + bias_quant_granularity=bias_quant_granularity, + out_quant_granularity=out_quant_granularity, enable_pruning=enable_pruning, activity_regularizer=activity_regularizer, **kwargs, @@ -958,89 +983,21 @@ def build(self, input_shape): else: self._bias = None super().build(input_shape) - if self.use_hgq: - self.input_quantizer.build(input_shape) - self.weight_quantizer.build(self._kernel.shape) - if self.use_bias: - self.bias_quantizer.build(self._bias.shape) - self.output_quantizer.build(self.compute_output_shape(input_shape)) - else: - if not self.input_quantizer.built: - self.input_quantizer.build(input_shape) - if not self.weight_quantizer.built: - self.weight_quantizer.build(self._kernel.shape) - if self.use_bias and not self.bias_quantizer.built: - self.bias_quantizer.build(self._bias.shape) - if self.quantize_output and not self.output_quantizer.built: - self.output_quantizer.build(self.compute_output_shape(input_shape)) - if self.enable_pruning and self.pruning_layer is not None and not self.pruning_layer.built: - pruning_shape = tuple(self._kernel.shape[i] for i in self.weight_transpose) - self.pruning_layer.build(pruning_shape) - - @property - def kernel(self): - if self.final_compression_done: - return self._kernel - if self.pruning_first: - weight = self.prune(self._kernel) - if self.enable_quantization: - weight = self.weight_quantizer(weight) - return weight - else: - weight = self._kernel - if self.enable_quantization: - weight = self.weight_quantizer(weight) - return self.prune(weight) - - @property - def bias(self): - if self.final_compression_done or self._bias is None: - return self._bias - bias = self._bias - if self.enable_quantization: - bias = self.bias_quantizer(self._bias) - return bias - - @bias.setter - def bias(self, bias): - self._bias = bias - - def ebops(self, include_mask=False): - bw_inp = self.input_quantizer.get_total_bits(self.input_shape) - bw_ker = self.weight_quantizer.get_total_bits(ops.shape(self._kernel)) - if include_mask: - mask = self.handle_transpose(self.pruning_layer.get_hard_mask(), self.weight_transpose_back, do_transpose=True) - bw_ker = bw_ker * mask - _, _, f = self.get_weight_quantization_bits() - quantization_step_size = 2 ** (-f - 1) - step_size_mask = ops.cast((ops.abs(self._kernel) > quantization_step_size), self._kernel.dtype) - bw_ker = bw_ker * step_size_mask - if self.parallelization_factor < 0: - ebops = ops.sum( - ops.conv( - bw_inp, - bw_ker, - strides=self.strides, - padding=self.padding, - data_format=None, - dilation_rate=self.dilation_rate, - ) - ) - else: - reduce_axis_kernel = tuple(range(0, 2)) - if self.do_transpose_data: # Is channels last - reduce_axis_input = reduce_axis_kernel - else: - reduce_axis_input = (0,) + tuple(range(2, 3)) - bw_inp = ops.max(bw_inp, axis=reduce_axis_input) - reduce_axis_kernel = tuple(range(0, 1)) - bw_ker = ops.sum(bw_ker, axis=reduce_axis_kernel) - ebops = ops.sum(bw_inp[:, None] * bw_ker) - if self.use_bias: - size = ops.cast(ops.prod(self.input_shape), self.dtype) - bw_bias = self.bias_quantizer.get_total_bits(ops.shape(self._bias)) - ebops += ops.mean(bw_bias) * size - return ebops + self._build_quantizers(input_shape) + self._build_pruning_layer() + + def ebops(self, include_mask=False): + def conv_bits(bw_inp, bw_ker): + return ops.conv( + bw_inp, + bw_ker, + strides=self.strides, + padding=self.padding, + data_format=None, + dilation_rate=self.dilation_rate, + ) + + return self._conv_ebops(conv_bits, rank=1, include_mask=include_mask) def compute_output_shape(self, input_shape): return compute_conv_output_shape( @@ -1053,12 +1010,6 @@ def compute_output_shape(self, input_shape): dilation_rate=self.dilation_rate, ) - def apply_final_compression(self): - self._kernel.assign(self.kernel) - if self._bias is not None: - self._bias.assign(self.bias) - self.final_compression_done = True - def call(self, x, training=None): x = self.pre_forward(x, training) x = ops.conv( @@ -1072,7 +1023,7 @@ def call(self, x, training=None): if self.use_bias: bias_shape = (1, 1, self.filters) if self.data_format == "channels_last" else (1, self.filters, 1) x = x + ops.reshape(self.bias, bias_shape) - x = self.post_forward(x, training) + x = self._post_forward(x, training) if self.use_hgq and self.enable_quantization: self.add_loss(self.hgq_loss()) return x @@ -1108,10 +1059,14 @@ def __init__( units, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, enable_pruning=None, use_bias=True, kernel_initializer="glorot_uniform", @@ -1131,6 +1086,10 @@ def __init__( weight_quant_bits=weight_quant_bits, bias_quant_bits=bias_quant_bits, out_quant_bits=out_quant_bits, + weight_quant_granularity=weight_quant_granularity, + in_quant_granularity=in_quant_granularity, + bias_quant_granularity=bias_quant_granularity, + out_quant_granularity=out_quant_granularity, enable_pruning=enable_pruning, **kwargs, ) @@ -1147,7 +1106,6 @@ def __init__( self.kernel_constraint = constraints.get(kernel_constraint) self.bias_constraint = constraints.get(bias_constraint) self.input_spec = InputSpec(min_ndim=2) - self._ebops = self.add_variable(shape=(), initializer="zeros", trainable=False) def build(self, input_shape): input_dim = input_shape[-1] @@ -1169,66 +1127,20 @@ def build(self, input_shape): else: self._bias = None super().build(input_shape) - if not self.input_quantizer.built: - self.input_quantizer.build(input_shape) - if not self.weight_quantizer.built: - self.weight_quantizer.build(self._kernel.shape) - if self.use_bias and not self.bias_quantizer.built: - self.bias_quantizer.build(self._bias.shape) - if self.quantize_output and not self.output_quantizer.built: - output_shape = input_shape[:-1] + (self.units,) - self.output_quantizer.build(output_shape) - if self.enable_pruning and self.pruning_layer is not None and not self.pruning_layer.built: - pruning_shape = tuple(self._kernel.shape[i] for i in self.weight_transpose) - self.pruning_layer.build(pruning_shape) - - @property - def kernel(self): - if self.final_compression_done: - return self._kernel - if self.pruning_first: - weight = self.prune(self._kernel) - if self.enable_quantization: - weight = self.weight_quantizer(weight) - return weight - else: - weight = self._kernel - if self.enable_quantization: - weight = self.weight_quantizer(weight) - return self.prune(weight) - - @property - def bias(self): - if self.final_compression_done or self._bias is None: - return self._bias - bias = self._bias - if self.enable_quantization: - bias = self.bias_quantizer(self._bias) - return bias + self._build_quantizers(input_shape) + self._build_pruning_layer() def ebops(self, include_mask=False): bw_inp = self.input_quantizer.get_total_bits(self.input_shape) bw_ker = self.weight_quantizer.get_total_bits(ops.shape(self._kernel)) if include_mask: - mask = self.handle_transpose(self.pruning_layer.get_hard_mask(), self.weight_transpose_back, do_transpose=True) - bw_ker = bw_ker * mask - _, _, f = self.get_weight_quantization_bits() - quantization_step_size = 2 ** (-f - 1) - step_size_mask = ops.cast((ops.abs(self._kernel) > quantization_step_size), self._kernel.dtype) - bw_ker = bw_ker * step_size_mask + bw_ker = self._masked_weight_bits(bw_ker) ebops = ops.sum(ops.matmul(bw_inp, bw_ker)) if self.use_bias: bw_bias = self.bias_quantizer.get_total_bits(ops.shape(self._bias)) size = ops.cast(ops.prod(self.input_shape[:-1]) * self.units, self.dtype) ebops += ops.mean(bw_bias) * size - ebops = ebops * self.parallelization_factor / self.n_parallel - return ebops - - def apply_final_compression(self): - self._kernel.assign(self.kernel) - if self._bias is not None: - self._bias.assign(self.bias) - self.final_compression_done = True + return ebops * self.parallelization_factor / self.n_parallel def compute_output_shape(self, input_shape): output_shape = list(input_shape) @@ -1236,13 +1148,11 @@ def compute_output_shape(self, input_shape): return tuple(output_shape) def call(self, x, training=None): - self.training = training x = self.pre_forward(x, training) x = ops.matmul(x, self.kernel) - bias = self.bias if self.use_bias: - x = ops.add(x, bias) - x = self.post_forward(x, training) + x = ops.add(x, self.bias) + x = self._post_forward(x, training) if self.use_hgq: self.add_loss(self.hgq_loss()) return x @@ -1274,25 +1184,28 @@ def __init__( synchronized=False, quantize_input=True, quantize_parameters=True, + in_quant_granularity=None, + weight_quant_granularity=None, + bias_quant_granularity=None, **kwargs, ): if isinstance(config, dict): config = PQConfig.load_from_config(config) super().__init__( - axis, - momentum, - epsilon, - center, - scale, - beta_initializer, - gamma_initializer, - moving_mean_initializer, - moving_variance_initializer, - beta_regularizer, - gamma_regularizer, - beta_constraint, - gamma_constraint, - synchronized, + axis=axis, + momentum=momentum, + epsilon=epsilon, + center=center, + scale=scale, + beta_initializer=beta_initializer, + gamma_initializer=gamma_initializer, + moving_mean_initializer=moving_mean_initializer, + moving_variance_initializer=moving_variance_initializer, + beta_regularizer=beta_regularizer, + gamma_regularizer=gamma_regularizer, + beta_constraint=beta_constraint, + gamma_constraint=gamma_constraint, + synchronized=synchronized, **kwargs, ) self.overflow_mode_parameters = config.quantization_parameters.overflow_mode_parameters @@ -1307,6 +1220,10 @@ def __init__( self.quantize_input = quantize_input self.quantize_parameters = quantize_parameters self.granularity = config.quantization_parameters.granularity + self.in_quant_granularity = in_quant_granularity + self.weight_quant_granularity = weight_quant_granularity + self.bias_quant_granularity = bias_quant_granularity + self.dynamic_data = config.quantization_parameters.dynamic_data_quantization self.config = config self.f_weight = self.f_bias = ops.convert_to_tensor(config.quantization_parameters.default_weight_fractional_bits) self.i_weight = self.i_bias = ops.convert_to_tensor(config.quantization_parameters.default_weight_integer_bits) @@ -1324,6 +1241,9 @@ def build(self, input_shape): trainable=False, dtype="float32", ) + in_granularity = self.in_quant_granularity if self.in_quant_granularity is not None else self.granularity + weight_granularity = self.weight_quant_granularity if self.weight_quant_granularity is not None else self.granularity + bias_granularity = self.bias_quant_granularity if self.bias_quant_granularity is not None else self.granularity self.input_quantizer = Quantizer( k=1.0, i=self.i_input, @@ -1332,27 +1252,31 @@ def build(self, input_shape): round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=in_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, ) self.weight_quantizer = Quantizer( k=1.0, i=self.i_weight, f=self.f_weight, - round_mode=self.round_mode, overflow=self.overflow_mode_parameters, - is_data=False, + round_mode=self.round_mode, is_heterogeneous=self.use_hgq, + is_data=False, + granularity=weight_granularity, place="weight", ) self.bias_quantizer = Quantizer( k=1.0, i=self.i_bias, f=self.f_bias, - round_mode=self.round_mode, overflow=self.overflow_mode_parameters, - is_data=False, + round_mode=self.round_mode, is_heterogeneous=self.use_hgq, + is_data=False, + granularity=bias_granularity, place="bias", ) self.input_quantizer.build(input_shape) @@ -1364,12 +1288,12 @@ def build(self, input_shape): self.input_shape = (1,) + tuple(input_shape[1:]) def apply_final_compression(self): - self.final_compression_done = True if self.enable_quantization and self.quantize_parameters: if self.gamma is not None: self.gamma.assign(self.weight_quantizer(self.gamma)) if self.beta is not None: self.beta.assign(self.bias_quantizer(self.beta)) + self.final_compression_done = True def ebops(self): bw_inp = self.input_quantizer.get_total_bits(self.input_shape) @@ -1443,7 +1367,8 @@ def call(self, inputs, training=None, mask=None): scale=gamma, epsilon=self.epsilon, ) - self.add_loss(self.hgq_loss()) + if self.use_hgq and self.enable_quantization: + self.add_loss(self.hgq_loss()) return ops.cast(outputs, self.compute_dtype) def get_input_quantization_bits(self): @@ -1474,6 +1399,9 @@ def get_config(self): "config": self.config.get_dict(), "quantize_input": self.quantize_input, "quantize_parameters": self.quantize_parameters, + "in_quant_granularity": self.in_quant_granularity, + "weight_quant_granularity": self.weight_quant_granularity, + "bias_quant_granularity": self.bias_quant_granularity, "final_compression_done": self.final_compression_done, } ) @@ -1487,8 +1415,10 @@ def __init__( config, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, **kwargs, ): @@ -1498,20 +1428,10 @@ def __init__( self.in_quant_bits = in_quant_bits self.out_quant_bits = out_quant_bits - - if in_quant_bits is not None: - self.k_input, self.i_input, self.f_input = in_quant_bits - else: - self.k_input = config.quantization_parameters.default_data_keep_negatives - self.i_input = config.quantization_parameters.default_data_integer_bits - self.f_input = config.quantization_parameters.default_data_fractional_bits - - if out_quant_bits is not None: - self.k_output, self.i_output, self.f_output = out_quant_bits - else: - self.k_output = config.quantization_parameters.default_data_keep_negatives - self.i_output = config.quantization_parameters.default_data_integer_bits - self.f_output = config.quantization_parameters.default_data_fractional_bits + self.in_quant_granularity = in_quant_granularity + self.out_quant_granularity = out_quant_granularity + self.k_input, self.i_input, self.f_input = resolve_data_quant_bits(in_quant_bits, config) + self.k_output, self.i_output, self.f_output = resolve_data_quant_bits(out_quant_bits, config) self.overflow_mode_data = config.quantization_parameters.overflow_mode_data self.config = config self.round_mode = config.quantization_parameters.round_mode @@ -1521,6 +1441,7 @@ def __init__( self.hgq_gamma = config.quantization_parameters.hgq_gamma self.hgq_beta = config.quantization_parameters.hgq_beta self.hgq_heterogeneous = config.quantization_parameters.hgq_heterogeneous + self.dynamic_data = config.quantization_parameters.dynamic_data_quantization self._is_pretraining = True self.quantize_input = quantize_input self.quantize_output = quantize_output @@ -1541,6 +1462,9 @@ def build(self, input_shape): trainable=False, dtype="float32", ) + config_granularity = self.config.quantization_parameters.granularity + in_granularity = self.in_quant_granularity if self.in_quant_granularity is not None else config_granularity + out_granularity = self.out_quant_granularity if self.out_quant_granularity is not None else config_granularity self.input_quantizer = Quantizer( k=1.0, i=self.i_input, @@ -1549,8 +1473,10 @@ def build(self, input_shape): round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=in_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, ) self.output_quantizer = Quantizer( k=1.0, @@ -1560,8 +1486,10 @@ def build(self, input_shape): round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=out_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, ) self.input_quantizer.build(input_shape) self.output_quantizer.build(self.compute_output_shape(input_shape)) @@ -1582,12 +1510,12 @@ def compute_output_shape(self, input_shape): self.data_format, ) - def pre_pooling(self, x, training): + def _pre_pooling(self, x, training): if self.quantize_input and self.enable_quantization: x = self.input_quantizer(x, training=training) return x - def post_pooling(self, x, training): + def _post_pooling(self, x, training): if self.quantize_output and self.enable_quantization: x = self.output_quantizer(x, training=training) return x @@ -1615,6 +1543,8 @@ def get_config(self): "quantize_output": self.quantize_output, "in_quant_bits": self.in_quant_bits, "out_quant_bits": self.out_quant_bits, + "in_quant_granularity": self.in_quant_granularity, + "out_quant_granularity": self.out_quant_granularity, } ) return config @@ -1628,8 +1558,10 @@ def __init__( pool_size, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, strides=None, padding="valid", data_format=None, @@ -1647,13 +1579,15 @@ def __init__( quantize_output=quantize_output, in_quant_bits=in_quant_bits, out_quant_bits=out_quant_bits, + in_quant_granularity=in_quant_granularity, + out_quant_granularity=out_quant_granularity, **kwargs, ) def call(self, x, training=None): - x = self.pre_pooling(x, training) + x = self._pre_pooling(x, training) x = super().call(x) - x = self.post_pooling(x, training) + x = self._post_pooling(x, training) if self.use_hgq and self.enable_quantization: self.add_loss(self.hgq_loss()) return x @@ -1670,8 +1604,10 @@ def __init__( pool_size, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, strides=None, padding="valid", data_format=None, @@ -1689,12 +1625,14 @@ def __init__( quantize_output=quantize_output, in_quant_bits=in_quant_bits, out_quant_bits=out_quant_bits, + in_quant_granularity=in_quant_granularity, + out_quant_granularity=out_quant_granularity, ) def call(self, x, training=None): - x = self.pre_pooling(x, training) + x = self._pre_pooling(x, training) x = super().call(x) - x = self.post_pooling(x, training) + x = self._post_pooling(x, training) if self.use_hgq and self.enable_quantization: self.add_loss(self.hgq_loss()) return x @@ -1703,143 +1641,407 @@ def get_config(self): return super().get_config() +@keras.saving.register_keras_serializable(package="PQuantML") +class PQMultiheadAttention(keras.layers.Layer): + """Multi-head attention with quantization support. + + Uses separate PQDense projections for Q, K, V, and output, and computes + scaled dot-product attention manually. + + Args: + config: PQuant configuration object. + embed_dim: Total embedding dimension. + num_heads: Number of attention heads. + dropout: Dropout probability on attention weights. + bias: Whether to add bias to projection layers. + kdim: Key feature dimension (defaults to embed_dim). + vdim: Value feature dimension (defaults to embed_dim). + quantize_input: Whether to quantize the Q/K/V projection inputs (the MHA inputs). + quantize_output: Whether to quantize the output projection's output (the MHA output). + The q/k/v projection outputs and the out_proj input (the context) are always + quantized, mirroring HGQ's QMultiHeadAttention. + approximate_softmax: Placeholder for approximate softmax (currently uses standard softmax). + in_quant_bits: (k, i, f) bits for input quantization. + weight_quant_bits: (k, i, f) bits for weight quantization. + bias_quant_bits: (k, i, f) bits for bias quantization. + out_quant_bits: (k, i, f) bits for output quantization. + attn_quant_bits: (k, i, f) bits for the softmax output quantizer (the attention + weights). The scores and context need no dedicated quantizers: the softmax's + input quantizer and the output projection's input quantizer cover them. + + Call args: + inputs: A tuple (query, key, value) of tensors with shape (batch, seq, features), + or a single tensor for self-attention. + training: Python boolean indicating whether the layer should behave in training mode. + key_padding_mask: Boolean tensor of shape (batch, key_seq). True means the position + should be ignored. + attn_mask: Additive mask of shape (query_seq, key_seq) or + (batch, num_heads, query_seq, key_seq). + need_weights: If True, returns (output, attn_weights). If False, returns (output, None). + """ + + def __init__( + self, + config, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + bias: bool = True, + kdim: int = None, + vdim: int = None, + quantize_input: bool = True, + quantize_output: bool = False, + approximate_softmax: bool = False, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + attn_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, + param_quant_granularity=None, + **kwargs, + ): + super().__init__(**kwargs) + assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads" + + if isinstance(config, dict): + config = PQConfig.load_from_config(config) + + self.config = config + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.dropout_rate = dropout + self.use_bias = bias + self.kdim = kdim if kdim is not None else embed_dim + self.vdim = vdim if vdim is not None else embed_dim + self.approximate_softmax = approximate_softmax + self.scale = self.head_dim**-0.5 + self.enable_quantization = config.quantization_parameters.enable_quantization + self.use_hgq = config.quantization_parameters.use_high_granularity_quantization + self.hgq_beta = config.quantization_parameters.hgq_beta + self.is_pretraining = True + + self.in_quant_bits = in_quant_bits + self.weight_quant_bits = weight_quant_bits + self.bias_quant_bits = bias_quant_bits + self.out_quant_bits = out_quant_bits + self.attn_quant_bits = attn_quant_bits + + self.in_quant_granularity = in_quant_granularity + self.out_quant_granularity = out_quant_granularity + self.param_quant_granularity = param_quant_granularity + + self.softmax = PQSoftmax(config, -1, quantize_input=True, quantize_output=True, out_quant_bits=attn_quant_bits) + proj_kwargs = dict( + use_bias=bias, + in_quant_bits=in_quant_bits, + weight_quant_bits=weight_quant_bits, + bias_quant_bits=bias_quant_bits, + out_quant_bits=out_quant_bits, + weight_quant_granularity=param_quant_granularity, + bias_quant_granularity=param_quant_granularity, + ) + + qkv_kwargs = dict( + quantize_input=quantize_input, quantize_output=True, in_quant_granularity=in_quant_granularity, **proj_kwargs + ) + self.q_proj = PQDense(config, embed_dim, enable_pruning=False, **qkv_kwargs) + self.k_proj = PQDense(config, embed_dim, enable_pruning=False, **qkv_kwargs) + self.v_proj = PQDense(config, embed_dim, enable_pruning=False, **qkv_kwargs) + self.out_proj = PQDense( + config, + embed_dim, + quantize_input=True, + quantize_output=quantize_output, + out_quant_granularity=out_quant_granularity, + **proj_kwargs, + ) + + self.attn_dropout = keras.layers.Dropout(dropout) if dropout > 0.0 else None + + def post_pre_train_function(self): + self.is_pretraining = False + for proj in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): + proj.post_pre_train_function() + self.softmax.post_pre_train_function() + + def pre_finetune_function(self): + for proj in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): + proj.pre_finetune_function() + + def pre_epoch_function(self, epoch, total_epochs): + for proj in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): + proj.pre_epoch_function(epoch, total_epochs) + + def post_epoch_function(self, epoch, total_epochs, **kwargs): + for proj in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): + proj.post_epoch_function(epoch, total_epochs, **kwargs) + + def post_round_function(self): + for proj in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): + proj.post_round_function() + + def _save_weights(self): + for proj in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): + proj._save_weights() + + def _rewind_weights(self): + for proj in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): + proj._rewind_weights() + + def _head_bits(self, proj, seq_len): + """Bitwidths of a projection's output, in per-head layout (1, H, seq, head_dim).""" + bw = proj.output_quantizer.get_total_bits((1, seq_len, self.embed_dim)) + bw = ops.reshape(bw, (1, seq_len, self.num_heads, self.head_dim)) + return ops.transpose(bw, (0, 2, 1, 3)) + + def attention_ebops(self): + """EBOPs of the q @ k^T and attn @ v einsums (mirrors HGQ's QMultiHeadAttention).""" + attn_shape = self.softmax.input_shape # (1, H, T, S), stored when the softmax was built + query_len, key_len = attn_shape[2], attn_shape[3] + bw_q = self._head_bits(self.q_proj, query_len) + bw_k = self._head_bits(self.k_proj, key_len) + bw_v = self._head_bits(self.v_proj, key_len) + + bw_attn = self.softmax.output_quantizer.get_total_bits(attn_shape) + ebops_qk = ops.einsum("bhtd,bhsd->", bw_q, bw_k) + ebops_av = ops.einsum("bhts,bhsd->", bw_attn, bw_v) + return ebops_qk + ebops_av + + def ebops(self): + ebops = self.attention_ebops() + self.softmax.ebops() + for proj in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): + ebops += proj.ebops(include_mask=proj.enable_pruning) + return ebops + + def hgq_loss(self): + if self.is_pretraining or not self.use_hgq: + return ops.convert_to_tensor(0.0) + return ops.convert_to_tensor(self.hgq_beta * self.attention_ebops() + self.softmax.hgq_loss()) + + def build(self, input_shape): + if isinstance(input_shape, (list, tuple)) and isinstance(input_shape[0], (list, tuple)): + if len(input_shape) == 3: + query_shape, key_shape, value_shape = input_shape + elif len(input_shape) == 2: + query_shape, key_shape = input_shape + value_shape = key_shape + else: + query_shape = key_shape = value_shape = input_shape[0] + else: + query_shape = key_shape = value_shape = input_shape + if not self.q_proj.built: + self.q_proj.build(query_shape) + if not self.k_proj.built: + self.k_proj.build(key_shape) + if not self.v_proj.built: + self.v_proj.build(value_shape) + scores_shape = (query_shape[0], self.num_heads, query_shape[1], key_shape[1]) + if not self.softmax.built: + self.softmax.build(scores_shape) + if not self.out_proj.built: + self.out_proj.build((query_shape[0], query_shape[1], self.embed_dim)) + super().build(input_shape) + + def compute_output_spec(self, inputs, training=None, key_padding_mask=None, attn_mask=None, need_weights=True): + if isinstance(inputs, (list, tuple)): + query = inputs[0] + key = inputs[1] if len(inputs) > 1 else inputs[0] + else: + query = key = inputs + out = keras.KerasTensor((query.shape[0], query.shape[1], self.embed_dim), dtype=self.compute_dtype) + if need_weights: + attn = keras.KerasTensor((query.shape[0], query.shape[1], key.shape[1]), dtype=self.compute_dtype) + return out, attn + return out, None + + def call( + self, + inputs, + training=None, + key_padding_mask=None, + attn_mask=None, + need_weights=True, + ): + if isinstance(inputs, (list, tuple)): + if len(inputs) == 3: + query, key, value = inputs + elif len(inputs) == 2: + query, key = inputs + value = key + else: + query = key = value = inputs[0] + else: + query = key = value = inputs + + batch_size = ops.shape(query)[0] + query_len = ops.shape(query)[1] + key_len = ops.shape(key)[1] + + q = self.q_proj(query, training=training) # (B, T, E) + k = self.k_proj(key, training=training) # (B, S, E) + v = self.v_proj(value, training=training) # (B, S, E) + + q = ops.reshape(q, (batch_size, query_len, self.num_heads, self.head_dim)) + q = ops.transpose(q, (0, 2, 1, 3)) + k = ops.reshape(k, (batch_size, key_len, self.num_heads, self.head_dim)) + k = ops.transpose(k, (0, 2, 1, 3)) + v = ops.reshape(v, (batch_size, key_len, self.num_heads, self.head_dim)) + v = ops.transpose(v, (0, 2, 1, 3)) + + # Scaled dot-product attention scores: (B, H, T, S) + attn_scores = ops.matmul(q, ops.transpose(k, (0, 1, 3, 2))) * self.scale + + if attn_mask is not None: + if ops.ndim(attn_mask) == 2: + # (T, S) -> (1, 1, T, S) + attn_mask = ops.reshape(attn_mask, (1, 1, query_len, key_len)) + elif ops.ndim(attn_mask) == 3: + # (B*H, T, S) -> (B, H, T, S) + attn_mask = ops.reshape(attn_mask, (batch_size, self.num_heads, query_len, key_len)) + attn_scores = attn_scores + ops.cast(attn_mask, attn_scores.dtype) + + mask = None + if key_padding_mask is not None: + mask = ops.logical_not(ops.cast(key_padding_mask, "bool")) + mask = ops.reshape(mask, (batch_size, 1, 1, key_len)) # (B, 1, 1, S) + + # The softmax's own input/output quantizers handle the scores and the attention weights; + attn_weights = self.softmax(attn_scores, mask=mask) + + if self.attn_dropout is not None: + attn_weights = self.attn_dropout(attn_weights, training=training) + + # Weighted sum of values: (B, H, T, head_dim) + out = ops.matmul(attn_weights, v) + + # Merge heads: (B, T, E) + out = ops.transpose(out, (0, 2, 1, 3)) + out = ops.reshape(out, (batch_size, query_len, self.embed_dim)) + out = self.out_proj(out, training=training) + + if self.use_hgq and self.enable_quantization: + self.add_loss(self.hgq_loss()) + + if need_weights: + # Average attention weights over heads: (B, T, S) + return out, ops.mean(attn_weights, axis=1) + return out, None + + def get_config(self): + config = super().get_config() + config.update( + { + "config": self.config.get_dict(), + "embed_dim": self.embed_dim, + "num_heads": self.num_heads, + "dropout": self.dropout_rate, + "bias": self.use_bias, + "kdim": self.kdim, + "vdim": self.vdim, + "quantize_input": self.q_proj.quantize_input, + "quantize_output": self.out_proj.quantize_output, + "approximate_softmax": self.approximate_softmax, + "in_quant_bits": self.in_quant_bits, + "weight_quant_bits": self.weight_quant_bits, + "bias_quant_bits": self.bias_quant_bits, + "out_quant_bits": self.out_quant_bits, + "attn_quant_bits": self.attn_quant_bits, + "in_quant_granularity": self.in_quant_granularity, + "out_quant_granularity": self.out_quant_granularity, + "param_quant_granularity": self.param_quant_granularity, + } + ) + return config + + @classmethod + def from_config(cls, config): + config = config.copy() + config.pop("q_proj", None) + config.pop("k_proj", None) + config.pop("v_proj", None) + config.pop("out_proj", None) + config.pop("softmax", None) + return cls(**config) + + +LAYERS_WITH_PRUNING_LAYER = (PQWeightBiasBase, PQSeparableConv2d, PQMultiheadAttention) + + +def _iter_weight_layers(model): + for layer in model.layers: + if isinstance(layer, PQWeightBiasBase): + yield layer + elif isinstance(layer, PQSeparableConv2d): + yield layer.depthwise_conv + yield layer.pointwise_conv + elif isinstance(layer, PQMultiheadAttention): + yield layer.q_proj + yield layer.k_proj + yield layer.v_proj + yield layer.out_proj + + def call_post_round_functions(model, rewind, rounds, r): last_round = r == rounds - 1 if rewind == "every-round": - rewind_weights_functions(model) + _rewind_weights_functions(model) elif rewind == "post-training-stage" and last_round: - rewind_weights_functions(model) + _rewind_weights_functions(model) elif not last_round: - post_round_functions(model) + _post_round_functions(model) def apply_final_compression(model): for layer in model.layers: - if isinstance(layer, (PQWeightBiasBase, PQSeparableConv2d, PQBatchNormalization, PQDepthwiseConv2d)): + if isinstance(layer, (PQWeightBiasBase, PQSeparableConv2d, PQBatchNormalization)): layer.apply_final_compression() if hasattr(layer, "input_quantizer"): layer.input_quantizer.apply_final_compression() if hasattr(layer, "output_quantizer"): layer.output_quantizer.apply_final_compression() + elif isinstance(layer, PQMultiheadAttention): + for proj in (layer.q_proj, layer.k_proj, layer.v_proj, layer.out_proj): + proj.apply_final_compression() return model -def _update_pruning_mask(layer): - if layer.enable_pruning and hasattr(layer.pruning_layer, "update_mask"): - kernel = layer.handle_transpose(layer._kernel, layer.weight_transpose, True) - layer.pruning_layer.update_mask(kernel) - - def post_epoch_functions(model, epoch, total_epochs, **kwargs): for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): - if layer.enable_pruning: - layer.pruning_layer.post_epoch_function(epoch, total_epochs, **kwargs) - _update_pruning_mask(layer) - elif isinstance(layer, PQSeparableConv2d): - if layer.enable_pruning: - layer.depthwise_conv.pruning_layer.post_epoch_function(epoch, total_epochs, **kwargs) - _update_pruning_mask(layer.depthwise_conv) - layer.pointwise_conv.pruning_layer.post_epoch_function(epoch, total_epochs, **kwargs) - _update_pruning_mask(layer.pointwise_conv) + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + layer.post_epoch_function(epoch, total_epochs, **kwargs) def pre_epoch_functions(model, epoch, total_epochs): for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): - if layer.enable_pruning: - layer.pruning_layer.pre_epoch_function(epoch, total_epochs) - elif isinstance(layer, PQSeparableConv2d): - if layer.enable_pruning: - layer.depthwise_conv.pruning_layer.pre_epoch_function(epoch, total_epochs) - layer.pointwise_conv.pruning_layer.pre_epoch_function(epoch, total_epochs) + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + layer.pre_epoch_function(epoch, total_epochs) -def post_round_functions(model): +def _post_round_functions(model): for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): - layer.pruning_layer.post_round_function() - elif isinstance(layer, PQSeparableConv2d): - layer.depthwise_conv.pruning_layer.post_round_function() - layer.pointwise_conv.pruning_layer.post_round_function() + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + layer.post_round_function() def save_weights_functions(model): for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): - layer.save_weights() - elif isinstance(layer, PQSeparableConv2d): - layer.depthwise_conv.save_weights() - layer.pointwise_conv.save_weights() + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + layer._save_weights() -def rewind_weights_functions(model): +def _rewind_weights_functions(model): for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): - layer.rewind_weights() - elif isinstance(layer, PQSeparableConv2d): - layer.depthwise_conv.rewind_weights() - layer.pointwise_conv.rewind_weights() + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + layer._rewind_weights() def pre_finetune_functions(model): for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): layer.pre_finetune_function() - layer.pruning_layer.pre_finetune_function() - elif isinstance(layer, PQSeparableConv2d): - layer.depthwise_conv.pre_finetune_function() - layer.depthwise_conv.pruning_layer.pre_finetune_function() - layer.pointwise_conv.pre_finetune_function() - layer.pointwise_conv.pruning_layer.pre_finetune_function() def post_pretrain_functions(model, config): @@ -1847,187 +2049,91 @@ def post_pretrain_functions(model, config): if isinstance( layer, ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, + PQWeightBiasBase, + PQSeparableConv2d, + PQActivation, + PQAvgPoolBase, + PQBatchNormalization, + PQSoftmax, + PQMultiheadAttention, ), ): layer.post_pre_train_function() - elif isinstance(layer, PQSeparableConv2d): - layer.depthwise_conv.post_pre_train_function() - layer.pointwise_conv.post_pre_train_function() - elif isinstance(layer, (PQActivation, PQAvgPoolBase, PQBatchNormalization)): - layer.post_pre_train_function() if config.pruning_parameters.pruning_method == "pdp" or ( config.pruning_parameters.pruning_method == "wanda" and config.pruning_parameters.calculate_pruning_budget ): - pdp_setup(model, config) + _pdp_setup(model, config) -def pdp_setup(model, config): +def _pdp_setup(model, config): """ Calculates a global sparsity threshold. Initializes target sparsity for each layer, which depends on how large percentage of weights in the layer is smaller than the global threshold """ - global_weights = None - for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): - if global_weights is None: - global_weights = ops.ravel(layer.kernel) - else: - global_weights = ops.concatenate((global_weights, ops.ravel(layer.kernel))) - elif isinstance(layer, PQSeparableConv2d): - if global_weights is None: - global_weights = ops.ravel(layer.depthwise_conv.kernel) - global_weights = ops.concatenate((global_weights, ops.ravel(layer.pointwise_conv.kernel))) - else: - global_weights = ops.concatenate((global_weights, ops.ravel(layer.depthwise_conv.kernel))) - global_weights = ops.concatenate((global_weights, ops.ravel(layer.pointwise_conv.kernel))) - + global_weights = ops.concatenate([ops.ravel(layer.kernel) for layer in _iter_weight_layers(model)]) abs_global_weights = ops.abs(global_weights) global_weight_topk, _ = ops.top_k(abs_global_weights, ops.size(abs_global_weights)) threshold = global_weight_topk[int((1 - config.pruning_parameters.sparsity) * float(ops.size(global_weight_topk)))] global_weights_below_threshold = ops.where(abs_global_weights < threshold, 1, 0) idx = 0 - for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): - weight_size = ops.size(layer.kernel) - w = ops.sum(global_weights_below_threshold[idx : idx + weight_size]) - layer.pruning_layer.init_r = ops.convert_to_tensor(w / weight_size, dtype=layer.kernel.dtype) - layer.pruning_layer.sparsity = ops.convert_to_tensor(w / weight_size, dtype=layer.kernel.dtype) # Wanda - idx += weight_size - elif isinstance(layer, PQSeparableConv2d): - weight_size = ops.size(layer.depthwise_conv.kernel) - w = ops.sum(global_weights_below_threshold[idx : idx + weight_size]) - layer.depthwise_conv.pruning_layer.init_r = ops.convert_to_tensor( - w / weight_size, dtype=layer.depthwise_conv.kernel.dtype - ) - layer.depthwise_conv.pruning_layer.sparsity = ops.convert_to_tensor( - w / weight_size, dtype=layer.depthwise_conv.kernel.dtype - ) # Wanda - idx += weight_size - - weight_size = ops.size(layer.pointwise_conv.kernel) - w = ops.sum(global_weights_below_threshold[idx : idx + weight_size]) - layer.pointwise_conv.pruning_layer.init_r = ops.convert_to_tensor( - w / weight_size, dtype=layer.pointwise_conv.kernel.dtype - ) - layer.pointwise_conv.pruning_layer.sparsity = ops.convert_to_tensor( - w / weight_size, dtype=layer.pointwise_conv.kernel.dtype - ) # Wanda - idx += weight_size + for layer in _iter_weight_layers(model): + weight_size = ops.size(layer.kernel) + w = ops.sum(global_weights_below_threshold[idx : idx + weight_size]) + sparsity = ops.convert_to_tensor(w / weight_size, dtype=layer.kernel.dtype) + layer.pruning_layer.init_r = sparsity + layer.pruning_layer.sparsity = sparsity # Wanda + idx += weight_size def get_layer_keep_ratio(model): total_w = 0 remaining_weights = 0 for layer in model.layers: - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): + if isinstance(layer, PQWeightBiasBase): weight = layer.kernel total_w += ops.size(weight) - rem = ops.count_nonzero(weight) - remaining_weights += rem - elif isinstance(layer, PQSeparableConv2d): - depthwise_weight = ops.cast(layer.depthwise_conv.kernel, layer.depthwise_conv.kernel.dtype) - pointwise_weight = ops.cast(layer.pointwise_conv.kernel, layer.pointwise_conv.kernel.dtype) - - depthwise_weight = layer.depthwise_conv.kernel - transpose = layer.depthwise_conv.weight_transpose - if layer.depthwise_conv.enable_pruning: - depthwise_weight = layer.depthwise_conv.pruning_layer.get_hard_mask( - ops.transpose(depthwise_weight, transpose) - ) * ops.transpose(depthwise_weight, transpose) - total_w += ops.size(layer.depthwise_conv.kernel) - rem = ops.count_nonzero(depthwise_weight) - remaining_weights += rem - - pointwise_weight = layer.pointwise_conv.kernel - transpose = layer.pointwise_conv.weight_transpose - if layer.pointwise_conv.enable_pruning: - pointwise_weight = layer.pointwise_conv.pruning_layer.get_hard_mask( - ops.transpose(pointwise_weight, transpose) - ) * ops.transpose(pointwise_weight, transpose) - total_w += ops.size(layer.pointwise_conv.kernel) - rem = ops.count_nonzero(pointwise_weight) - remaining_weights += rem - + remaining_weights += ops.count_nonzero(weight) + elif isinstance(layer, (PQSeparableConv2d, PQMultiheadAttention)): + if isinstance(layer, PQSeparableConv2d): + sublayers = (layer.depthwise_conv, layer.pointwise_conv) + else: + sublayers = (layer.q_proj, layer.k_proj, layer.v_proj, layer.out_proj) + for sublayer in sublayers: + weight = sublayer.kernel + total_w += ops.size(weight) + remaining_weights += ops.count_nonzero(weight) elif isinstance(layer, (Conv2D, Conv1D, DepthwiseConv2D, Dense)): weight = layer.kernel total_w += ops.size(weight) remaining_weights += ops.count_nonzero(weight) elif isinstance(layer, SeparableConv2D): - depthwise_weight = layer.depthwise_kernel - pointwise_weight = layer.pointwise_kernel - total_w += ops.size(depthwise_weight) - total_w += ops.size(pointwise_weight) - remaining_weights += ops.count_nonzero(depthwise_weight) - remaining_weights += ops.count_nonzero(pointwise_weight) + total_w += ops.size(layer.depthwise_kernel) + total_w += ops.size(layer.pointwise_kernel) + remaining_weights += ops.count_nonzero(layer.depthwise_kernel) + remaining_weights += ops.count_nonzero(layer.pointwise_kernel) if total_w != 0: return remaining_weights / total_w return 0.0 -def is_training_stage(layer): - return False if layer.pruning_layer.is_finetuning or layer.pruning_layer.is_pretraining else True +def _is_training_stage(layer): + return not (layer.pruning_layer._is_finetuning or layer.pruning_layer._is_pretraining) def get_model_losses(model, losses): + for layer in _iter_weight_layers(model): + if layer.enable_pruning and _is_training_stage(layer): + losses += layer.pruning_layer.calculate_additional_loss() + if layer.enable_quantization and layer.use_hgq: + losses += layer.hgq_loss() for layer in model.layers: - loss = 0.0 - if isinstance( - layer, - ( - PQDepthwiseConv2d, - PQConv2d, - PQConv1d, - PQDense, - ), - ): - if layer.enable_pruning and is_training_stage(layer): - loss += layer.pruning_layer.calculate_additional_loss() - if layer.enable_quantization and layer.use_hgq: - loss += layer.hgq_loss() - losses += loss - elif isinstance(layer, PQSeparableConv2d): - if layer.enable_pruning and is_training_stage(layer): - loss += layer.depthwise_conv.pruning_layer.calculate_additional_loss() - loss += layer.pointwise_conv.pruning_layer.calculate_additional_loss() - if layer.enable_quantization and layer.use_hgq: - loss += layer.depthwise_conv.hgq_loss() - loss += layer.pointwise_conv.hgq_loss() - losses += loss - elif isinstance(layer, (PQActivation, PQAvgPoolBase, PQBatchNormalization)): + if isinstance(layer, (PQActivation, PQAvgPoolBase, PQBatchNormalization, PQSoftmax)): if layer.enable_quantization and layer.use_hgq: losses += layer.hgq_loss() return losses -def check_activation(layer, config): +def _check_activation(layer, config): """ Replaces activations with quantized activations. The activation can be a part of another layer such as Conv2D, or an Activation layer @@ -2044,7 +2150,7 @@ def check_activation(layer, config): else ReLU() ) if quantization_enabled: - set_quantization_bits_activations(config, layer, act) + _set_quantization_bits_activations(config, layer, act) act.build(layer.input.shape) elif layer.activation.__name__ == "tanh": type_of_tanh = "tanh" if config.quantization_parameters.use_real_tanh else "hard_tanh" @@ -2054,13 +2160,22 @@ def check_activation(layer, config): else Activation(activation="tanh") ) if quantization_enabled: - set_quantization_bits_activations(config, layer, act) + _set_quantization_bits_activations(config, layer, act) act.build(layer.input.shape) - else: - act = None return act +def _build_pruning_layer_from_kernel(new_layer, kernel): + transposed_kernel = ops.transpose(kernel, new_layer.weight_transpose) + new_layer.pruning_layer.build(transposed_kernel.shape) + + +def _copy_kernel_and_bias(new_layer, layer): + new_layer._kernel.assign(layer._kernel) + if layer.use_bias: + new_layer._bias.assign(layer.bias) + + def add_compression_layers(model, config, input_shape=None): # Pruning algorithms assume channels_first format # Creates a new functional model from model, replacing certain layers with compressed / quantized variants @@ -2090,17 +2205,11 @@ def add_compression_layers(model, config, input_shape=None): quantize_input=quantize_input, quantize_output=quantize_output, ) - set_quantization_bits_weight_layers(config, layer, new_layer) - - enable_pruning = get_enable_pruning(layer, config) - new_layer.set_enable_pruning(enable_pruning) - pruning_layer_input = layer.kernel - transpose_shape = new_layer.weight_transpose - pruning_layer_input = ops.transpose(pruning_layer_input, transpose_shape) - new_layer.pruning_layer.build(pruning_layer_input.shape) - + _set_quantization_bits_weight_layers(config, layer, new_layer) + new_layer.set_enable_pruning(_get_enable_pruning(layer, config)) + _build_pruning_layer_from_kernel(new_layer, layer.kernel) x = new_layer(x) - act = check_activation(layer, config) + act = _check_activation(layer, config) elif isinstance(layer, Conv2D): new_layer = PQConv2d( config=config, @@ -2122,19 +2231,13 @@ def add_compression_layers(model, config, input_shape=None): quantize_input=quantize_input, quantize_output=quantize_output, ) - set_quantization_bits_weight_layers(config, layer, new_layer) - enable_pruning = get_enable_pruning(layer, config) - new_layer.set_enable_pruning(enable_pruning) - pruning_layer_input = layer.kernel - transpose_shape = new_layer.weight_transpose - pruning_layer_input = ops.transpose(pruning_layer_input, transpose_shape) - new_layer.pruning_layer.build(pruning_layer_input.shape) + _set_quantization_bits_weight_layers(config, layer, new_layer) + new_layer.set_enable_pruning(_get_enable_pruning(layer, config)) + _build_pruning_layer_from_kernel(new_layer, layer.kernel) new_layer.build(x.shape) x = new_layer(x) - new_layer._kernel.assign(layer._kernel) - if layer.use_bias: - new_layer._bias.assign(layer.bias) - act = check_activation(layer, config) + _copy_kernel_and_bias(new_layer, layer) + act = _check_activation(layer, config) elif isinstance(layer, SeparableConv2D): new_layer = PQSeparableConv2d( config, @@ -2158,26 +2261,20 @@ def add_compression_layers(model, config, input_shape=None): quantize_input=quantize_input, quantize_output=quantize_output, ) - set_quantization_bits_weight_layers(config, layer, new_layer) + _set_quantization_bits_weight_layers(config, layer, new_layer) - enable_pruning_depthwise, enable_pruning_pointwise = get_enable_pruning(layer, config) + enable_pruning_depthwise, enable_pruning_pointwise = _get_enable_pruning(layer, config) new_layer.depthwise_conv.set_enable_pruning(enable_pruning_depthwise) new_layer.pointwise_conv.set_enable_pruning(enable_pruning_pointwise) - - pruning_layer_input = layer.depthwise_kernel - pruning_layer_input = ops.transpose(pruning_layer_input, new_layer.depthwise_conv.weight_transpose) - new_layer.depthwise_conv.pruning_layer.build(pruning_layer_input.shape) - - pointwise_pruning_layer_input = layer.pointwise_kernel - pointwise_pruning_layer_input = ops.transpose( - pointwise_pruning_layer_input, new_layer.pointwise_conv.weight_transpose - ) - new_layer.pointwise_conv.pruning_layer.build(pointwise_pruning_layer_input.shape) - new_layer.depthwise_conv.build(x.shape) - y = new_layer.depthwise_conv(x).shape - new_layer.pointwise_conv.build(y) + _build_pruning_layer_from_kernel(new_layer.depthwise_conv, layer.depthwise_kernel) + _build_pruning_layer_from_kernel(new_layer.pointwise_conv, layer.pointwise_kernel) + new_layer.build(x.shape) x = new_layer(x) - act = check_activation(layer, config) + new_layer.depthwise_conv._kernel.assign(layer.depthwise_kernel) + new_layer.pointwise_conv._kernel.assign(layer.pointwise_kernel) + if layer.use_bias: + new_layer.pointwise_conv._bias.assign(layer.bias) + act = _check_activation(layer, config) elif isinstance(layer, Conv1D): new_layer = PQConv1d( config=config, @@ -2193,19 +2290,13 @@ def add_compression_layers(model, config, input_shape=None): quantize_input=quantize_input, quantize_output=quantize_output, ) - set_quantization_bits_weight_layers(config, layer, new_layer) - enable_pruning = get_enable_pruning(layer, config) - new_layer.set_enable_pruning(enable_pruning) - pruning_layer_input = layer.kernel - transpose_shape = new_layer.weight_transpose - pruning_layer_input = ops.transpose(pruning_layer_input, transpose_shape) - new_layer.pruning_layer.build(pruning_layer_input.shape) + _set_quantization_bits_weight_layers(config, layer, new_layer) + new_layer.set_enable_pruning(_get_enable_pruning(layer, config)) + _build_pruning_layer_from_kernel(new_layer, layer.kernel) new_layer.build(x.shape) x = new_layer(x) - new_layer._kernel.assign(layer._kernel) - if layer.use_bias: - new_layer._bias.assign(layer.bias) - act = check_activation(layer, config) + _copy_kernel_and_bias(new_layer, layer) + act = _check_activation(layer, config) elif isinstance(layer, Dense): new_layer = PQDense( config=config, @@ -2221,30 +2312,24 @@ def add_compression_layers(model, config, input_shape=None): quantize_input=quantize_input, quantize_output=quantize_output, ) - set_quantization_bits_weight_layers(config, layer, new_layer) - enable_pruning = get_enable_pruning(layer, config) - new_layer.set_enable_pruning(enable_pruning) - pruning_layer_input = layer.kernel - transpose_shape = new_layer.weight_transpose - pruning_layer_input = ops.transpose(pruning_layer_input, transpose_shape) - new_layer.pruning_layer.build(pruning_layer_input.shape) + _set_quantization_bits_weight_layers(config, layer, new_layer) + new_layer.set_enable_pruning(_get_enable_pruning(layer, config)) + _build_pruning_layer_from_kernel(new_layer, layer.kernel) x = new_layer(x) - new_layer._kernel.assign(layer._kernel) - if layer.use_bias: - new_layer._bias.assign(layer.bias) - act = check_activation(layer, config) + _copy_kernel_and_bias(new_layer, layer) + act = _check_activation(layer, config) # Activation layers elif isinstance(layer, ReLU): if config.quantization_parameters.enable_quantization: new_layer = PQActivation(config, "relu", quantize_input=quantize_input, quantize_output=quantize_output) - set_quantization_bits_activations(config, layer, new_layer) + _set_quantization_bits_activations(config, layer, new_layer) new_layer.build(layer.input.shape) x = new_layer(x) else: x = layer(x) elif isinstance(layer, Activation): - new_layer = check_activation(layer, config) + new_layer = _check_activation(layer, config) if new_layer is not None: x = new_layer(x) @@ -2257,7 +2342,7 @@ def add_compression_layers(model, config, input_shape=None): padding=layer.padding, data_format=layer.data_format, ) - set_quantization_bits_activations(config, layer, new_layer) + _set_quantization_bits_activations(config, layer, new_layer) new_layer.build(x.shape) x = new_layer(x) elif isinstance(layer, AveragePooling2D): @@ -2269,7 +2354,7 @@ def add_compression_layers(model, config, input_shape=None): padding=layer.padding, data_format=layer.data_format, ) - set_quantization_bits_activations(config, layer, new_layer) + _set_quantization_bits_activations(config, layer, new_layer) new_layer.build(x.shape) x = new_layer(x) elif isinstance(layer, (BatchNormalization)): @@ -2292,7 +2377,7 @@ def add_compression_layers(model, config, input_shape=None): layer.synchronized, quantize_input=True, ) - set_quantization_bits_activations(config, layer, new_layer) + _set_quantization_bits_activations(config, layer, new_layer) new_layer.build(x.shape) x = new_layer(x) else: @@ -2305,55 +2390,39 @@ def add_compression_layers(model, config, input_shape=None): return replaced_model -def set_quantization_bits_activations(config, layer, new_layer): - i_input = i_output = i_weight = i_bias = config.quantization_parameters.default_data_integer_bits - f_input = f_output = f_weight = f_bias = config.quantization_parameters.default_data_fractional_bits +def _get_quant_section(section, i_default, f_default, target=None, quantize_attr=None): + """Read integer/fractional bits from one layer_specific section, optionally applying its quantize flag.""" + if section is None: + return i_default, f_default + if quantize_attr is not None and "quantize" in section: + setattr(target, quantize_attr, section["quantize"]) + return section.get("integer_bits", i_default), section.get("fractional_bits", f_default) + + +def _set_quantization_bits_activations(config, layer, new_layer): + quant_params = config.quantization_parameters + i_input = i_output = i_weight = i_bias = quant_params.default_data_integer_bits + f_input = f_output = f_weight = f_bias = quant_params.default_data_fractional_bits if isinstance(layer, ReLU): f_input += 1 f_output += 1 # Unsigned, add 1 bit to default value only - layer_specific = config.quantization_parameters.layer_specific - if layer.name in layer_specific: - layer_config = layer_specific[layer.name] + layer_config = quant_params.layer_specific.get(layer.name) + if layer_config is not None: if hasattr(layer, "activation") and layer.activation.__name__ in layer_config: - if "input" in layer_config[layer.activation.__name__]: - if "integer_bits" in layer_config[layer.activation.__name__]["input"]: - i_input = layer_config[layer.activation.__name__]["input"]["integer_bits"] - if "integer_bits" in layer_config[layer.activation.__name__]["input"]: - f_input = layer_config[layer.activation.__name__]["input"]["fractional_bits"] - if "quantize" in layer_config[layer.activation.__name__]["input"]: - new_layer.quantize_input = layer_config[layer.activation.__name__]["input"]["quantize"] - if "output" in layer_config[layer.activation.__name__]: - if "integer_bits" in layer_config[layer.activation.__name__]["output"]: - i_output = layer_config[layer.activation.__name__]["output"]["integer_bits"] - if "fractional_bits" in layer_config[layer.activation.__name__]["output"]: - f_output = layer_config[layer.activation.__name__]["output"]["fractional_bits"] - if "quantize" in layer_config[layer.activation.__name__]["output"]: - new_layer.quantize_output = layer_config[layer.activation.__name__]["output"]["quantize"] + activation_config = layer_config[layer.activation.__name__] + i_input, f_input = _get_quant_section( + activation_config.get("input"), i_input, f_input, new_layer, "quantize_input" + ) + i_output, f_output = _get_quant_section( + activation_config.get("output"), i_output, f_output, new_layer, "quantize_output" + ) else: - if "input" in layer_config: - if "integer_bits" in layer_config["input"]: - i_input = layer_config["input"]["integer_bits"] - if "fractional_bits" in layer_config["input"]: - f_input = layer_config["input"]["fractional_bits"] - if "quantize" in layer_config["input"]: - new_layer.quantize_input = layer_config["input"]["quantize"] - if "weight" in layer_config: - if "integer_bits" in layer_config["weight"]: - i_weight = layer_config["weight"]["integer_bits"] - if "fractional_bits" in layer_config["weight"]: - f_weight = layer_config["weight"]["fractional_bits"] - if "bias" in layer_config: - if "integer_bits" in layer_config["bias"]: - i_bias = layer_config["bias"]["integer_bits"] - if "fractional_bits" in layer_config["bias"]: - f_bias = layer_config["bias"]["fractional_bits"] - if "output" in layer_config: - if "integer_bits" in layer_config["output"]: - i_output = layer_config["output"]["integer_bits"] - if "fractional_bits" in layer_config["output"]: - f_output = layer_config["output"]["fractional_bits"] - if "quantize" in layer_config["output"]: - new_layer.quantize_output = layer_config["output"]["quantize"] + i_input, f_input = _get_quant_section(layer_config.get("input"), i_input, f_input, new_layer, "quantize_input") + i_weight, f_weight = _get_quant_section(layer_config.get("weight"), i_weight, f_weight) + i_bias, f_bias = _get_quant_section(layer_config.get("bias"), i_bias, f_bias) + i_output, f_output = _get_quant_section( + layer_config.get("output"), i_output, f_output, new_layer, "quantize_output" + ) if isinstance(layer, BatchNormalization): new_layer.i_weight = i_weight new_layer.f_weight = f_weight @@ -2365,86 +2434,59 @@ def set_quantization_bits_activations(config, layer, new_layer): new_layer.f_output = f_output -def set_quantization_bits_weight_layers(config, layer, new_layer): - layer_specific = config.quantization_parameters.layer_specific +def _set_quantization_bits_weight_layers(config, layer, new_layer): + quant_params = config.quantization_parameters + layer_config = quant_params.layer_specific.get(layer.name) if isinstance(layer, SeparableConv2D): - dw_i_bits_w = pw_i_bits_w = pw_i_bits_b = config.quantization_parameters.default_weight_integer_bits - dw_f_bits_w = pw_f_bits_w = pw_f_bits_b = config.quantization_parameters.default_weight_fractional_bits - i_input = i_output = config.quantization_parameters.default_data_integer_bits - f_input = f_output = config.quantization_parameters.default_data_fractional_bits - if layer.name in layer_specific: - layer_config = layer_specific[layer.name] - if "input" in layer_config: - if "quantize" in layer_config["input"]: - new_layer.depthwise_conv.quantize_input = layer_config["input"]["quantize"] - if "integer_bits" in layer_config["input"]: - i_input = layer_config["input"]["integer_bits"] - if "fractional_bits" in layer_config["input"]: - f_input = layer_config["input"]["fractional_bits"] - if "depthwise" in layer_config: - if "weight" in layer_config["depthwise"]: - dw_i_bits_w = layer_config["depthwise"]["weight"]["integer_bits"] - dw_f_bits_w = layer_config["depthwise"]["weight"]["fractional_bits"] - if "pointwise" in layer_config: - if "weight" in layer_config["pointwise"]: - pw_i_bits_w = layer_config["pointwise"]["weight"]["integer_bits"] - pw_f_bits_w = layer_config["pointwise"]["weight"]["fractional_bits"] - if "bias" in layer_config: - pw_i_bits_b = layer_config["pointwise"]["bias"]["integer_bits"] - pw_f_bits_b = layer_config["pointwise"]["bias"]["fractional_bits"] - if "output" in layer_config: - if "quantize" in layer_config["output"]: - new_layer.quantize_output = layer_config["output"]["quantize"] - if "integer_bits" in layer_config["output"]: - i_output = layer_config["output"]["integer_bits"] - if "fractional_bits" in layer_config["output"]: - f_output = layer_config["output"]["fractional_bits"] + dw_i_weight = pw_i_weight = pw_i_bias = quant_params.default_weight_integer_bits + dw_f_weight = pw_f_weight = pw_f_bias = quant_params.default_weight_fractional_bits + i_input = i_output = quant_params.default_data_integer_bits + f_input = f_output = quant_params.default_data_fractional_bits + if layer_config is not None: + i_input, f_input = _get_quant_section( + layer_config.get("input"), i_input, f_input, new_layer.depthwise_conv, "quantize_input" + ) + depthwise_config = layer_config.get("depthwise", {}) + dw_i_weight, dw_f_weight = _get_quant_section(depthwise_config.get("weight"), dw_i_weight, dw_f_weight) + pointwise_config = layer_config.get("pointwise", {}) + pw_i_weight, pw_f_weight = _get_quant_section(pointwise_config.get("weight"), pw_i_weight, pw_f_weight) + pw_i_bias, pw_f_bias = _get_quant_section(pointwise_config.get("bias"), pw_i_bias, pw_f_bias) + i_output, f_output = _get_quant_section( + layer_config.get("output"), i_output, f_output, new_layer, "quantize_output" + ) new_layer.depthwise_conv.i_input = i_input new_layer.depthwise_conv.f_input = f_input - new_layer.depthwise_conv.i_weight = dw_i_bits_w - new_layer.depthwise_conv.f_weight = dw_f_bits_w - new_layer.pointwise_conv.i_weight = pw_i_bits_w - new_layer.pointwise_conv.f_weight = pw_f_bits_w - new_layer.pointwise_conv.i_bias = pw_i_bits_b - new_layer.pointwise_conv.f_bias = pw_f_bits_b + new_layer.depthwise_conv.i_weight = dw_i_weight + new_layer.depthwise_conv.f_weight = dw_f_weight + new_layer.pointwise_conv.i_weight = pw_i_weight + new_layer.pointwise_conv.f_weight = pw_f_weight + new_layer.pointwise_conv.i_bias = pw_i_bias + new_layer.pointwise_conv.f_bias = pw_f_bias new_layer.pointwise_conv.i_output = i_output new_layer.pointwise_conv.f_output = f_output else: - i_bits_w = i_bits_b = config.quantization_parameters.default_weight_integer_bits - f_bits_w = f_bits_b = config.quantization_parameters.default_weight_fractional_bits - if layer.name in layer_specific: - layer_config = layer_specific[layer.name] - if "input" in layer_config: - if "quantize" in layer_config["input"]: - new_layer.quantize_input = layer_config["input"]["quantize"] - if "integer_bits" in layer_config["input"]: - new_layer.i_input = layer_config["input"]["integer_bits"] - if "fractional_bits" in layer_config["input"]: - new_layer.f_input = layer_config["input"]["fractional_bits"] - if "weight" in layer_config: - i_bits_w = layer_config["weight"]["integer_bits"] - f_bits_w = layer_config["weight"]["fractional_bits"] - if "bias" in layer_config: - i_bits_b = layer_config["bias"]["integer_bits"] - f_bits_b = layer_config["bias"]["fractional_bits"] - if "output" in layer_config: - if "quantize" in layer_config["output"]: - new_layer.quantize_output = layer_config["output"]["quantize"] - if "integer_bits" in layer_config["output"]: - new_layer.i_output = layer_config["output"]["integer_bits"] - if "fractional_bits" in layer_config["output"]: - new_layer.f_output = layer_config["output"]["fractional_bits"] - new_layer.i_weight = i_bits_w - new_layer.f_weight = f_bits_w - new_layer.i_bias = i_bits_b - new_layer.f_bias = f_bits_b - new_layer.weight_quantizer.i_init = float(i_bits_w) - new_layer.weight_quantizer.f_init = float(f_bits_w) - new_layer.bias_quantizer.i_init = float(i_bits_b) - new_layer.bias_quantizer.f_init = float(f_bits_b) - - -def get_enable_pruning(layer, config): + i_weight = i_bias = quant_params.default_weight_integer_bits + f_weight = f_bias = quant_params.default_weight_fractional_bits + if layer_config is not None: + new_layer.i_input, new_layer.f_input = _get_quant_section( + layer_config.get("input"), new_layer.i_input, new_layer.f_input, new_layer, "quantize_input" + ) + i_weight, f_weight = _get_quant_section(layer_config.get("weight"), i_weight, f_weight) + i_bias, f_bias = _get_quant_section(layer_config.get("bias"), i_bias, f_bias) + new_layer.i_output, new_layer.f_output = _get_quant_section( + layer_config.get("output"), new_layer.i_output, new_layer.f_output, new_layer, "quantize_output" + ) + new_layer.i_weight = i_weight + new_layer.f_weight = f_weight + new_layer.i_bias = i_bias + new_layer.f_bias = f_bias + new_layer.weight_quantizer.i_init = float(i_weight) + new_layer.weight_quantizer.f_init = float(f_weight) + new_layer.bias_quantizer.i_init = float(i_bias) + new_layer.bias_quantizer.f_init = float(f_bias) + + +def _get_enable_pruning(layer, config): enable_pruning = config.pruning_parameters.enable_pruning if isinstance(layer, (SeparableConv2D, PQSeparableConv2d)): enable_pruning_depthwise = enable_pruning_pointwise = True @@ -2459,7 +2501,7 @@ def get_enable_pruning(layer, config): return enable_pruning -def populate_config_with_all_layers(model, config): +def _populate_config_with_all_layers(model, config): """Create a default config, where all the layers are added to the disable_pruning list, and have their own default quantization bits in layer_specific. By default input/output quantization is disabled. """ @@ -2552,8 +2594,8 @@ def post_training_prune(model, config, calibration_data): def get_ebops(model, **kwargs): ebops = 0 for m in model.layers: - if isinstance(m, (PQWeightBiasBase)): + if isinstance(m, PQWeightBiasBase): ebops += m.ebops(include_mask=m.enable_pruning) - elif isinstance(m, (PQAvgPoolBase, PQBatchNormalization, PQActivation)): + elif isinstance(m, (PQAvgPoolBase, PQBatchNormalization, PQActivation, PQSoftmax, PQMultiheadAttention)): ebops += m.ebops() return ebops diff --git a/src/pquant/core/keras/pruning_methods/__init__.py b/src/pquant/core/keras/pruning_methods/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/pquant/pruning_methods/activation_pruning.py b/src/pquant/core/keras/pruning_methods/activation_pruning.py similarity index 100% rename from src/pquant/pruning_methods/activation_pruning.py rename to src/pquant/core/keras/pruning_methods/activation_pruning.py diff --git a/src/pquant/pruning_methods/autosparse.py b/src/pquant/core/keras/pruning_methods/autosparse.py similarity index 97% rename from src/pquant/pruning_methods/autosparse.py rename to src/pquant/core/keras/pruning_methods/autosparse.py index 07c1252..910f09a 100644 --- a/src/pquant/pruning_methods/autosparse.py +++ b/src/pquant/core/keras/pruning_methods/autosparse.py @@ -113,7 +113,9 @@ def call(self, weight): is_training = ops.logical_not(ops.logical_or(self.is_pretraining, self.is_finetuning)) self.mask.assign(ops.where(is_training, new_binary_mask, ops.convert_to_tensor(self.mask))) - sparse_weight = ops.sign(weight) * ops.reshape(autosparse_prune(w_t, self.alpha), weight.shape) + sparse_weight = ops.sign(weight) * ops.reshape( + autosparse_prune(w_t, ops.convert_to_tensor(self.alpha)), weight.shape + ) return ops.where( self.is_pretraining, diff --git a/src/pquant/pruning_methods/constraint_functions.py b/src/pquant/core/keras/pruning_methods/constraint_functions.py similarity index 96% rename from src/pquant/pruning_methods/constraint_functions.py rename to src/pquant/core/keras/pruning_methods/constraint_functions.py index 431cdd2..f835346 100644 --- a/src/pquant/pruning_methods/constraint_functions.py +++ b/src/pquant/core/keras/pruning_methods/constraint_functions.py @@ -24,19 +24,19 @@ def __init__(self, lmbda_init=1.0, scale=1.0, damping=1.0, **kwargs): super().__init__(**kwargs) self.scale = self.add_weight( - name='scale', + name="scale", shape=(), initializer=lambda shape, dtype: ops.convert_to_tensor(scale, dtype=dtype), trainable=False, ) self.damping = self.add_weight( - name='damping', + name="damping", shape=(), initializer=lambda shape, dtype: ops.convert_to_tensor(damping, dtype=dtype), trainable=False, ) self.lmbda = self.add_weight( - name=f'{self.name}_lmbda', + name=f"{self.name}_lmbda", shape=(), initializer=lambda shape, dtype: ops.convert_to_tensor(lmbda_init, dtype=dtype), trainable=self.use_grad_, @@ -44,7 +44,7 @@ def __init__(self, lmbda_init=1.0, scale=1.0, damping=1.0, **kwargs): if not self.use_grad_: self.prev_infs = self.add_weight( - name=f'{self.name}_prev_infs', + name=f"{self.name}_prev_infs", shape=(), initializer=lambda shape, dtype: ops.convert_to_tensor(0.0, dtype=dtype), trainable=False, diff --git a/src/pquant/pruning_methods/cs.py b/src/pquant/core/keras/pruning_methods/cs.py similarity index 94% rename from src/pquant/pruning_methods/cs.py rename to src/pquant/core/keras/pruning_methods/cs.py index 97e64ff..eddf561 100644 --- a/src/pquant/pruning_methods/cs.py +++ b/src/pquant/core/keras/pruning_methods/cs.py @@ -45,6 +45,7 @@ def call(self, weight): use_current_mask = ops.logical_or(self.is_pretraining, self.is_finetuning) updated_mask = ops.where(use_current_mask, stored_mask, new_mask) self.mask.assign(updated_mask) + self.add_loss(self.calculate_additional_loss()) return updated_mask * weight def pre_finetune_function(self): @@ -82,9 +83,11 @@ def post_round_function(self): self.beta.assign(1.0) def calculate_additional_loss(self): - return ops.convert_to_tensor( + penalty = ops.convert_to_tensor( self.config.pruning_parameters.threshold_decay * ops.norm(ops.ravel(self.get_mask()), ord=1) ) + inactive = ops.logical_or(self.is_pretraining, self.is_finetuning) + return ops.where(inactive, ops.zeros_like(penalty), penalty) def get_layer_sparsity(self, weight): return ops.sum(self.get_hard_mask()) / ops.size(weight) diff --git a/src/pquant/pruning_methods/dst.py b/src/pquant/core/keras/pruning_methods/dst.py similarity index 95% rename from src/pquant/pruning_methods/dst.py rename to src/pquant/core/keras/pruning_methods/dst.py index e2d5ca9..c040db9 100644 --- a/src/pquant/pruning_methods/dst.py +++ b/src/pquant/core/keras/pruning_methods/dst.py @@ -108,9 +108,9 @@ def get_layer_sparsity(self, weight): return ops.sum(self.get_mask(weight)) / ops.size(weight) def calculate_additional_loss(self): - if self._is_pretraining or self._is_finetuning: - return ops.cast(0.0, self.threshold.dtype) - return self.config.pruning_parameters.alpha * ops.sum(ops.exp(-self.threshold)) + penalty = self.config.pruning_parameters.alpha * ops.sum(ops.exp(-self.threshold)) + inactive = ops.logical_or(self.is_pretraining, self.is_finetuning) + return ops.where(inactive, ops.zeros_like(penalty), penalty) def pre_finetune_function(self): self._is_finetuning = True diff --git a/src/pquant/pruning_methods/mdmm.py b/src/pquant/core/keras/pruning_methods/mdmm.py similarity index 77% rename from src/pquant/pruning_methods/mdmm.py rename to src/pquant/core/keras/pruning_methods/mdmm.py index 5837dd0..62b5e01 100644 --- a/src/pquant/pruning_methods/mdmm.py +++ b/src/pquant/core/keras/pruning_methods/mdmm.py @@ -8,7 +8,26 @@ import keras from keras import ops -from pquant.core.constants import CONSTRAINT_REGISTRY, METRIC_REGISTRY +from pquant.core.keras.pruning_methods.constraint_functions import ( + EqualityConstraint, + GreaterThanOrEqualConstraint, + LessThanOrEqualConstraint, +) +from pquant.core.keras.pruning_methods.metric_functions import ( + StructuredSparsityMetric, + UnstructuredSparsityMetric, +) + +METRIC_REGISTRY = { + "UnstructuredSparsity": UnstructuredSparsityMetric, + "StructuredSparsity": StructuredSparsityMetric, +} + +CONSTRAINT_REGISTRY = { + "Equality": EqualityConstraint, + "LessThanOrEqual": LessThanOrEqualConstraint, + "GreaterThanOrEqual": GreaterThanOrEqualConstraint, +} # ------------------------------------------------------------------- # MDMM Layer @@ -28,6 +47,10 @@ def __init__(self, config, layer_type, *args, **kwargs): self.constraint_layer = None self._is_finetuning = False self._is_pretraining = True + # TEMP: cache last penalty so calculate_additional_loss() works in + # custom training loops via get_model_losses(). Remove once the + # add_loss()/model.fit path is the only supported path. + self._last_penalty = None def build(self, input_shape): pruning_parameters = self.config.pruning_parameters @@ -94,8 +117,11 @@ def call(self, weight): self.mask.assign(ops.where(not_active, ops.convert_to_tensor(self.mask), hard_mask)) penalty = ops.sum(self.constraint_layer(weight)) - self.add_loss(ops.where(not_active, ops.zeros_like(penalty), penalty)) - + gated_penalty = ops.where(not_active, ops.zeros_like(penalty), penalty) + self.add_loss(gated_penalty) + # TEMP: cache for calculate_additional_loss() — remove with the + # _last_penalty attribute once custom-loop callers move to model.losses. + self._last_penalty = gated_penalty return ops.where(self.is_finetuning, weight * hard_mask, weight) def get_hard_mask(self, weight=None): @@ -109,7 +135,12 @@ def get_layer_sparsity(self, weight): def calculate_additional_loss(self): # Loss is added via self.add_loss() in call() for model.fit. - # For custom training loops, accumulate model.losses from the last forward pass instead. + # TEMP: also return the cached penalty so custom training loops using + # get_model_losses() see the constraint term. Remove this branch (and + # the _last_penalty cache) once those callers switch to model.losses; + # then this can revert to `return 0.0`. + if self._last_penalty is not None: + return self._last_penalty return 0.0 def pre_epoch_function(self, epoch, total_epochs): diff --git a/src/pquant/pruning_methods/metric_functions.py b/src/pquant/core/keras/pruning_methods/metric_functions.py similarity index 88% rename from src/pquant/pruning_methods/metric_functions.py rename to src/pquant/core/keras/pruning_methods/metric_functions.py index 0f22b5e..071b84b 100644 --- a/src/pquant/pruning_methods/metric_functions.py +++ b/src/pquant/core/keras/pruning_methods/metric_functions.py @@ -6,10 +6,10 @@ class UnstructuredSparsityMetric: """Calculates the ratio of non-zero weights in a tensor.""" - def __init__(self, l0_mode='coarse', scale_mode="mean", epsilon=1e-3, target_sparsity=0.8, alpha=100.0): + def __init__(self, l0_mode="coarse", scale_mode="mean", epsilon=1e-3, target_sparsity=0.8, alpha=100.0): # Note: scale_mode:"sum" give very high losses for large model - assert l0_mode in ['coarse', 'smooth'], "Mode must be 'coarse' or 'smooth'" - assert scale_mode in ['sum', 'mean'], "Scale mode must be 'sum' or 'mean'" + assert l0_mode in ["coarse", "smooth"], "Mode must be 'coarse' or 'smooth'" + assert scale_mode in ["sum", "mean"], "Scale mode must be 'sum' or 'mean'" assert 0 <= target_sparsity <= 1, "target_sparsity must be between 0 and 1" self.l0_mode = l0_mode self.scale_mode = scale_mode @@ -24,14 +24,14 @@ def __init__(self, l0_mode='coarse', scale_mode="mean", epsilon=1e-3, target_spa def build(self): # l0 term -> number of zero weights/number of weights - if self.l0_mode == 'coarse': + if self.l0_mode == "coarse": self.l0_fn = self._coarse_l0 - elif self.l0_mode == 'smooth': + elif self.l0_mode == "smooth": self.l0_fn = self._smooth_l0 - if self.scale_mode == 'mean': + if self.scale_mode == "mean": self._scaling = self._mean_scaling - elif self.scale_mode == 'sum': + elif self.scale_mode == "sum": self._scaling = self._sum_scaling def _sum_scaling(self, fn_value, num): diff --git a/src/pquant/pruning_methods/pdp.py b/src/pquant/core/keras/pruning_methods/pdp.py similarity index 100% rename from src/pquant/pruning_methods/pdp.py rename to src/pquant/core/keras/pruning_methods/pdp.py diff --git a/src/pquant/pruning_methods/wanda.py b/src/pquant/core/keras/pruning_methods/wanda.py similarity index 100% rename from src/pquant/pruning_methods/wanda.py rename to src/pquant/core/keras/pruning_methods/wanda.py diff --git a/src/pquant/core/keras/quantizer.py b/src/pquant/core/keras/quantizer.py index 38b96af..c3f0405 100644 --- a/src/pquant/core/keras/quantizer.py +++ b/src/pquant/core/keras/quantizer.py @@ -1,9 +1,10 @@ -from enum import Enum - import keras +from hgq.quantizer import Quantizer as HGQQuantizer +from hgq.quantizer import QuantizerConfig from keras import ops +from quantizers import get_fixed_quantizer -from pquant.core.quantizer_functions import create_quantizer +from pquant.core.constants import QuantizationGranularity @keras.saving.register_keras_serializable(package="PQuantML") @@ -18,9 +19,10 @@ def __init__( round_mode="RND", is_heterogeneous=False, is_data=False, - granularity="per_tensor", + granularity=QuantizationGranularity.PER_TENSOR, hgq_gamma=0, place="datalane", + dynamic_data=True, ): super().__init__() self.k_init = float(k) @@ -31,19 +33,42 @@ def __init__( self.round_mode = round_mode self.use_hgq = is_heterogeneous self.is_data = is_data + self.dynamic_data = dynamic_data self.place = place + self.granularity = QuantizationGranularity(granularity).value self.quantizer = create_quantizer( - self.k_init, self.i_init, self.f_init, self.overflow, self.round_mode, self.use_hgq, self.is_data, place + self.k_init, + self.i_init, + self.f_init, + self.overflow, + self.round_mode, + self.use_hgq, + self.is_data, + place, + granularity=self.granularity, ) self.is_pretraining = True self.hgq_gamma = hgq_gamma - if isinstance(granularity, Enum): - self.granularity = granularity.value - else: - self.granularity = granularity - def compute_dynamic_bits(self, x): - if self.granularity == "per_channel": + def calculate_bits_from_abs(self, abs_x): + m = ops.ceil(ops.log(abs_x + 1e-6) / ops.log(2.0)) + int_bits = ops.maximum(m, 0.0) + int_bits = ops.minimum(int_bits, self.b - self.k) + frac_bits = ops.maximum(self.b - int_bits - self.k, 0.0) + return int_bits, frac_bits + + def compute_data_dynamic_bits(self, x): + if not self.dynamic_data: + _, i, f = self.get_quantization_bits() + return i, f + abs_x = ops.max(ops.abs(x)) + return self.calculate_bits_from_abs(abs_x) + + def compute_weight_dynamic_bits(self, x): + if self.granularity == QuantizationGranularity.PER_TENSOR or ops.ndim(x) == 1: + _, i, f = self.get_quantization_bits() + return i, f + if self.granularity == QuantizationGranularity.PER_CHANNEL: if ops.ndim(x) == 2: abs_x = ops.max(ops.abs(x), axis=0, keepdims=True) elif ops.ndim(x) == 3: @@ -52,50 +77,50 @@ def compute_dynamic_bits(self, x): abs_x = ops.max(ops.abs(x), axis=(0, 1, 2), keepdims=True) else: raise ValueError("Unsupported tensor rank") - elif self.granularity == "per_weight": + elif self.granularity == QuantizationGranularity.PER_WEIGHT: abs_x = ops.abs(x) else: raise ValueError(f"compute_dynamic_bits called for granularity={self.granularity}") - m = ops.ceil(ops.log(abs_x + 1e-6) / ops.log(2.0)) - int_bits = ops.maximum(m, 0.0) - b = self.b if hasattr(self, "b") else self.b_init - frac_bits = ops.maximum(b - int_bits - self.k_init, 0.0) - return int_bits, frac_bits + return self.calculate_bits_from_abs(abs_x) + + def compute_dynamic_bits(self, x): + if self.is_data: + return self.compute_data_dynamic_bits(x) + return self.compute_weight_dynamic_bits(x) + + def _dynamic_bits_shape(self, input_shape): + """Shape the dynamically computed i/f will have, without computing them (used by build).""" + if self.is_data or len(input_shape) == 1: + return () # data and bias are per_tensor, so scalar shape + if self.granularity == QuantizationGranularity.PER_CHANNEL: + if len(input_shape) not in (2, 3, 4): + raise ValueError("Unsupported tensor rank") + return (1,) * (len(input_shape) - 1) + (input_shape[-1],) + return tuple(input_shape) # per_weight def build(self, input_shape): if self.use_hgq: shape = tuple(input_shape) if not self.is_data else (1,) + tuple(input_shape[1:]) - self.k = self.add_weight(shape=shape, initializer=keras.initializers.Constant(self.k_init), trainable=False) - self.i = self.add_weight(shape=shape, initializer=keras.initializers.Constant(self.i_init), trainable=False) - self.f = self.add_weight(shape=shape, initializer=keras.initializers.Constant(self.f_init), trainable=False) - self.b = self.add_weight( - shape=shape, - initializer=keras.initializers.Constant(self.k_init + self.i_init + self.f_init), - trainable=False, - ) if not self.quantizer.built: self.quantizer.build(shape) self.set_quantization_bits(self.i_init, self.f_init) - elif self.granularity == "per_tensor": - self.k = self.add_weight(shape=(), initializer=keras.initializers.Constant(self.k_init), trainable=False) - self.i = self.add_weight(shape=(), initializer=keras.initializers.Constant(self.i_init), trainable=False) - self.f = self.add_weight(shape=(), initializer=keras.initializers.Constant(self.f_init), trainable=False) - self.b = self.add_weight( - shape=(), initializer=keras.initializers.Constant(self.k_init + self.f_init + self.f_init), trainable=False - ) + elif self.granularity == QuantizationGranularity.PER_TENSOR: + self._build_params(shape=()) else: - i, _ = self.compute_dynamic_bits(keras.ops.ones(input_shape)) - self.k = self.add_weight(shape=i.shape, initializer=keras.initializers.Constant(self.k_init), trainable=False) - self.i = self.add_weight(shape=i.shape, initializer=keras.initializers.Constant(self.i_init), trainable=False) - self.f = self.add_weight(shape=i.shape, initializer=keras.initializers.Constant(self.f_init), trainable=False) - self.b = self.add_weight( - shape=i.shape, - initializer=keras.initializers.Constant(self.k_init + self.f_init + self.f_init), - trainable=False, - ) - + shape = self._dynamic_bits_shape(input_shape) + self._build_params(shape) super().build(input_shape) + def _build_params(self, shape): + self.k = self.add_weight(shape=shape, initializer=keras.initializers.Constant(self.k_init), trainable=False) + self.i = self.add_weight(shape=shape, initializer=keras.initializers.Constant(self.i_init), trainable=False) + self.f = self.add_weight(shape=shape, initializer=keras.initializers.Constant(self.f_init), trainable=False) + self.b = self.add_weight( + shape=shape, + initializer=keras.initializers.Constant(self.k_init + self.i_init + self.f_init), + trainable=False, + ) + def get_total_bits(self, shape): if self.use_hgq: return self.quantizer.bits_(shape) @@ -106,22 +131,29 @@ def get_total_bits(self, shape): def get_quantization_bits(self): if self.use_hgq: return self.quantizer.quantizer.k, self.quantizer.quantizer.i, self.quantizer.quantizer.f + if not hasattr(self, "i"): + return self.k_init, self.i_init, self.f_init return self.k, self.i, self.f def set_quantization_bits(self, i, f): if self.use_hgq: self.quantizer.quantizer._i.assign(self.quantizer.quantizer._i * 0.0 + i) self.quantizer.quantizer._f.assign(self.quantizer.quantizer._f * 0.0 + f) - self.i = i - self.f = f + elif hasattr(self, "i") and hasattr(self.i, "assign"): + self.i.assign(ops.zeros_like(self.i) + i) + self.f.assign(ops.zeros_like(self.f) + f) + self.b.assign(self.k + self.i + self.f) + else: + self.i = i + self.f = f def apply_final_compression(self): - if self.use_hgq and not self.quantizer.built or not self.built: + if self.use_hgq or not self.built: return k, i, f = self.get_quantization_bits() - self.i.assign(i) - self.f.assign(f) - self.b.assign(k + i + f) + self.i.assign(ops.zeros_like(self.i) + i) + self.f.assign(ops.zeros_like(self.f) + f) + self.b.assign(ops.zeros_like(self.b) + k + i + f) self.final_compression_done = True def post_pre_train_function(self): @@ -132,12 +164,9 @@ def call(self, x, training=None): return self.quantizer(x, training=training) if not training: return self.quantizer(x, k=self.k, i=self.i, f=self.f, training=training) - elif self.granularity == "per_tensor": - i, f = self.i, self.f - else: - i, f = self.compute_dynamic_bits(x) - self.i.assign(i) - self.f.assign(f) + i, f = self.compute_dynamic_bits(x) + self.i.assign(i) + self.f.assign(f) return self.quantizer(x, k=self.k, i=i, f=f, training=training) def hgq_loss(self): @@ -158,6 +187,7 @@ def from_config(cls, config): is_data=config.pop("is_data"), granularity=config.pop("granularity"), place=config.pop("place"), + dynamic_data=config.pop("dynamic_data", True), ) if use_hgq: @@ -179,8 +209,79 @@ def get_config(self): "is_heterogeneous": self.use_hgq, "granularity": self.granularity, "place": self.place, + "dynamic_data": self.dynamic_data, } ) if self.use_hgq: config.update({"quantizer": keras.saving.serialize_keras_object(self.quantizer)}) return config + + +def axis_kwargs_for_granularity(granularity, is_data): + """Translate a granularity into HGQ's (mutually exclusive) homogeneous/heterogeneous axis spec. + + HGQ only supports per_tensor and per_weight from the granularity enum: + - per_tensor: nothing varies -> heterogeneous_axis=() (one bitwidth for the whole tensor) + - per_weight: every element varies. For data we keep the batch axis (0) homogeneous via + homogeneous_axis=(0,); for weights nothing is shared via homogeneous_axis=(). + + per_channel is intentionally NOT supported for HGQ (the channel axis is layout-dependent, so we + don't guess it). + """ + if granularity == QuantizationGranularity.PER_TENSOR: + return {"heterogeneous_axis": ()} + if granularity == QuantizationGranularity.PER_WEIGHT: + return {"homogeneous_axis": (0,) if is_data else ()} + if granularity == QuantizationGranularity.PER_CHANNEL: + raise ValueError("per_channel granularity is not supported for HGQ. Use 'per_tensor' or 'per_weight'.") + raise ValueError(f"Unsupported granularity: {granularity}") + + +def create_hgq_parameters_quantizer(k, i, f, overflow, round_mode, place, axis_kwargs, gamma=1e-8): + quantizer_config = QuantizerConfig( + q_type="kif", + place=place, + k0=k, + i0=i, + f0=f, + overflow_mode=overflow, + round_mode=round_mode, + **axis_kwargs, + ) + return HGQQuantizer(config=quantizer_config) + + +def create_hgq_data_quantizer(k, i, f, overflow, round_mode, axis_kwargs, gamma=1e-8): + quantizer_config = QuantizerConfig( + q_type="kif", + place="datalane", + k0=k, + i0=i, + f0=f, + overflow_mode=overflow, + round_mode=round_mode, + **axis_kwargs, + ) + return HGQQuantizer(config=quantizer_config) + + +def create_quantizer( + k, + i, + f, + overflow, + round_mode, + is_heterogeneous, + is_data, + place="datalane", + granularity=QuantizationGranularity.PER_WEIGHT, + gamma=1e-8, +): + if is_heterogeneous: + axis_kwargs = axis_kwargs_for_granularity(granularity, is_data) + if is_data: + return create_hgq_data_quantizer(k, i, f, overflow, round_mode, axis_kwargs, gamma=gamma) + else: + return create_hgq_parameters_quantizer(k, i, f, overflow, round_mode, place, axis_kwargs, gamma=gamma) + else: + return get_fixed_quantizer(round_mode=round_mode, overflow_mode=overflow) diff --git a/src/pquant/core/keras/utils.py b/src/pquant/core/keras/utils.py new file mode 100644 index 0000000..a0e4faa --- /dev/null +++ b/src/pquant/core/keras/utils.py @@ -0,0 +1,25 @@ +from pquant.core.keras.pruning_methods.activation_pruning import ActivationPruning +from pquant.core.keras.pruning_methods.autosparse import AutoSparse +from pquant.core.keras.pruning_methods.cs import ContinuousSparsification +from pquant.core.keras.pruning_methods.dst import DST +from pquant.core.keras.pruning_methods.mdmm import MDMM +from pquant.core.keras.pruning_methods.pdp import PDP +from pquant.core.keras.pruning_methods.wanda import Wanda + + +def get_pruning_layer(config, layer_type): + pruning_method = config.pruning_parameters.pruning_method + if pruning_method == "dst": + return DST(config, layer_type) + elif pruning_method == "autosparse": + return AutoSparse(config, layer_type) + elif pruning_method == "cs": + return ContinuousSparsification(config, layer_type) + elif pruning_method == "pdp": + return PDP(config, layer_type) + elif pruning_method == "activation_pruning": + return ActivationPruning(config, layer_type) + elif pruning_method == "wanda": + return Wanda(config, layer_type) + elif pruning_method == "mdmm": + return MDMM(config, layer_type) diff --git a/src/pquant/core/quantizer_functions.py b/src/pquant/core/quantizer_functions.py deleted file mode 100644 index 81922d7..0000000 --- a/src/pquant/core/quantizer_functions.py +++ /dev/null @@ -1,48 +0,0 @@ -import keras - - -def create_fixed_quantizer(k, i, f, overflow, round_mode): - if keras.backend.backend() == "torch": - from pquant.core.torch.fixed_point_quantizer import get_fixed_quantizer - else: - from quantizers import get_fixed_quantizer - - quantizer = get_fixed_quantizer(round_mode=round_mode, overflow_mode=overflow) - return quantizer - - -def create_hgq_parameters_quantizer(k, i, f, overflow, round_mode, place): - from hgq.quantizer import Quantizer, QuantizerConfig - - quantizer_config = QuantizerConfig( - q_type="kif", place=place, k0=k, i0=i, f0=f, overflow_mode=overflow, round_mode=round_mode, homogeneous_axis=() - ) - - return Quantizer(config=quantizer_config) - - -def create_hgq_data_quantizer(k, i, f, overflow, round_mode): - from hgq.quantizer import Quantizer, QuantizerConfig - - quantizer_config = QuantizerConfig( - q_type="kif", - place="datalane", - k0=k, - i0=i, - f0=f, - overflow_mode=overflow, - round_mode=round_mode, - homogeneous_axis=(0,), - ) - - return Quantizer(config=quantizer_config) - - -def create_quantizer(k, i, f, overflow, round_mode, is_heterogeneous, is_data, place="datalane"): - if is_heterogeneous: - if is_data: - return create_hgq_data_quantizer(k, i, f, overflow, round_mode) - else: - return create_hgq_parameters_quantizer(k, i, f, overflow, round_mode, place) - else: - return create_fixed_quantizer(k, i, f, overflow, round_mode) diff --git a/src/pquant/core/torch/activations.py b/src/pquant/core/torch/activations.py index 4630a04..d7e356d 100644 --- a/src/pquant/core/torch/activations.py +++ b/src/pquant/core/torch/activations.py @@ -1,4 +1,5 @@ -from typing import Tuple, TypeVar +from math import prod +from typing import TypeVar import torch import torch.nn as nn @@ -22,7 +23,13 @@ def hard_tanh(x): return 2.0 * hard_sigmoid(x) - 1.0 -activation_registry = {"relu": relu, "tanh": tanh, "hard_tanh": hard_tanh} +activation_registry = { + "relu": relu, + "tanh": tanh, + "hard_tanh": hard_tanh, + "leaky_relu": nn.LeakyReLU(negative_slope=0.1015625), + "gelu": nn.GELU(), +} class PQActivation(nn.Module): @@ -30,10 +37,11 @@ def __init__( self, config, activation="relu", - in_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, quantize_input=True, quantize_output=False, + enable_ebops=True, ): super().__init__() if isinstance(config, dict): @@ -55,8 +63,13 @@ def __init__( else: self.k_output, self.i_output, self.f_output = out_quant_bits - self.activation_name = activation.lower() - self.activation_function = activation_registry.get(self.activation_name) + if isinstance(activation, str): + self.activation_name = activation.lower() + self.activation_function = activation_registry.get(self.activation_name) + else: + # An activation function/callable was passed directly instead of a registry key. + self.activation_function = activation + self.activation_name = getattr(activation, "__name__", activation.__class__.__name__).lower() self.enable_quantization = config.quantization_parameters.enable_quantization self.use_hgq = config.quantization_parameters.use_high_granularity_quantization @@ -69,12 +82,13 @@ def __init__( self.hgq_gamma = config.quantization_parameters.hgq_gamma self.hgq_heterogeneous = config.quantization_parameters.hgq_heterogeneous self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress - self.granularity = config.quantization_parameters.granularity + self.dynamic_data = config.quantization_parameters.dynamic_data_quantization self.post_fitcompress_calibration = False self.saved_inputs = [] self.quantize_input = quantize_input self.quantize_output = quantize_output + self.enable_ebops = enable_ebops self.built = False def check_is_built(self, input_shape): @@ -92,6 +106,8 @@ def check_is_built(self, input_shape): is_heterogeneous=self.use_hgq, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, + granularity=self.config.quantization_parameters.granularity, ) self.input_quantizer = Quantizer( k=self.k_input, @@ -103,6 +119,8 @@ def check_is_built(self, input_shape): is_heterogeneous=self.use_hgq, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.dynamic_data, + granularity=self.config.quantization_parameters.granularity, ) if self.use_hgq: self.input_quantizer.quantizer.build(input_shape) @@ -127,6 +145,8 @@ def post_pre_train_function(self): self.is_pretraining = False def ebops(self): + if not self.enable_ebops: + return torch.tensor(0.0) bw_inp = self.input_quantizer.get_total_bits(self.input_shape) bw_out = self.output_quantizer.get_total_bits(self.input_shape) return torch.sum((2.0**bw_inp) * bw_out) * 1e-4 # type: ignore @@ -182,3 +202,205 @@ def get_config(self): def extra_repr(self): return f"quantize_input = {self.quantize_input}, quantize_output = {self.quantize_output}" + + +class PQSoftmax(nn.Module): + """Quantized softmax that mirrors HGQ's ``QSoftmax``. + + Args: + config: PQuant configuration object (or a serialized dict). + axis: Axis (or axes) the softmax normalizes over. + stable: If True, subtract the max before exponentiating for numerical stability. + input_scaler: Scalar multiplied with the logits before exponentiating. + parallelization_factor: hls4ml parallelization factor used in the ebops cost model. + quantize_input: Whether to quantize the softmax logits before the exp table. + quantize_output: Whether to quantize the exp * inv product (the softmax output). + in_quant_bits: (k, i, f) bits for the softmax input quantizer. + out_quant_bits: (k, i, f) bits for the softmax output quantizer. + exp_in_quant_bits / exp_out_quant_bits: (k, i, f) bits for the exp table. + inv_in_quant_bits / inv_out_quant_bits: (k, i, f) bits for the inv table. + """ + + def __init__( + self, + config, + axis: int | tuple[int, ...] = -1, + stable: bool = True, + input_scaler: float = 1.0, + parallelization_factor: int = -1, + quantize_input: bool = True, + quantize_output: bool = False, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + exp_in_quant_bits: tuple[T, T, T] = None, + exp_out_quant_bits: tuple[T, T, T] = None, + inv_in_quant_bits: tuple[T, T, T] = None, + inv_out_quant_bits: tuple[T, T, T] = None, + **kwargs, + ): + super().__init__(**kwargs) + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + + self._axis = tuple(axis) if isinstance(axis, (tuple, list)) else (axis,) + self.axes = self._axis + self.stable = stable + self.input_scaler = input_scaler + self.parallelization_factor = parallelization_factor + self.quantize_input = quantize_input + self.quantize_output = quantize_output + self.epsilon = 1e-7 + + if in_quant_bits is not None: + self.k_input, self.i_input, self.f_input = in_quant_bits + else: + self.k_input = config.quantization_parameters.default_data_keep_negatives + self.i_input = config.quantization_parameters.default_data_integer_bits + self.f_input = config.quantization_parameters.default_data_fractional_bits + + if out_quant_bits is not None: + self.k_output, self.i_output, self.f_output = out_quant_bits + else: + self.k_output = 0 + self.i_output = config.quantization_parameters.default_data_integer_bits + self.f_output = config.quantization_parameters.default_data_fractional_bits + + self.overflow_mode_data = config.quantization_parameters.overflow_mode_data + self.round_mode = config.quantization_parameters.round_mode + self.use_hgq = config.quantization_parameters.use_high_granularity_quantization + self.hgq_gamma = config.quantization_parameters.hgq_gamma + self.hgq_beta = config.quantization_parameters.hgq_beta + self.enable_quantization = config.quantization_parameters.enable_quantization + self.is_pretraining = True + self.built = False + + i_data = config.quantization_parameters.default_data_integer_bits + f_data = config.quantization_parameters.default_data_fractional_bits + k_data = config.quantization_parameters.default_data_keep_negatives + exp_in_quant_bits = exp_in_quant_bits if exp_in_quant_bits is not None else (k_data, i_data, f_data) + exp_out_quant_bits = exp_out_quant_bits if exp_out_quant_bits is not None else (0, i_data, f_data) + inv_in_quant_bits = inv_in_quant_bits if inv_in_quant_bits is not None else (k_data, i_data, f_data) + inv_out_quant_bits = inv_out_quant_bits if inv_out_quant_bits is not None else (0, i_data, f_data) + + def _exp(x): + if self.stable: + return torch.exp(-x * self.input_scaler) + return torch.exp(x * self.input_scaler) + + def _inv(x): + return 1.0 / (x + self.epsilon) + + self.exp_table = PQActivation( + config, + _exp, + in_quant_bits=exp_in_quant_bits, + out_quant_bits=exp_out_quant_bits, + quantize_input=stable, + quantize_output=True, + enable_ebops=stable, + ) + self.inv_table = PQActivation( + config, + _inv, + in_quant_bits=inv_in_quant_bits, + out_quant_bits=inv_out_quant_bits, + quantize_input=True, + quantize_output=True, + ) + + def check_is_built(self, input_shape): + if self.built: + return + self.built = True + ndim = len(input_shape) + self.axes = tuple(sorted(a if a >= 0 else a + ndim for a in self._axis)) + self.input_shape = (1,) + tuple(input_shape[1:]) + + def _data_quantizer(k, i, f): + return Quantizer( + k=torch.tensor(k), + i=torch.tensor(i), + f=torch.tensor(f), + overflow=self.overflow_mode_data, + round_mode=self.round_mode, + is_heterogeneous=self.use_hgq, + is_data=True, + hgq_gamma=self.hgq_gamma, + place="datalane", + dynamic_data=self.config.quantization_parameters.dynamic_data_quantization, + ) + + self.input_quantizer = _data_quantizer(self.k_input, self.i_input, self.f_input) + self.output_quantizer = _data_quantizer(self.k_output, self.i_output, self.f_output) + if self.use_hgq: + self.input_quantizer.quantizer.build(input_shape) + self.output_quantizer.quantizer.build(input_shape) + + def get_input_quantization_bits(self): + return self.input_quantizer.get_quantization_bits() + + def get_output_quantization_bits(self): + return self.output_quantizer.get_quantization_bits() + + def post_pre_train_function(self): + self.is_pretraining = False + + def ebops(self): + shape = self.input_shape + accum_shape = tuple(1 if i in self.axes else s for i, s in enumerate(shape)) + max_instance = prod(accum_shape) + n_instance = self.parallelization_factor if self.parallelization_factor > 0 else max_instance + factor = n_instance / max_instance + + inp_bits = self.input_quantizer.get_total_bits(shape) + exp_bits = self.exp_table.output_quantizer.get_total_bits(shape) + inv_bits = self.inv_table.output_quantizer.get_total_bits(accum_shape) + + substract_ebops = torch.sum(inp_bits) if self.stable else 0.0 + accum_ebops = torch.sum(exp_bits) - torch.sum(torch.amin(exp_bits, dim=self.axes)) + mult_ebops = torch.sum(exp_bits * inv_bits) + + ebops = substract_ebops + accum_ebops + mult_ebops + if not self.stable: + ebops = ebops + torch.sum((2.0**inp_bits) * exp_bits) * 1e-4 + return ebops * factor + + def hgq_loss(self): + if self.is_pretraining or not self.use_hgq: + return torch.tensor(0.0) + loss = self.hgq_beta * self.ebops() + if self.quantize_input: + loss += self.input_quantizer.hgq_loss() + if self.quantize_output: + loss += self.output_quantizer.hgq_loss() + return loss + + def forward(self, inputs, mask=None): + self.check_is_built(inputs.shape) + if self.quantize_input and self.enable_quantization: + inputs = self.input_quantizer(inputs) + + if self.stable: + inputs = torch.amax(inputs, dim=self.axes, keepdim=True) - inputs + + exp_inp = self.exp_table(inputs) + + if mask is not None: + exp_inp = mask.to(exp_inp.dtype) * exp_inp + + sums = torch.sum(exp_inp, dim=self.axes, keepdim=True) + divisor = self.inv_table(sums) + + out = exp_inp * divisor + if self.quantize_output and self.enable_quantization: + out = self.output_quantizer(out) + return out + + def extra_repr(self) -> str: + return ( + f"axis={self.axes}, stable={self.stable}, input_scaler={self.input_scaler}, " + f"quantize_input={self.quantize_input}, quantize_output={self.quantize_output}" + ) diff --git a/src/pquant/core/torch/convert_to_onnx.py b/src/pquant/core/torch/convert_to_onnx.py new file mode 100644 index 0000000..1601b73 --- /dev/null +++ b/src/pquant/core/torch/convert_to_onnx.py @@ -0,0 +1,1857 @@ +""" +Convert a PQuant model to ONNX or QONNX format. + +Pass ``use_qonnx=True`` to emit QONNX ``Quant`` custom nodes (requires the +qonnx runtime). Pass ``use_qonnx=False`` (default) to emit standard +``Clip + QuantizeLinear + DequantizeLinear`` nodes runnable with plain +onnxruntime. + +Fixed-point (k, i, f) mapping +------------------------------ +QONNX: + scale = 2^(-f) + zero_point = 0 + bit_width = k + i + f + signed = int(k) + +Standard ONNX (QDQ): + scale = 2^(-f) + zero_point = 0 (int8 signed, uint8 unsigned) + clip range = [-2^i, 2^i - 2^(-f)] signed + = [0, 2^i - 2^(-f)] unsigned + Rounding is always nearest-even (QuantizeLinear behaviour). + Weights are stored as plain float32 initializers — after + apply_final_compression() they are already on the fixed-point grid. +""" + +import functools +import logging +import operator as _operator +import os + +import numpy as np +import onnx +import onnx.helper as oh +import onnx.numpy_helper as onh +import torch +import torch.fx as _fx +import torch.nn as nn +import torch.nn.functional as _F +from onnx import TensorProto + +os.environ["KERAS_BACKEND"] = "torch" # must be set before any keras/pquant import + +from pquant.core.torch.activations import PQActivation # noqa: E402 +from pquant.core.torch.layers import ( # noqa: E402 + PQAvgPool1d, + PQAvgPool2d, + PQBatchNorm1d, + PQBatchNorm2d, + PQConv1d, + PQConv2d, + PQDense, + PQLayerNorm, + PQMultiheadAttention, +) +from pquant.core.torch.quantizer import Quantizer # noqa: E402 + +# --------------------------------------------------------------------------- +# QONNX Quant node +# --------------------------------------------------------------------------- + +ROUND_MODE_MAP = { + "TRN": "FLOOR", + "RND": "ROUND", + "RND_CONV": "ROUND", + "TRN_ZERO": "TRUNCATE", + "RND_ZERO": "ROUND", + "RND_MIN_INF": "FLOOR", + "RND_INF": "ROUND", +} + + +def _quant_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT"): + """Build a QONNX Quant node. Returns ([node], output_name). + + QONNX Quant is per-tensor only. If i/f are per-channel or per-weight tensors + (non-scalar), collapse to the broadest range: min(f) / max(i) ensures no channel + overflows at the cost of slightly coarser quantization for small-value channels. + """ + k_val = int(k.item()) + if hasattr(f, "numel") and f.numel() > 1: + i = i.reshape(-1).max() + f = f.reshape(-1).min() + i_val = float(i.item()) + f_val = float(f.item()) + scale = float(2.0 ** (-f_val)) + bit_width = float(k_val + i_val + f_val) + qonnx_rnd = ROUND_MODE_MAP.get(rounding_mode, "ROUND") + narrow = 1 if (k_val == 1 and overflow_mode == "SAT_SYM") else 0 + + scale_name = f"{name_prefix}_scale" + zp_name = f"{name_prefix}_zero_point" + bw_name = f"{name_prefix}_bit_width" + out_name = f"{name_prefix}_quantized" + + initializers.append(onh.from_array(np.array(scale, dtype=np.float32), name=scale_name)) + initializers.append(onh.from_array(np.array(0.0, dtype=np.float32), name=zp_name)) + initializers.append(onh.from_array(np.array(bit_width, dtype=np.float32), name=bw_name)) + + node = oh.make_node( + op_type="Quant", + inputs=[input_name, scale_name, zp_name, bw_name], + outputs=[out_name], + domain="qonnx.custom_op.general", + signed=k_val, + narrow=narrow, + rounding_mode=qonnx_rnd, + ) + return [node], out_name + + +# --------------------------------------------------------------------------- +# Standard ONNX QDQ triple +# --------------------------------------------------------------------------- + + +def _qdq_node(name_prefix, input_name, rounding_mode, k, i, f, initializers, overflow_mode="SAT", include_clip=True): + """Build QuantizeLinear+DequantizeLinear nodes, optionally preceded by a Clip. + + Returns ([nodes], output_name). Set include_clip=False to skip the Clip node + (safe when values are guaranteed in-range at inference time, since + QuantizeLinear saturates naturally). + """ + k_val = int(k.item()) + i_val = float(i.item()) + f_val = float(f.item()) + scale = float(2.0 ** (-f_val)) + signed = k_val == 1 + + clip_max = float(2.0**i_val - 2.0 ** (-f_val)) + if not signed: + clip_min = 0.0 + elif overflow_mode == "SAT_SYM": + clip_min = -clip_max + else: + clip_min = float(-(2.0**i_val)) + zp_val = np.int8(0) if signed else np.uint8(0) + + scale_name = f"{name_prefix}_scale" + zp_name = f"{name_prefix}_zero_point" + quantized_name = f"{name_prefix}_quantized" + out_name = f"{name_prefix}_dequantized" + + initializers += [ + onh.from_array(np.array(scale, dtype=np.float32), name=scale_name), + onh.from_array(np.array(zp_val), name=zp_name), + ] + + if include_clip: + clip_min_name = f"{name_prefix}_clip_min" + clip_max_name = f"{name_prefix}_clip_max" + clipped_name = f"{name_prefix}_clipped" + initializers += [ + onh.from_array(np.array(clip_min, dtype=np.float32), name=clip_min_name), + onh.from_array(np.array(clip_max, dtype=np.float32), name=clip_max_name), + ] + nodes = [ + oh.make_node("Clip", inputs=[input_name, clip_min_name, clip_max_name], outputs=[clipped_name]), + oh.make_node("QuantizeLinear", inputs=[clipped_name, scale_name, zp_name], outputs=[quantized_name]), + ] + else: + nodes = [ + oh.make_node("QuantizeLinear", inputs=[input_name, scale_name, zp_name], outputs=[quantized_name]), + ] + nodes.append(oh.make_node("DequantizeLinear", inputs=[quantized_name, scale_name, zp_name], outputs=[out_name])) + return nodes, out_name + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _int_weight_node(name_prefix, weight_np, k, i, f, initializers): # noqa: ARG001 (i unused) + """ + Store a weight tensor as int8/uint8 + DequantizeLinear. + + weight_np must already be on the fixed-point grid (guaranteed after + apply_final_compression). Converts by dividing by the scale and casting — + no re-rounding needed. + + Granularity handling: + - per-tensor (f is scalar): single scale, standard DequantizeLinear. + - per-channel (f has shape [out, 1, ...]): 1D scale with axis=0. + All weights in a channel share the same f so the conversion is exact. + - per-weight (f is fully per-element): ONNX has no per-weight quantization; + falls back to float32 storage (no DequantizeLinear node). + + Returns ([node], output_name). + """ + k_val = int(k.item()) if hasattr(k, "item") else int(k) + dtype = np.int8 if k_val == 1 else np.uint8 + out_channels = weight_np.shape[0] + out_name = f"{name_prefix}_dequantized" + + f_t = f.detach().cpu() if hasattr(f, "detach") else torch.as_tensor(f) + + if f_t.numel() == 1: + # per-tensor + scale_np = np.array(float(2.0 ** (-f_t.item())), dtype=np.float32) + int_weights = np.round(weight_np / float(scale_np)).astype(dtype) + per_channel = False + else: + f_np = f_t.float().numpy().reshape(out_channels, -1) + if np.allclose(f_np, f_np[:, :1]): + # per-channel: all elements within an output channel share one f + f_1d = f_np[:, 0] + scale_np = (2.0 ** (-f_1d)).astype(np.float32) + bcast = scale_np.reshape((out_channels,) + (1,) * (weight_np.ndim - 1)) + int_weights = np.round(weight_np / bcast).astype(dtype) + per_channel = True + else: + # per-weight: ONNX cannot represent this; store as float32 + float_name = f"{name_prefix}_float" + initializers.append(onh.from_array(weight_np, name=float_name)) + return [], float_name + + int_name = f"{name_prefix}_int" + scale_name = f"{name_prefix}_dq_scale" + zp_name = f"{name_prefix}_dq_zp" + + zp_np = np.zeros(out_channels if per_channel else 1, dtype=dtype) + initializers += [ + onh.from_array(int_weights, name=int_name), + onh.from_array(scale_np, name=scale_name), + onh.from_array(zp_np if per_channel else np.array(dtype(0)), name=zp_name), + ] + node_kwargs = {"axis": 0} if per_channel else {} + node = oh.make_node("DequantizeLinear", inputs=[int_name, scale_name, zp_name], outputs=[out_name], **node_kwargs) + return [node], out_name + + +def _torch_padding_to_onnx(padding, ndim): + if isinstance(padding, int): + padding = (padding,) * ndim + return list(padding) + list(padding) + + +def _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn): + if ( + getattr(module, "input_quantizer", None) is not None + and getattr(module, "quantize_input", True) + and getattr(module, "enable_quantization", True) + ): + q = module.input_quantizer + k, i, f = q.get_quantization_bits() + new_nodes, current = quant_fn( + f"{prefix}_in", current, q.round_mode, k, i, f, initializers, overflow_mode=getattr(q, "overflow", "SAT") + ) + nodes.extend(new_nodes) + return current + + +def _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn): + if ( + getattr(module, "output_quantizer", None) is not None + and getattr(module, "quantize_output", False) + and getattr(module, "enable_quantization", True) + ): + q = module.output_quantizer + k, i, f = q.get_quantization_bits() + new_nodes, current = quant_fn( + f"{prefix}_out", current, q.round_mode, k, i, f, initializers, overflow_mode=getattr(q, "overflow", "SAT") + ) + nodes.extend(new_nodes) + return current + + +# --------------------------------------------------------------------------- +# per-layer graph builders +# --------------------------------------------------------------------------- + + +def _add_dense_integer(module, prefix, current, nodes, initializers): + """Dense layer using MatMulInteger for true integer arithmetic. + + Flow: + float → Clip+QuantizeLinear → int8 ─┐ + ├─ MatMulInteger → int32 + int8 weights (pre-transposed) ───────┘ + → Add int32 bias + → DequantizeLinear(scale = s_x * s_w) → float + + The inner product accumulates in int32; there is no float Gemm. + A single DequantizeLinear at the end converts back to float for activations. + Per-channel weights use axis=1 on the output DequantizeLinear. + """ + if not (getattr(module, "input_quantizer", None) and getattr(module, "quantize_input", True)): + raise ValueError(f"{prefix}: integer_ops requires quantize_input=True on the layer") + + # --- Input: Clip + QuantizeLinear → int8 (stop before DequantizeLinear) --- + k_x, i_x, f_x = module.input_quantizer.get_quantization_bits() + k_x_val = int(k_x.item()) + i_x_val = float(i_x.item()) + f_x_val = float(f_x.item()) + s_x = float(2.0 ** (-f_x_val)) + signed_x = k_x_val == 1 + + clip_min_x = float(-(2.0**i_x_val)) if signed_x else 0.0 + clip_max_x = float(2.0**i_x_val - 2.0 ** (-f_x_val)) + zp_x_np = np.int8(0) if signed_x else np.uint8(0) + + clip_min_name = f"{prefix}_in_clip_min" + clip_max_name = f"{prefix}_in_clip_max" + scale_x_name = f"{prefix}_in_scale" + zp_x_name = f"{prefix}_in_zp" + x_int_name = f"{prefix}_in_int" + + initializers += [ + onh.from_array(np.array(clip_min_x, dtype=np.float32), name=clip_min_name), + onh.from_array(np.array(clip_max_x, dtype=np.float32), name=clip_max_name), + onh.from_array(np.array(s_x, dtype=np.float32), name=scale_x_name), + onh.from_array(np.array(zp_x_np), name=zp_x_name), + ] + nodes += [ + oh.make_node("Clip", inputs=[current, clip_min_name, clip_max_name], outputs=[f"{prefix}_in_clipped"]), + oh.make_node("QuantizeLinear", inputs=[f"{prefix}_in_clipped", scale_x_name, zp_x_name], outputs=[x_int_name]), + ] + + # --- Weights: stored pre-transposed as int8 so MatMulInteger needs no Transpose --- + # PyTorch weight shape: [out, in]. MatMulInteger(A, B) = A @ B, so we need [in, out]. + weight_np = module._weight.detach().cpu().numpy().astype(np.float32) + k_w, _, f_w = module.weight_quantizer.get_quantization_bits() + k_w_val = int(k_w.item()) if hasattr(k_w, "item") else int(k_w) + dtype_w = np.int8 if k_w_val == 1 else np.uint8 + out_ch = weight_np.shape[0] + + f_w_t = f_w.detach().cpu() if hasattr(f_w, "detach") else torch.as_tensor(f_w) + if f_w_t.numel() == 1: + f_w_1d = np.array([float(f_w_t.item())]) + per_channel_w = False + else: + f_w_2d = f_w_t.float().numpy().reshape(out_ch, -1) + f_w_1d = f_w_2d.min(axis=1) # min f → max scale → covers all values + per_channel_w = True + + s_w_1d = (2.0 ** (-f_w_1d)).astype(np.float32) # shape [1] or [out] + bcast_s_w = s_w_1d.reshape((out_ch,) + (1,) * (weight_np.ndim - 1)) if per_channel_w else float(s_w_1d[0]) + # Transpose before storing so MatMulInteger can use it without a runtime Transpose node + int_weights_T = np.round(weight_np / bcast_s_w).astype(dtype_w).T # [in, out] + + zp_w_np = np.array(dtype_w(0)) # scalar zero-point; zero for symmetric quantization + w_int_name = f"{prefix}_weight_int" + w_zp_name = f"{prefix}_weight_zp" + initializers += [ + onh.from_array(int_weights_T, name=w_int_name), + onh.from_array(zp_w_np, name=w_zp_name), + ] + + # --- MatMulInteger([batch, in], [in, out]) → int32 [batch, out] --- + y_int_name = f"{prefix}_matmul_int" + nodes.append( + oh.make_node( + "MatMulInteger", + inputs=[x_int_name, w_int_name, zp_x_name, w_zp_name], + outputs=[y_int_name], + ) + ) + + # --- Bias added in int32 domain: bias_int[c] = round(bias[c] / (s_x * s_w[c])) --- + current_int32 = y_int_name + if module._bias is not None: + bias_np = module._bias.detach().cpu().numpy().astype(np.float32) + combined_s = s_x * s_w_1d # shape [1] or [out] + bias_int32 = np.round(bias_np / (combined_s if per_channel_w else float(combined_s[0]))).astype(np.int32) + bias_int_name = f"{prefix}_bias_int" + y_biased_name = f"{prefix}_matmul_biased" + initializers.append(onh.from_array(bias_int32, name=bias_int_name)) + nodes.append(oh.make_node("Add", inputs=[current_int32, bias_int_name], outputs=[y_biased_name])) + current_int32 = y_biased_name + + # --- DequantizeLinear: int32 → float32 using combined scale s_x * s_w --- + # Per-channel: axis=1 because the output tensor is [batch, out] and out is axis 1. + combined_scale_name = f"{prefix}_combined_scale" + combined_zp_name = f"{prefix}_combined_zp" + + if per_channel_w: + combined_scale_np = (s_x * s_w_1d).astype(np.float32) # [out] + combined_zp_np = np.zeros(out_ch, dtype=np.int32) + dql_kwargs = {"axis": 1} + else: + combined_scale_np = np.array(float(s_x * s_w_1d[0]), dtype=np.float32) + combined_zp_np = np.array(np.int32(0)) + dql_kwargs = {} + + initializers += [ + onh.from_array(combined_scale_np, name=combined_scale_name), + onh.from_array(combined_zp_np, name=combined_zp_name), + ] + y_float_name = f"{prefix}_dequantized" + nodes.append( + oh.make_node( + "DequantizeLinear", + inputs=[current_int32, combined_scale_name, combined_zp_name], + outputs=[y_float_name], + **dql_kwargs, + ) + ) + current = y_float_name + + # Optional output quantization (e.g. last layer with quantize_output=True) + current = _maybe_quant_output(module, prefix, current, nodes, initializers, _qdq_node) + return current + + +def _add_dense_nd(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """Dense (linear) projection via MatMul, supporting input of any rank ≥ 2. + + Identical logic to _add_dense but emits ``MatMul(input, W_T)`` instead of + ``Gemm(input, W, transB=1)`` so it accepts (B, T, E) inputs (e.g. from MHA + projections) as well as the usual 2-D (batch, features) inputs. + Weight is stored pre-transposed as [in, out] to avoid a runtime Transpose node. + """ + current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + weight_np = module._weight.detach().cpu().numpy().astype(np.float32) # [out, in] + if use_qonnx: + weight_fp_name = f"{prefix}_weight_fp" + initializers.append(onh.from_array(weight_np, name=weight_fp_name)) + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + w_nodes, q_weight_raw = _quant_node( + f"{prefix}_weight", + weight_fp_name, + module.weight_quantizer.round_mode, + k_w, + i_w, + f_w, + initializers, + overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(w_nodes) + q_weight_t = f"{prefix}_weight_T" + nodes.append(oh.make_node("Transpose", inputs=[q_weight_raw], outputs=[q_weight_t], perm=[1, 0])) + q_weight = q_weight_t + elif store_integer_weights: + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + w_nodes, q_weight_stored = _int_weight_node(f"{prefix}_weight", weight_np, k_w, i_w, f_w, initializers) + nodes.extend(w_nodes) + q_weight_t = f"{prefix}_weight_T" + nodes.append(oh.make_node("Transpose", inputs=[q_weight_stored], outputs=[q_weight_t], perm=[1, 0])) + q_weight = q_weight_t + else: + q_weight = f"{prefix}_weight_T" + initializers.append(onh.from_array(weight_np.T, name=q_weight)) # pre-transposed [in, out] + + matmul_out = f"{prefix}_matmul" + nodes.append(oh.make_node("MatMul", inputs=[current, q_weight], outputs=[matmul_out])) + current = matmul_out + + if module._bias is not None: + bias_np = module._bias.detach().cpu().numpy().astype(np.float32) + if use_qonnx: + bias_fp_name = f"{prefix}_bias_fp" + initializers.append(onh.from_array(bias_np, name=bias_fp_name)) + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _quant_node( + f"{prefix}_bias", + bias_fp_name, + module.bias_quantizer.round_mode, + k_b, + i_b, + f_b, + initializers, + overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights: + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, k_b, i_b, f_b, initializers) + nodes.extend(b_nodes) + else: + q_bias = f"{prefix}_bias" + initializers.append(onh.from_array(bias_np, name=q_bias)) + biased_out = f"{prefix}_biased" + nodes.append(oh.make_node("Add", inputs=[matmul_out, q_bias], outputs=[biased_out])) + current = biased_out + + current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) + return current + + +def _add_dense(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops=False): + if integer_ops and not use_qonnx: + return _add_dense_integer(module, prefix, current, nodes, initializers) + current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + weight_np = module._weight.detach().cpu().numpy().astype(np.float32) + if use_qonnx: + weight_fp_name = f"{prefix}_weight_fp" + initializers.append(onh.from_array(weight_np, name=weight_fp_name)) + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + w_nodes, q_weight = _quant_node( + f"{prefix}_weight", + weight_fp_name, + module.weight_quantizer.round_mode, + k_w, + i_w, + f_w, + initializers, + overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(w_nodes) + elif store_integer_weights: + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + w_nodes, q_weight = _int_weight_node(f"{prefix}_weight", weight_np, k_w, i_w, f_w, initializers) + nodes.extend(w_nodes) + else: + q_weight = f"{prefix}_weight" + initializers.append(onh.from_array(weight_np, name=q_weight)) + + # Use Gemm with transB=1 — weight stays in its native [out, in] layout, + # no Transpose node needed. Bias (if any) is fused as the third Gemm input. + gemm_inputs = [current, q_weight] + + if module._bias is not None: + bias_np = module._bias.detach().cpu().numpy().astype(np.float32) + if use_qonnx: + bias_fp_name = f"{prefix}_bias_fp" + initializers.append(onh.from_array(bias_np, name=bias_fp_name)) + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _quant_node( + f"{prefix}_bias", + bias_fp_name, + module.bias_quantizer.round_mode, + k_b, + i_b, + f_b, + initializers, + overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights: + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, k_b, i_b, f_b, initializers) + nodes.extend(b_nodes) + else: + q_bias = f"{prefix}_bias" + initializers.append(onh.from_array(bias_np, name=q_bias)) + gemm_inputs.append(q_bias) + + gemm_out = f"{prefix}_gemm" + nodes.append(oh.make_node("Gemm", inputs=gemm_inputs, outputs=[gemm_out], transB=1)) + current = gemm_out + + current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) + return current + + +def _add_conv(module, prefix, current, nodes, initializers, ndim, quant_fn, use_qonnx, store_integer_weights): + current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + weight_np = module._weight.detach().cpu().numpy().astype(np.float32) + if use_qonnx: + weight_fp_name = f"{prefix}_weight_fp" + initializers.append(onh.from_array(weight_np, name=weight_fp_name)) + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + w_nodes, q_weight = _quant_node( + f"{prefix}_weight", + weight_fp_name, + module.weight_quantizer.round_mode, + k_w, + i_w, + f_w, + initializers, + overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(w_nodes) + elif store_integer_weights: + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + w_nodes, q_weight = _int_weight_node(f"{prefix}_weight", weight_np, k_w, i_w, f_w, initializers) + nodes.extend(w_nodes) + else: + q_weight = f"{prefix}_weight" + initializers.append(onh.from_array(weight_np, name=q_weight)) + + conv_inputs = [current, q_weight] + + if module._bias is not None: + bias_np = module._bias.detach().cpu().numpy().astype(np.float32) + if use_qonnx: + bias_fp_name = f"{prefix}_bias_fp" + initializers.append(onh.from_array(bias_np, name=bias_fp_name)) + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _quant_node( + f"{prefix}_bias", + bias_fp_name, + module.bias_quantizer.round_mode, + k_b, + i_b, + f_b, + initializers, + overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights: + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_bias = _int_weight_node(f"{prefix}_bias", bias_np, k_b, i_b, f_b, initializers) + nodes.extend(b_nodes) + else: + q_bias = f"{prefix}_bias" + initializers.append(onh.from_array(bias_np, name=q_bias)) + conv_inputs.append(q_bias) + + padding = module.padding + if isinstance(padding, str): + auto_pad = "SAME_UPPER" if padding == "same" else "VALID" + pads = None + else: + auto_pad = "NOTSET" + pads = _torch_padding_to_onnx(padding, ndim) + + to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 + conv_attrs = dict( + kernel_shape=to_list(module.kernel_size, ndim), + strides=to_list(module.stride, ndim), + dilations=to_list(module.dilation, ndim), + group=module.groups, + auto_pad=auto_pad, + ) + if pads is not None: + conv_attrs["pads"] = pads + + conv_out = f"{prefix}_conv" + nodes.append(oh.make_node("Conv", inputs=conv_inputs, outputs=[conv_out], **conv_attrs)) + current = conv_out + + current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) + return current + + +def _add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + gamma_np = module._weight.detach().cpu().numpy().astype(np.float32) + beta_np = module._bias.detach().cpu().numpy().astype(np.float32) + + if use_qonnx: + gamma_fp_name = f"{prefix}_gamma_fp" + initializers.append(onh.from_array(gamma_np, name=gamma_fp_name)) + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + g_nodes, q_gamma = _quant_node( + f"{prefix}_gamma", + gamma_fp_name, + module.weight_quantizer.round_mode, + k_w, + i_w, + f_w, + initializers, + overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(g_nodes) + + beta_fp_name = f"{prefix}_beta_fp" + initializers.append(onh.from_array(beta_np, name=beta_fp_name)) + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_beta = _quant_node( + f"{prefix}_beta", + beta_fp_name, + module.bias_quantizer.round_mode, + k_b, + i_b, + f_b, + initializers, + overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights: + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + g_nodes, q_gamma = _int_weight_node(f"{prefix}_gamma", gamma_np, k_w, i_w, f_w, initializers) + nodes.extend(g_nodes) + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_beta = _int_weight_node(f"{prefix}_beta", beta_np, k_b, i_b, f_b, initializers) + nodes.extend(b_nodes) + else: + q_gamma = f"{prefix}_gamma" + q_beta = f"{prefix}_beta" + initializers.append(onh.from_array(gamma_np, name=q_gamma)) + initializers.append(onh.from_array(beta_np, name=q_beta)) + + mean_name = f"{prefix}_running_mean" + var_name = f"{prefix}_running_var" + initializers.append(onh.from_array(module.running_mean.detach().cpu().numpy().astype(np.float32), name=mean_name)) + initializers.append(onh.from_array(module.running_var.detach().cpu().numpy().astype(np.float32), name=var_name)) + + bn_out = f"{prefix}_bn" + nodes.append( + oh.make_node( + "BatchNormalization", + inputs=[current, q_gamma, q_beta, mean_name, var_name], + outputs=[bn_out], + epsilon=float(module.eps), + ) + ) + return bn_out + + +def _add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """PQLayerNorm. Emits LayerNormalization (opset >= 17 required).""" + current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + ns = ( + tuple(int(d) for d in module.normalized_shape) + if hasattr(module.normalized_shape, "__iter__") + else (int(module.normalized_shape),) + ) + axis = -len(ns) + + has_weight = module._weight is not None + has_bias = module._bias is not None + + gamma_np = module._weight.detach().cpu().numpy().astype(np.float32) if has_weight else np.ones(ns, dtype=np.float32) + beta_np = module._bias.detach().cpu().numpy().astype(np.float32) if has_bias else None + + if use_qonnx and has_weight: + gamma_fp_name = f"{prefix}_gamma_fp" + initializers.append(onh.from_array(gamma_np, name=gamma_fp_name)) + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + g_nodes, q_gamma = _quant_node( + f"{prefix}_gamma", + gamma_fp_name, + module.weight_quantizer.round_mode, + k_w, + i_w, + f_w, + initializers, + overflow_mode=getattr(module.weight_quantizer, "overflow", "SAT"), + ) + nodes.extend(g_nodes) + if has_bias: + beta_fp_name = f"{prefix}_beta_fp" + initializers.append(onh.from_array(beta_np, name=beta_fp_name)) + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_beta = _quant_node( + f"{prefix}_beta", + beta_fp_name, + module.bias_quantizer.round_mode, + k_b, + i_b, + f_b, + initializers, + overflow_mode=getattr(module.bias_quantizer, "overflow", "SAT"), + ) + nodes.extend(b_nodes) + elif store_integer_weights and has_weight: + k_w, i_w, f_w = module.weight_quantizer.get_quantization_bits() + g_nodes, q_gamma = _int_weight_node(f"{prefix}_gamma", gamma_np, k_w, i_w, f_w, initializers) + nodes.extend(g_nodes) + if has_bias: + k_b, i_b, f_b = module.bias_quantizer.get_quantization_bits() + b_nodes, q_beta = _int_weight_node(f"{prefix}_beta", beta_np, k_b, i_b, f_b, initializers) + nodes.extend(b_nodes) + else: + q_gamma = f"{prefix}_gamma" + initializers.append(onh.from_array(gamma_np, name=q_gamma)) + if has_bias: + q_beta = f"{prefix}_beta" + initializers.append(onh.from_array(beta_np, name=q_beta)) + + ln_inputs = [current, q_gamma] + if has_bias: + ln_inputs.append(q_beta) + ln_out = f"{prefix}_ln" + nodes.append( + oh.make_node( + "LayerNormalization", + inputs=ln_inputs, + outputs=[ln_out], + axis=axis, + epsilon=float(module.eps), + ) + ) + current = ln_out + current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) + return current + + +def _add_avgpool(module, prefix, current, nodes, initializers, ndim, quant_fn): + current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + + to_list = lambda v, n: list(v) if hasattr(v, "__iter__") else [v] * n # noqa: E731 + pool_out = f"{prefix}_pool" + nodes.append( + oh.make_node( + "AveragePool", + inputs=[current], + outputs=[pool_out], + kernel_shape=to_list(module.kernel_size, ndim), + strides=to_list(module.stride, ndim), + pads=_torch_padding_to_onnx(module.padding, ndim), + ceil_mode=int(module.ceil_mode), + count_include_pad=int(module.count_include_pad), + ) + ) + current = pool_out + + current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) + return current + + +# --------------------------------------------------------------------------- +# multi-head attention graph builder +# --------------------------------------------------------------------------- + + +def _add_mha(module, prefix, q_input, k_input, v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights): + """Build ONNX nodes for PQMultiheadAttention. + + Decomposes multi-head attention into primitive ONNX ops: + + [optional transpose if not batch_first] + Q/K/V Gemm projections + Reshape (B, L, E) → (B, H, L, head_dim) + Transpose + MatMul(Q, K^T) * scale → optional Quant + Softmax → optional Quant + MatMul(attn_weights, V) → optional Quant + Transpose + Reshape (B, T, E) + out_proj Gemm + [optional transpose back if not batch_first] + + Returns (out_name, avg_attn_weights_name): the projected output and the + attention weights averaged over heads (B, T, S). Both names are valid ONNX + value names so downstream getitem(mha, 0) / getitem(mha, 1) work in the FX + converter. + + Note: if ``approximate_softmax=True`` the module uses a polynomial + approximation in PyTorch, but ONNX has no equivalent standard op — a plain + ``Softmax`` node is emitted instead. + """ + H = module.num_heads + head_dim = module.head_dim + E = module.embed_dim + scale_val = float(module.scale) + + # --- Optional transpose for seq-first inputs (T, B, E) → (B, T, E) --- + if not module.batch_first: + q_t = f"{prefix}_q_in_t" + k_t = f"{prefix}_k_in_t" + v_t = f"{prefix}_v_in_t" + nodes.append(oh.make_node("Transpose", inputs=[q_input], outputs=[q_t], perm=[1, 0, 2])) + nodes.append(oh.make_node("Transpose", inputs=[k_input], outputs=[k_t], perm=[1, 0, 2])) + nodes.append(oh.make_node("Transpose", inputs=[v_input], outputs=[v_t], perm=[1, 0, 2])) + q_input, k_input, v_input = q_t, k_t, v_t + + # --- Q / K / V projections: (B, L, E) → (B, L, E) via MatMul (input is rank-3) --- + q_proj_out = _add_dense_nd( + module.q_proj, f"{prefix}_q_proj", q_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + k_proj_out = _add_dense_nd( + module.k_proj, f"{prefix}_k_proj", k_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + v_proj_out = _add_dense_nd( + module.v_proj, f"{prefix}_v_proj", v_input, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + # --- Helper: (B, L, E) → (B, H, L, head_dim) using dynamic shapes --- + def _split_heads(x_name, pfx): + shape_out = f"{pfx}_shape" + b_scalar = f"{pfx}_b_sc" + l_scalar = f"{pfx}_l_sc" + b_1d = f"{pfx}_b_1d" + l_1d = f"{pfx}_l_1d" + h_1d_const = f"{pfx}_H_1d" + hd_1d_const = f"{pfx}_hd_1d" + shape_4d = f"{pfx}_shape4d" + reshaped = f"{pfx}_reshaped" + transposed = f"{pfx}_transposed" + idx0 = f"{pfx}_gi0" + idx1 = f"{pfx}_gi1" + ax0 = f"{pfx}_ax0" + + nodes.append(oh.make_node("Shape", inputs=[x_name], outputs=[shape_out])) + initializers.extend( + [ + onh.from_array(np.array(0, dtype=np.int64), name=idx0), + onh.from_array(np.array(1, dtype=np.int64), name=idx1), + onh.from_array(np.array([0], dtype=np.int64), name=ax0), + onh.from_array(np.array([H], dtype=np.int64), name=h_1d_const), + onh.from_array(np.array([head_dim], dtype=np.int64), name=hd_1d_const), + ] + ) + nodes.append(oh.make_node("Gather", inputs=[shape_out, idx0], outputs=[b_scalar])) + nodes.append(oh.make_node("Gather", inputs=[shape_out, idx1], outputs=[l_scalar])) + nodes.append(oh.make_node("Unsqueeze", inputs=[b_scalar, ax0], outputs=[b_1d])) + nodes.append(oh.make_node("Unsqueeze", inputs=[l_scalar, ax0], outputs=[l_1d])) + nodes.append(oh.make_node("Concat", inputs=[b_1d, l_1d, h_1d_const, hd_1d_const], outputs=[shape_4d], axis=0)) + nodes.append(oh.make_node("Reshape", inputs=[x_name, shape_4d], outputs=[reshaped])) + # (B, L, H, head_dim) → (B, H, L, head_dim) + nodes.append(oh.make_node("Transpose", inputs=[reshaped], outputs=[transposed], perm=[0, 2, 1, 3])) + return transposed + + q_h = _split_heads(q_proj_out, f"{prefix}_q") + k_h = _split_heads(k_proj_out, f"{prefix}_k") + v_h = _split_heads(v_proj_out, f"{prefix}_v") + + # --- k^T: (B, H, S, head_dim) → (B, H, head_dim, S) --- + k_t_name = f"{prefix}_k_T" + nodes.append(oh.make_node("Transpose", inputs=[k_h], outputs=[k_t_name], perm=[0, 1, 3, 2])) + + # --- Scaled dot-product scores: (B, H, T, head_dim) @ (B, H, head_dim, S) → (B, H, T, S) --- + raw_scores = f"{prefix}_scores_raw" + scaled_scores = f"{prefix}_scores_scaled" + scale_cst = f"{prefix}_attn_scale" + nodes.append(oh.make_node("MatMul", inputs=[q_h, k_t_name], outputs=[raw_scores])) + initializers.append(onh.from_array(np.array(scale_val, dtype=np.float32), name=scale_cst)) + nodes.append(oh.make_node("Mul", inputs=[raw_scores, scale_cst], outputs=[scaled_scores])) + current = scaled_scores + + # --- Softmax input quantization (the MHA enables the softmax's input quantizer) --- + if module.softmax.quantize_input and getattr(module, "enable_quantization", True): + q = module.softmax.input_quantizer + k_q, i_q, f_q = q.get_quantization_bits() + q_nodes, current = quant_fn( + f"{prefix}_attn_score_q", + current, + q.round_mode, + k_q, + i_q, + f_q, + initializers, + overflow_mode=getattr(q, "overflow", "SAT"), + ) + nodes.extend(q_nodes) + + # --- Softmax (dim=-1); approximate_softmax falls back to standard Softmax in ONNX --- + attn_w_name = f"{prefix}_attn_weights" + nodes.append(oh.make_node("Softmax", inputs=[current], outputs=[attn_w_name], axis=-1)) + current = attn_w_name + + # --- Softmax output quantization (the MHA enables the softmax's output quantizer) --- + if module.softmax.quantize_output and getattr(module, "enable_quantization", True): + q = module.softmax.output_quantizer + k_q, i_q, f_q = q.get_quantization_bits() + q_nodes, current = quant_fn( + f"{prefix}_attn_weight_q", + current, + q.round_mode, + k_q, + i_q, + f_q, + initializers, + overflow_mode=getattr(q, "overflow", "SAT"), + ) + nodes.extend(q_nodes) + + # --- Context: (B, H, T, S) @ (B, H, S, head_dim) → (B, H, T, head_dim) --- + # No dedicated quantizer: out_proj's input quantizer (exported with the dense) covers it. + ctx_raw = f"{prefix}_ctx_raw" + nodes.append(oh.make_node("MatMul", inputs=[current, v_h], outputs=[ctx_raw])) + current_ctx = ctx_raw + + # --- Merge heads: (B, H, T, head_dim) → (B, T, E) using dynamic shapes --- + ctx_t = f"{prefix}_ctx_t" # after Transpose → (B, T, H, head_dim) + ctx_shape = f"{prefix}_ctx_shape" + ctx_b_sc = f"{prefix}_ctx_b_sc" + ctx_t_sc = f"{prefix}_ctx_t_sc" + ctx_b_1d = f"{prefix}_ctx_b_1d" + ctx_t_1d = f"{prefix}_ctx_t_1d" + ctx_E_1d = f"{prefix}_ctx_E_1d" + ctx_ax0 = f"{prefix}_ctx_ax0" + ctx_gi0 = f"{prefix}_ctx_gi0" + ctx_gi1 = f"{prefix}_ctx_gi1" + ctx_3d = f"{prefix}_ctx_shape3d" + ctx_merged = f"{prefix}_ctx_merged" + + nodes.append(oh.make_node("Transpose", inputs=[current_ctx], outputs=[ctx_t], perm=[0, 2, 1, 3])) + nodes.append(oh.make_node("Shape", inputs=[ctx_t], outputs=[ctx_shape])) + initializers += [ + onh.from_array(np.array(0, dtype=np.int64), name=ctx_gi0), + onh.from_array(np.array(1, dtype=np.int64), name=ctx_gi1), + onh.from_array(np.array([0], dtype=np.int64), name=ctx_ax0), + onh.from_array(np.array([E], dtype=np.int64), name=ctx_E_1d), + ] + nodes.append(oh.make_node("Gather", inputs=[ctx_shape, ctx_gi0], outputs=[ctx_b_sc])) + nodes.append(oh.make_node("Gather", inputs=[ctx_shape, ctx_gi1], outputs=[ctx_t_sc])) + nodes.append(oh.make_node("Unsqueeze", inputs=[ctx_b_sc, ctx_ax0], outputs=[ctx_b_1d])) + nodes.append(oh.make_node("Unsqueeze", inputs=[ctx_t_sc, ctx_ax0], outputs=[ctx_t_1d])) + nodes.append(oh.make_node("Concat", inputs=[ctx_b_1d, ctx_t_1d, ctx_E_1d], outputs=[ctx_3d], axis=0)) + nodes.append(oh.make_node("Reshape", inputs=[ctx_t, ctx_3d], outputs=[ctx_merged])) + + # --- Output projection (rank-3 input: (B, T, E)) --- + out = _add_dense_nd( + module.out_proj, f"{prefix}_out_proj", ctx_merged, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + + # --- Average attention weights over heads: (B, H, T, S) → (B, T, S) --- + # Emitted so that getitem(mha, 1) has a valid ONNX value name. + avg_attn = f"{prefix}_avg_attn_weights" + nodes.append(oh.make_node("ReduceMean", inputs=[attn_w_name], outputs=[avg_attn], axes=[1], keepdims=0)) + + # --- Optional transpose back for seq-first output --- + if not module.batch_first: + out_final = f"{prefix}_out_seq_first" + nodes.append(oh.make_node("Transpose", inputs=[out], outputs=[out_final], perm=[1, 0, 2])) + return out_final, avg_attn + + return out, avg_attn + + +# --------------------------------------------------------------------------- +# shared module dispatch (used by both sequential and FX converters) +# --------------------------------------------------------------------------- + + +def _emit_module( + module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops=False +): + """Emit ONNX nodes for a single PQuant or standard torch.nn module.""" + if isinstance(module, PQDense): + return _add_dense( + module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops + ) + if isinstance(module, PQConv2d): + return _add_conv( + module, + prefix, + current, + nodes, + initializers, + ndim=2, + quant_fn=quant_fn, + use_qonnx=use_qonnx, + store_integer_weights=store_integer_weights, + ) + if isinstance(module, PQConv1d): + return _add_conv( + module, + prefix, + current, + nodes, + initializers, + ndim=1, + quant_fn=quant_fn, + use_qonnx=use_qonnx, + store_integer_weights=store_integer_weights, + ) + if isinstance(module, (PQBatchNorm2d, PQBatchNorm1d)): + return _add_batchnorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + if isinstance(module, PQLayerNorm): + return _add_layernorm(module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights) + if isinstance(module, PQAvgPool2d): + return _add_avgpool(module, prefix, current, nodes, initializers, ndim=2, quant_fn=quant_fn) + if isinstance(module, PQAvgPool1d): + return _add_avgpool(module, prefix, current, nodes, initializers, ndim=1, quant_fn=quant_fn) + if isinstance(module, nn.ReLU): + out = f"{prefix}_relu" + nodes.append(oh.make_node("Relu", inputs=[current], outputs=[out])) + return out + if isinstance(module, nn.Flatten): + out = f"{prefix}_flatten" + nodes.append(oh.make_node("Flatten", inputs=[current], outputs=[out], axis=module.start_dim)) + return out + if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)): + gamma_name = f"{prefix}_bn_gamma" + beta_name = f"{prefix}_bn_beta" + mean_name = f"{prefix}_bn_mean" + var_name = f"{prefix}_bn_var" + initializers += [ + onh.from_array(module.weight.detach().cpu().numpy().astype(np.float32), name=gamma_name), + onh.from_array(module.bias.detach().cpu().numpy().astype(np.float32), name=beta_name), + onh.from_array(module.running_mean.detach().cpu().numpy().astype(np.float32), name=mean_name), + onh.from_array(module.running_var.detach().cpu().numpy().astype(np.float32), name=var_name), + ] + out = f"{prefix}_bn" + nodes.append( + oh.make_node( + "BatchNormalization", + inputs=[current, gamma_name, beta_name, mean_name, var_name], + outputs=[out], + epsilon=float(module.eps), + ) + ) + return out + if isinstance(module, (nn.Dropout, nn.Dropout2d)): + return current # identity at inference + if isinstance(module, nn.LeakyReLU): + out = f"{prefix}_leakyrelu" + nodes.append(oh.make_node("LeakyRelu", inputs=[current], outputs=[out], alpha=module.negative_slope)) + return out + if isinstance(module, nn.MaxPool2d): + out = f"{prefix}_maxpool" + kernel = module.kernel_size if isinstance(module.kernel_size, (list, tuple)) else [module.kernel_size] * 2 + stride = module.stride if isinstance(module.stride, (list, tuple)) else [module.stride] * 2 + pad = module.padding if isinstance(module.padding, (list, tuple)) else [module.padding] * 2 + nodes.append( + oh.make_node( + "MaxPool", + inputs=[current], + outputs=[out], + kernel_shape=list(kernel), + strides=list(stride), + pads=[pad[0], pad[1], pad[0], pad[1]], + ) + ) + return out + if isinstance(module, nn.Upsample): + # Emit a Resize node with nearest/bilinear mode and scale factors. + roi_name = f"{prefix}_upsample_roi" + scales_name = f"{prefix}_upsample_scales" + initializers.append(onh.from_array(np.array([], dtype=np.float32), name=roi_name)) + scale_factor = module.scale_factor + if isinstance(scale_factor, (int, float)): + scale_factor = (scale_factor, scale_factor) + scales = np.array([1.0, 1.0, float(scale_factor[0]), float(scale_factor[1])], dtype=np.float32) + initializers.append(onh.from_array(scales, name=scales_name)) + mode = "nearest" if module.mode == "nearest" else "linear" + out = f"{prefix}_upsample" + nodes.append( + oh.make_node( + "Resize", + inputs=[current, roi_name, scales_name], + outputs=[out], + mode=mode, + coordinate_transformation_mode="asymmetric", + ) + ) + return out + if isinstance(module, PQActivation): + current = _maybe_quant_input(module, prefix, current, nodes, initializers, quant_fn) + act = module.activation_name + act_out = f"{prefix}_act" + if act == "relu": + nodes.append(oh.make_node("Relu", inputs=[current], outputs=[act_out])) + elif act == "tanh": + nodes.append(oh.make_node("Tanh", inputs=[current], outputs=[act_out])) + elif act == "hard_tanh": + cmin_name = f"{prefix}_htanh_min" + cmax_name = f"{prefix}_htanh_max" + initializers += [ + onh.from_array(np.array(-1.0, dtype=np.float32), name=cmin_name), + onh.from_array(np.array(1.0, dtype=np.float32), name=cmax_name), + ] + nodes.append(oh.make_node("Clip", inputs=[current, cmin_name, cmax_name], outputs=[act_out])) + elif act == "leaky_relu": + nodes.append( + oh.make_node( + "LeakyRelu", inputs=[current], outputs=[act_out], alpha=module.activation_function.negative_slope + ) + ) + elif act == "gelu": + # Decompose so the default opset (13) works; ONNX added a Gelu op only in opset 20. + approximate = getattr(module.activation_function, "approximate", "none") + half_name = f"{prefix}_gelu_half" + one_name = f"{prefix}_gelu_one" + initializers += [ + onh.from_array(np.array(0.5, dtype=np.float32), name=half_name), + onh.from_array(np.array(1.0, dtype=np.float32), name=one_name), + ] + if approximate == "tanh": + # 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) + c0_name = f"{prefix}_gelu_sqrt2_over_pi" + c1_name = f"{prefix}_gelu_c1" + three_name = f"{prefix}_gelu_three" + initializers += [ + onh.from_array(np.array(np.sqrt(2.0 / np.pi), dtype=np.float32), name=c0_name), + onh.from_array(np.array(0.044715, dtype=np.float32), name=c1_name), + onh.from_array(np.array(3.0, dtype=np.float32), name=three_name), + ] + x3 = f"{prefix}_gelu_x3" + cx3 = f"{prefix}_gelu_cx3" + inner = f"{prefix}_gelu_inner" + scaled = f"{prefix}_gelu_scaled" + tanh_out = f"{prefix}_gelu_tanh" + plus_one = f"{prefix}_gelu_plus1" + x_times = f"{prefix}_gelu_xprod" + nodes += [ + oh.make_node("Pow", inputs=[current, three_name], outputs=[x3]), + oh.make_node("Mul", inputs=[x3, c1_name], outputs=[cx3]), + oh.make_node("Add", inputs=[current, cx3], outputs=[inner]), + oh.make_node("Mul", inputs=[inner, c0_name], outputs=[scaled]), + oh.make_node("Tanh", inputs=[scaled], outputs=[tanh_out]), + oh.make_node("Add", inputs=[tanh_out, one_name], outputs=[plus_one]), + oh.make_node("Mul", inputs=[current, plus_one], outputs=[x_times]), + oh.make_node("Mul", inputs=[x_times, half_name], outputs=[act_out]), + ] + else: + # Exact: 0.5 * x * (1 + erf(x / sqrt(2))) + inv_sqrt2_name = f"{prefix}_gelu_inv_sqrt2" + initializers.append(onh.from_array(np.array(1.0 / np.sqrt(2.0), dtype=np.float32), name=inv_sqrt2_name)) + scaled = f"{prefix}_gelu_scaled" + erf_out = f"{prefix}_gelu_erf" + plus_one = f"{prefix}_gelu_plus1" + x_times = f"{prefix}_gelu_xprod" + nodes += [ + oh.make_node("Mul", inputs=[current, inv_sqrt2_name], outputs=[scaled]), + oh.make_node("Erf", inputs=[scaled], outputs=[erf_out]), + oh.make_node("Add", inputs=[erf_out, one_name], outputs=[plus_one]), + oh.make_node("Mul", inputs=[current, plus_one], outputs=[x_times]), + oh.make_node("Mul", inputs=[x_times, half_name], outputs=[act_out]), + ] + else: + raise TypeError(f"PQActivation: unsupported activation {act!r} for ONNX export") + current = act_out + current = _maybe_quant_output(module, prefix, current, nodes, initializers, quant_fn) + return current + if isinstance(module, PQMultiheadAttention): + # Sequential converter: treat as self-attention (Q = K = V = current). + # Returns (out_name, avg_attn_name); expose only the attention output. + out, _ = _add_mha( + module, prefix, current, current, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights + ) + return out + if isinstance(module, Quantizer): + # Standalone quantizer (e.g. an auto-inserted missing quantizer or a + # constant-matrix quantizer): emit a single QDQ node from its k/i/f. + k, i, f = module.get_quantization_bits() + new_nodes, out = quant_fn( + prefix, current, module.round_mode, k, i, f, initializers, overflow_mode=getattr(module, "overflow", "SAT") + ) + nodes.extend(new_nodes) + return out + raise TypeError(f"Unsupported module type for ONNX export: {type(module).__name__}") + + +# --------------------------------------------------------------------------- +# main conversion +# --------------------------------------------------------------------------- + + +def convert_to_onnx( + model: nn.Sequential, + input_shape: tuple, + output_path: str = "model.onnx", + opset: int = 13, + use_qonnx: bool = False, + store_integer_weights: bool = False, + integer_ops: bool = False, + include_clip: bool = True, + batch_size: int | None = None, +) -> onnx.ModelProto: + """ + Convert a Sequential model of PQuant layers to ONNX or QONNX. + + Args: + model: Trained nn.Sequential. Call apply_final_compression() + on all PQ modules before passing here. + input_shape: Shape of a single sample (excluding batch), e.g. (3, 32, 32). + output_path: Where to save the .onnx file. + opset: ONNX opset version (≥13 required for per-channel DequantizeLinear). + use_qonnx: If True, emit QONNX Quant custom nodes (requires qonnx runtime). + If False (default), emit Clip+QuantizeLinear+DequantizeLinear + nodes runnable with plain onnxruntime. + store_integer_weights: If True (and use_qonnx=False), store weight/bias initializers + as int8/uint8 followed by DequantizeLinear instead of float32. + Ignored when use_qonnx=True or integer_ops=True. + integer_ops: If True (and use_qonnx=False), use MatMulInteger for Dense layers + so the inner product runs in int32 arithmetic. Weights are stored + as int8 (pre-transposed) and a single DequantizeLinear converts the + int32 accumulator back to float using the combined scale s_x * s_w. + Implies integer weight storage; store_integer_weights is ignored. + include_clip: Prepend a Clip node before each QuantizeLinear when True (default). + Set to False to emit bare QuantizeLinear+DequantizeLinear pairs — + safe when values are guaranteed in-range at inference time since + QuantizeLinear saturates naturally. Ignored when use_qonnx=True. + batch_size: If not None, fix the batch dimension of all graph inputs and + outputs to this value. If None (default), the batch dimension + is left dynamic. + + Returns: + The constructed onnx.ModelProto. + """ + model.eval() + quant_fn = _quant_node if use_qonnx else functools.partial(_qdq_node, include_clip=include_clip) + + nodes: list[onnx.NodeProto] = [] + initializers: list[onnx.TensorProto] = [] + current = "input" + + for layer_idx, module in enumerate(model): + prefix = f"layer{layer_idx}" + current = _emit_module( + module, prefix, current, nodes, initializers, quant_fn, use_qonnx, store_integer_weights, integer_ops + ) + + with torch.no_grad(): + dummy_out = model(torch.zeros(1, *input_shape)) + batch_dim = batch_size # None → dynamic, int → fixed + output_shape = [batch_dim] + list(dummy_out.shape[1:]) + + batch_dim_vi = oh.make_tensor_value_info("input", TensorProto.FLOAT, [batch_dim, *input_shape]) + output_vi = oh.make_tensor_value_info(current, TensorProto.FLOAT, output_shape) + + graph = oh.make_graph( + nodes=nodes, + name="pquant_onnx", + inputs=[batch_dim_vi], + outputs=[output_vi], + initializer=initializers, + ) + + opset_imports = [oh.make_opsetid("", opset)] + if use_qonnx: + opset_imports.append(oh.make_opsetid("qonnx.custom_op.general", 1)) + model_proto = oh.make_model(graph, opset_imports=opset_imports) + model_proto.ir_version = 6 + + onnx.checker.check_model(model_proto) + onnx.save(model_proto, output_path) + fmt = "QONNX" if use_qonnx else "ONNX (QDQ)" + logging.info("Saved %s model → %s", fmt, output_path) + return model_proto + + +# --------------------------------------------------------------------------- +# Hardware-targeted static-QDQ LayerNormalization graph +# --------------------------------------------------------------------------- + + +def _is_pow2(n: int) -> bool: + return n > 0 and (n & (n - 1)) == 0 + + +def export_qdq_layernorm( + output_path: str, + input_shape, + gamma: np.ndarray, + beta: np.ndarray, + input_scale_log2: int, + output_scale_log2: int, + eps_q0: int = 1, + opset: int = 17, +) -> onnx.ModelProto: + """Build and save a single-LayerNormalization ONNX graph using static QDQ quantization. + + Graph layout:: + + int8 input -> DequantizeLinear -> LayerNormalization -> QuantizeLinear -> DequantizeLinear -> output + + All quantization parameters are explicit float32 initializers (no dynamic + tensors). Per-tensor quantization only; activation zero-points are 0. + + Constraints (validated at build time, not in the graph): + * ``input_shape`` is rank-2 or rank-3 with no dynamic dims. + * Last dim ``D`` is a power of two AND a multiple of 32. + * ``gamma``/``beta`` are 1-D float arrays of length ``D``. + * ``gamma`` is exactly representable as int16 with scale ``2**-7`` (Q7). + * ``beta`` is exactly representable as int16 with scale ``2**-15`` (Q15). + * Input/output scales are exact powers of two, given as log2 exponents. + * ``epsilon = eps_q0 * input_scale**2`` with integer ``eps_q0 >= 1``. + * Normalization axis is the last axis. + + Args: + output_path: Where to save the .onnx file. + input_shape: Static shape of the int8 graph input, e.g. ``(4, 64)`` or ``(1, 4, 64)``. + gamma: Constant gamma initializer, shape ``(D,)``. + beta: Constant beta initializer, shape ``(D,)``. + input_scale_log2: Integer ``a`` with input scale ``= 2**a``. + output_scale_log2: Integer ``b`` with output scale ``= 2**b``. + eps_q0: Positive integer ``>= 1``; ``epsilon = eps_q0 * (2**a)**2``. + opset: ONNX opset version (must be ``>= 17`` for LayerNormalization). + + Returns: + The constructed ``onnx.ModelProto``. + """ + # ----- validate shape ----- + input_shape = tuple(int(d) for d in input_shape) + if len(input_shape) not in (2, 3): + raise ValueError(f"input_shape rank must be 2 or 3, got {len(input_shape)} ({input_shape})") + for d in input_shape: + if d <= 0: + raise ValueError(f"input_shape must be fully static and positive, got {input_shape}") + D = input_shape[-1] + if not _is_pow2(D): + raise ValueError(f"last dim must be a power of two, got {D}") + if D % 32 != 0: + raise ValueError(f"last dim must be a multiple of 32, got {D}") + + # ----- validate gamma / beta ----- + gamma = np.asarray(gamma, dtype=np.float32) + beta = np.asarray(beta, dtype=np.float32) + if gamma.shape != (D,): + raise ValueError(f"gamma must have shape ({D},), got {gamma.shape}") + if beta.shape != (D,): + raise ValueError(f"beta must have shape ({D},), got {beta.shape}") + + GAMMA_F = 7 # Q7 in int16 -> scale = 2**-7 + BETA_F = 15 # Q15 in int16 -> scale = 2**-15 + INT16_MIN, INT16_MAX = -(2**15), 2**15 - 1 + + def _check_q_int16(arr: np.ndarray, frac_bits: int, name: str) -> None: + scaled = arr.astype(np.float64) * (2**frac_bits) + rounded = np.round(scaled) + # Exactly representable: rounding is a no-op (within fp slack). + if not np.allclose(scaled, rounded, atol=1e-4): + raise ValueError( + f"{name} not exactly representable as int16 Q{frac_bits} " + f"(max abs round error = {np.max(np.abs(scaled - rounded)):.6g})" + ) + if rounded.min() < INT16_MIN or rounded.max() > INT16_MAX: + raise ValueError(f"{name} overflows int16 at Q{frac_bits} (range [{rounded.min()}, {rounded.max()}])") + + _check_q_int16(gamma, GAMMA_F, "gamma") + _check_q_int16(beta, BETA_F, "beta") + + # ----- validate quant params ----- + input_scale_log2 = int(input_scale_log2) + output_scale_log2 = int(output_scale_log2) + eps_q0 = int(eps_q0) + if eps_q0 < 1: + raise ValueError(f"eps_q0 must be >= 1, got {eps_q0}") + + if opset < 17: + raise ValueError(f"opset must be >= 17 for LayerNormalization, got {opset}") + + input_scale = float(2.0**input_scale_log2) + output_scale = float(2.0**output_scale_log2) + epsilon = float(eps_q0) * input_scale * input_scale + + # ----- build initializers ----- + initializers = [ + onh.from_array(np.array(input_scale, dtype=np.float32), name="input_scale"), + onh.from_array(np.array(0, dtype=np.int8), name="input_zero_point"), + onh.from_array(np.array(output_scale, dtype=np.float32), name="output_scale"), + onh.from_array(np.array(0, dtype=np.int8), name="output_zero_point"), + onh.from_array(gamma.astype(np.float32), name="gamma"), + onh.from_array(beta.astype(np.float32), name="beta"), + ] + + # ----- build nodes ----- + nodes = [ + oh.make_node( + "DequantizeLinear", + inputs=["input_q", "input_scale", "input_zero_point"], + outputs=["x_dq"], + name="input_dq", + ), + oh.make_node( + "LayerNormalization", + inputs=["x_dq", "gamma", "beta"], + outputs=["ln_out"], + name="layernorm", + axis=-1, + epsilon=epsilon, + ), + oh.make_node( + "QuantizeLinear", + inputs=["ln_out", "output_scale", "output_zero_point"], + outputs=["y_q"], + name="output_q", + ), + oh.make_node( + "DequantizeLinear", + inputs=["y_q", "output_scale", "output_zero_point"], + outputs=["output"], + name="output_dq", + ), + ] + + # ----- build graph + model ----- + input_vi = oh.make_tensor_value_info("input_q", TensorProto.INT8, list(input_shape)) + output_vi = oh.make_tensor_value_info("output", TensorProto.FLOAT, list(input_shape)) + + graph = oh.make_graph( + nodes=nodes, + name="qdq_layernorm", + inputs=[input_vi], + outputs=[output_vi], + initializer=initializers, + ) + + model_proto = oh.make_model(graph, opset_imports=[oh.make_opsetid("", opset)]) + model_proto.ir_version = 8 + + # Strip any initializer names that the onnx library may have added to graph.input. + _init_names = {t.name for t in model_proto.graph.initializer} + _data_inputs = [vi for vi in model_proto.graph.input if vi.name not in _init_names] + del model_proto.graph.input[:] + model_proto.graph.input.extend(_data_inputs) + + onnx.checker.check_model(model_proto) + onnx.save(model_proto, output_path) + return model_proto + + +# --------------------------------------------------------------------------- +# FX-based conversion (supports arbitrary nn.Module topology / skip connections) +# --------------------------------------------------------------------------- + + +class _PQTracer(_fx.Tracer): + """Tracer that treats all PQuant layer types (and standard torch.nn leaves) as atomic.""" + + _LEAF_TYPES = ( + PQDense, + PQConv2d, + PQConv1d, + PQBatchNorm1d, + PQBatchNorm2d, + PQLayerNorm, + PQAvgPool1d, + PQAvgPool2d, + PQMultiheadAttention, + PQActivation, + Quantizer, + ) + + def is_leaf_module(self, m: nn.Module, qualname: str) -> bool: + return isinstance(m, self._LEAF_TYPES) or super().is_leaf_module(m, qualname) + + +def convert_to_onnx_fx( + model: nn.Module, + input_shape: tuple, + output_path: str = "model.onnx", + opset: int = 13, + use_qonnx: bool = False, + store_integer_weights: bool = False, + integer_ops: bool = False, + include_clip: bool = True, +) -> onnx.ModelProto: + """ + Convert any PQuant nn.Module to ONNX using torch.fx symbolic tracing. + + Unlike convert_to_onnx(), this function works with arbitrary model topologies + including residual/skip connections, branches, and concatenations. It requires + the model to be symbolically traceable (no data-dependent control flow). + + Args match convert_to_onnx() exactly; see that function for parameter docs. + """ + model.eval() + quant_fn = _quant_node if use_qonnx else functools.partial(_qdq_node, include_clip=include_clip) + + graph = _PQTracer().trace(model) + gm = _fx.GraphModule(model, graph) + + # ShapeProp populates node.meta["tensor_meta"], which transpose/permute + # need to expand torch's two-arg .transpose(d0, d1) into a full ONNX perm. + from torch.fx.passes.shape_prop import ShapeProp + + # Build the probe tensor on the model's own device so ShapeProp doesn't hit a + # device mismatch when a default device (e.g. CUDA) is set via torch.set_default_device. + device = next((p.device for p in model.parameters()), None) + with torch.no_grad(): + ShapeProp(gm).propagate(torch.zeros(1, *input_shape, device=device)) + + onnx_nodes: list[onnx.NodeProto] = [] + initializers: list[onnx.TensorProto] = [] + node_to_name: dict[_fx.Node, str] = {} + output_names: list[str] = [] + + def _res(arg) -> str: + if isinstance(arg, _fx.Node): + return node_to_name[arg] + raise TypeError(f"Expected fx.Node, got {type(arg)}") + + def _binop_inputs(node: _fx.Node) -> list[str]: + # Like _res for both args, but lifts scalar literals (int/float/bool) + # to float32 initializers so patterns like ``x / 2.0`` work. + names: list[str] = [] + for i, a in enumerate(node.args[:2]): + if isinstance(a, _fx.Node): + names.append(node_to_name[a]) + elif isinstance(a, (int, float, bool)): + cname = f"{node.name}_arg{i}_const" + initializers.append(onh.from_array(np.array(float(a), dtype=np.float32), name=cname)) + names.append(cname) + else: + raise TypeError(f"FX export: unsupported binary-op arg type {type(a).__name__}") + return names + + def _rank(n: _fx.Node) -> int: + meta = n.meta.get("tensor_meta") + if meta is None or not hasattr(meta, "shape"): + raise RuntimeError(f"FX export: ShapeProp did not produce tensor_meta for {n.name!r}") + return len(meta.shape) + + def _swap_perm(rank: int, d0: int, d1: int) -> list[int]: + perm = list(range(rank)) + a, b = d0 % rank, d1 % rank + perm[a], perm[b] = perm[b], perm[a] + return perm + + def _resolve_perm_dims(args, rank: int) -> list[int]: + # Accept both permute(d0, d1, ...) and permute([d0, d1, ...]) shapes. + if len(args) == 1 and isinstance(args[0], (list, tuple)): + dims = args[0] + else: + dims = args + return [int(d) % rank for d in dims] + + for node in gm.graph.nodes: + if node.op == "placeholder": + node_to_name[node] = "input" + + elif node.op == "get_attr": + # Constant tensor attributes — store as initializer on first use. + # Retrieve the actual tensor from the GraphModule. + obj = gm + for part in node.target.split("."): + obj = getattr(obj, part) + attr_name = node.name + if isinstance(obj, torch.Tensor): + initializers.append(onh.from_array(obj.detach().cpu().numpy(), name=attr_name)) + node_to_name[node] = attr_name + + elif node.op == "call_module": + mod = gm.get_submodule(node.target) + mod_prefix = node.name.replace(".", "_") + if isinstance(mod, PQMultiheadAttention): + # node.args = (query, key, value[, key_padding_mask, attn_mask, ...]) + q_name = node_to_name[node.args[0]] + k_name = node_to_name[node.args[1]] if len(node.args) > 1 else q_name + v_name = node_to_name[node.args[2]] if len(node.args) > 2 else q_name + out_name, avg_attn_name = _add_mha( + mod, + mod_prefix, + q_name, + k_name, + v_name, + onnx_nodes, + initializers, + quant_fn, + use_qonnx, + store_integer_weights, + ) + # Store tuple so operator.getitem(node, 0/1) resolves correctly. + node_to_name[node] = (out_name, avg_attn_name) + else: + current = _emit_module( + mod, + mod_prefix, + node_to_name[node.args[0]], + onnx_nodes, + initializers, + quant_fn, + use_qonnx, + store_integer_weights, + integer_ops, + ) + node_to_name[node] = current + + elif node.op == "call_function": + fn = node.target + + if fn is _operator.getitem: + # Unpack a tuple output (e.g. from PQMultiheadAttention). + container = node_to_name[node.args[0]] + if not isinstance(container, tuple): + raise TypeError( + f"operator.getitem on non-tuple node {node.args[0].name!r} is not supported in FX ONNX export" + ) + node_to_name[node] = container[node.args[1]] + continue + + if fn in (torch.add, _operator.add, _operator.iadd): + out = f"{node.name}_add" + onnx_nodes.append(oh.make_node("Add", inputs=_binop_inputs(node), outputs=[out])) + node_to_name[node] = out + + elif fn in (torch.mul, _operator.mul): + out = f"{node.name}_mul" + onnx_nodes.append(oh.make_node("Mul", inputs=_binop_inputs(node), outputs=[out])) + node_to_name[node] = out + + elif fn in (torch.sub, _operator.sub, _operator.isub): + out = f"{node.name}_sub" + onnx_nodes.append(oh.make_node("Sub", inputs=_binop_inputs(node), outputs=[out])) + node_to_name[node] = out + + elif fn in (torch.div, _operator.truediv, _operator.itruediv): + out = f"{node.name}_div" + onnx_nodes.append(oh.make_node("Div", inputs=_binop_inputs(node), outputs=[out])) + node_to_name[node] = out + + elif fn in (torch.matmul, _operator.matmul): + out = f"{node.name}_matmul" + onnx_nodes.append(oh.make_node("MatMul", inputs=_binop_inputs(node), outputs=[out])) + node_to_name[node] = out + + elif fn is torch.transpose: + # torch.transpose(t, d0, d1) swaps two dims; ONNX needs a full perm. + rank = _rank(node.args[0]) + perm = _swap_perm(rank, int(node.args[1]), int(node.args[2])) + out = f"{node.name}_transpose" + onnx_nodes.append(oh.make_node("Transpose", inputs=[_res(node.args[0])], outputs=[out], perm=perm)) + node_to_name[node] = out + + elif fn is torch.permute: + rank = _rank(node.args[0]) + perm = _resolve_perm_dims(node.args[1:], rank) + out = f"{node.name}_permute" + onnx_nodes.append(oh.make_node("Transpose", inputs=[_res(node.args[0])], outputs=[out], perm=perm)) + node_to_name[node] = out + + elif fn is torch.cat: + tensors = [_res(a) for a in node.args[0]] + dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("dim", 0) + out = f"{node.name}_concat" + onnx_nodes.append(oh.make_node("Concat", inputs=tensors, outputs=[out], axis=int(dim))) + node_to_name[node] = out + + elif fn in (_F.relu, torch.relu): + out = f"{node.name}_relu" + onnx_nodes.append(oh.make_node("Relu", inputs=[_res(node.args[0])], outputs=[out])) + node_to_name[node] = out + + elif fn in (_F.sigmoid, torch.sigmoid): + out = f"{node.name}_sigmoid" + onnx_nodes.append(oh.make_node("Sigmoid", inputs=[_res(node.args[0])], outputs=[out])) + node_to_name[node] = out + + elif fn is torch.flatten: + start_dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("start_dim", 0) + out = f"{node.name}_flatten" + onnx_nodes.append(oh.make_node("Flatten", inputs=[_res(node.args[0])], outputs=[out], axis=int(start_dim))) + node_to_name[node] = out + + else: + raise TypeError(f"Unsupported call_function for FX ONNX export: {fn}") + + elif node.op == "call_method": + x = _res(node.args[0]) + + if node.target == "relu": + out = f"{node.name}_relu" + onnx_nodes.append(oh.make_node("Relu", inputs=[x], outputs=[out])) + node_to_name[node] = out + + elif node.target == "flatten": + start_dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("start_dim", 1) + out = f"{node.name}_flatten" + onnx_nodes.append(oh.make_node("Flatten", inputs=[x], outputs=[out], axis=int(start_dim))) + node_to_name[node] = out + + elif node.target in ("view", "reshape"): + shape_vals = [] + for a in node.args[1:]: + if not isinstance(a, int): + raise TypeError("Dynamic reshape (non-constant shape) is not supported in FX ONNX export") + shape_vals.append(a) + shape_name = f"{node.name}_shape" + out = f"{node.name}_reshape" + initializers.append(onh.from_array(np.array(shape_vals, dtype=np.int64), name=shape_name)) + onnx_nodes.append(oh.make_node("Reshape", inputs=[x, shape_name], outputs=[out])) + node_to_name[node] = out + + elif node.target == "transpose": + rank = _rank(node.args[0]) + perm = _swap_perm(rank, int(node.args[1]), int(node.args[2])) + out = f"{node.name}_transpose" + onnx_nodes.append(oh.make_node("Transpose", inputs=[x], outputs=[out], perm=perm)) + node_to_name[node] = out + + elif node.target == "permute": + rank = _rank(node.args[0]) + perm = _resolve_perm_dims(node.args[1:], rank) + out = f"{node.name}_permute" + onnx_nodes.append(oh.make_node("Transpose", inputs=[x], outputs=[out], perm=perm)) + node_to_name[node] = out + + elif node.target == "matmul": + out = f"{node.name}_matmul" + onnx_nodes.append(oh.make_node("MatMul", inputs=[x, _res(node.args[1])], outputs=[out])) + node_to_name[node] = out + + else: + raise TypeError(f"Unsupported call_method for FX ONNX export: {node.target!r}") + + elif node.op == "output": + ret = node.args[0] + rets = list(ret) if isinstance(ret, (tuple, list)) else [ret] + for r in rets: + if not isinstance(r, _fx.Node): + raise TypeError("FX ONNX export: unsupported (non-tensor) model output") + val = node_to_name[r] + # MHA nodes store a tuple (out, avg_attn); expose the attention output. + output_names.append(val[0] if isinstance(val, tuple) else val) + + with torch.no_grad(): + dummy_out = model(torch.zeros(1, *input_shape, device=device)) + dummy_outs = list(dummy_out) if isinstance(dummy_out, (tuple, list)) else [dummy_out] + + batch_dim = oh.make_tensor_value_info("input", TensorProto.FLOAT, [None, *input_shape]) + output_vis = [ + oh.make_tensor_value_info(name, TensorProto.FLOAT, [None] + list(t.shape[1:])) + for name, t in zip(output_names, dummy_outs) + ] + + onnx_graph = oh.make_graph( + nodes=onnx_nodes, + name="pquant_onnx_fx", + inputs=[batch_dim], + outputs=output_vis, + initializer=initializers, + ) + + opset_imports = [oh.make_opsetid("", opset)] + if use_qonnx: + opset_imports.append(oh.make_opsetid("qonnx.custom_op.general", 1)) + model_proto = oh.make_model(onnx_graph, opset_imports=opset_imports) + model_proto.ir_version = 6 + + onnx.checker.check_model(model_proto) + onnx.save(model_proto, output_path) + fmt = "QONNX" if use_qonnx else "ONNX (QDQ)" + logging.info("Saved %s model (FX) → %s", fmt, output_path) + return model_proto + + +# --------------------------------------------------------------------------- +# usage example +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + import onnxruntime as ort + + import pquant + + cfg = pquant.cs_config() + cfg.quantization_parameters.granularity = "per-channel" + + model = nn.Sequential( + PQConv2d(cfg, in_channels=3, out_channels=16, kernel_size=3, padding=1), + PQBatchNorm2d(cfg, num_features=16), + nn.ReLU(), + PQAvgPool2d(cfg, kernel_size=2, stride=2), + nn.Flatten(), + PQDense(cfg, in_features=16 * 16 * 16, out_features=64), + nn.ReLU(), + PQDense(cfg, in_features=64, out_features=10), + ) + + x = torch.randn(4, 3, 32, 32) + with torch.no_grad(): + model(x) + + for module in model.modules(): + if hasattr(module, "apply_final_compression"): + module.apply_final_compression() + + model.eval() + with torch.no_grad(): + torch_out = model(x).numpy() + + qonnx_path = "model_qonnx.onnx" + onnx_path = "model_qdq.onnx" + convert_to_onnx(model, input_shape=(3, 32, 32), output_path=qonnx_path, use_qonnx=True) + convert_to_onnx(model, input_shape=(3, 32, 32), output_path=onnx_path, use_qonnx=False) + + from qonnx.core.modelwrapper import ModelWrapper + from qonnx.core.onnx_exec import execute_onnx + from qonnx.transformation.infer_shapes import InferShapes + + qmodel = ModelWrapper(qonnx_path) + qmodel.graph.input[0].type.tensor_type.shape.dim[0].dim_value = x.shape[0] + qmodel = qmodel.transform(InferShapes()) + input_name = qmodel.graph.input[0].name + output_name = qmodel.graph.output[0].name + qonnx_out = execute_onnx(qmodel, {input_name: x.numpy()})[output_name] + + sess = ort.InferenceSession(onnx_path) + onnx_out = sess.run(None, {sess.get_inputs()[0].name: x.numpy()})[0] + + print(f"\n{'':=<55}") # noqa: T201 + print(f" max |torch - qonnx| : {np.abs(torch_out - qonnx_out).max():.6f}") # noqa: T201 + print(f" max |torch - onnx| : {np.abs(torch_out - onnx_out).max():.6f}") # noqa: T201 + print(f" max |qonnx - onnx| : {np.abs(qonnx_out - onnx_out).max():.6f}") # noqa: T201 + print(f"{'':=<55}") # noqa: T201 diff --git a/src/pquant/core/torch/distillers.py b/src/pquant/core/torch/distillers.py new file mode 100644 index 0000000..dcb11ca --- /dev/null +++ b/src/pquant/core/torch/distillers.py @@ -0,0 +1,686 @@ +from __future__ import annotations + +import os +import tempfile +from collections.abc import Callable, Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data + +from pquant.core.torch.layers import ( + PQWeightBiasBase, + get_model_losses, + post_epoch_functions, + post_pretrain_functions, + post_round_functions, + pre_epoch_functions, + pre_finetune_functions, +) + + +def get_module(model: nn.Module, name: str) -> nn.Module: + for n, m in model.named_modules(): + if n == name: + return m + raise ValueError(f"Module '{name}' not found in model") + + +def pq_layer_names(model: nn.Module) -> list[str]: + """Return names of all PQWeightBiasBase submodules in forward order.""" + return [name for name, m in model.named_modules() if isinstance(m, PQWeightBiasBase)] + + +class CachedDataset(torch.utils.data.Dataset): + """Per-batch cache backed by files in a temporary directory. + + Each file stores one ``(input, output)`` batch as a tuple. + Since each file has one batch, uses batch_size of 1 here. + """ + + def __init__(self, cache_dir: str, n_batches: int) -> None: + self.cache_dir = cache_dir + self.n_batches = n_batches + + def __len__(self) -> int: + return self.n_batches + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: + return torch.load( + os.path.join(self.cache_dir, f"{idx:08d}.pt"), + weights_only=True, + ) + + +class LayerwiseDistiller: + """ + Distills a teacher model into a student model one PQ layer at a time. + + For each layer, all student parameters are frozen except that layer's, + and the layer is trained to match the teacher's outputs at that point. + + Args: + teacher: Reference model. Will be set to eval() and + gradients disabled during distillation. + student: PQuantML model. + loss_fn: Loss function between student and teacher activations. + precompute_layer_inputs: If True, run one pass over the dataloader at + the start of ``distill_layer``, capturing the teacher + layer's input and output. Pairs + are saved to a temporary directory on disk. All epoch loops + then call ``student_layer(layer_input)`` directly, skipping the full model forward entirely. + prefetch_workers: Number of DataLoader worker processes used to + prefetch cached batches from disk when + ``precompute_layer_inputs=True``. + cache_dir: Directory under which temporary per-layer cache + subdirectories are created when ``precompute_layer_inputs=True``. + """ + + def __init__( + self, + teacher: nn.Module, + student: nn.Module, + device: torch.device | None, + loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + precompute_layer_inputs: bool = True, + prefetch_workers: int = 2, + cache_dir: str | None = None, + ): + self.teacher = teacher + self.student = student + self.device = device + self.loss_fn = loss_fn or F.mse_loss + self.precompute_layer_inputs = precompute_layer_inputs + self.prefetch_workers = prefetch_workers + self.cache_dir = cache_dir + + self.teacher.eval() + for param in self.teacher.parameters(): + param.requires_grad_(False) + + def precompute_layer_io( + self, + teacher_layer: nn.Module, + dataloader: Iterable, + cache_dir: str, + ) -> torch.utils.data.DataLoader: + """Cache ``(layer_input, layer_output)``. + + Runs one teacher pass and one frozen student pass per batch, saving + each pair. Returns a DataLoader which uses this saved data. + """ + teacher_captured_data: dict[str, torch.Tensor] = {} + + def teacher_post(m: nn.Module, inp: tuple, out: torch.Tensor) -> None: + teacher_output = out[0] if isinstance(out, tuple) else out + teacher_captured_data["out"] = teacher_output.detach().cpu() + + def teacher_pre(m: nn.Module, inp: tuple) -> None: + teacher_captured_data["inp"] = inp[0].detach().cpu() + + teacher_output_hook = teacher_layer.register_forward_hook(teacher_post) + teacher_input_hook = teacher_layer.register_forward_pre_hook(teacher_pre) + + n_batches = 0 + try: + with torch.no_grad(): + for x, _ in dataloader: + if self.device is not None: + x = x.to(self.device) + self.teacher(x) + torch.save( + (teacher_captured_data["inp"], teacher_captured_data["out"]), + os.path.join(cache_dir, f"{n_batches:08d}.pt"), + ) + n_batches += 1 + finally: + teacher_output_hook.remove() + teacher_input_hook.remove() + + dataset = CachedDataset(cache_dir, n_batches) + return torch.utils.data.DataLoader( + dataset, + batch_size=1, + shuffle=True, + num_workers=self.prefetch_workers, + collate_fn=lambda b: b[0], + ) + + def val_layer_loss( + self, + teacher_layer: nn.Module, + student_layer: nn.Module, + val_dataloader: Iterable, + ) -> float: + """Compute mean validation loss for a single layer (no grad, eval mode). + + Always runs a full model forward pass with hooks — the val dataset is + never cached to disk, avoiding large intermediate feature map storage. + """ + t_captured: dict[str, torch.Tensor] = {} + s_captured: dict[str, torch.Tensor] = {} + + def teacher_post(m: nn.Module, inp: tuple, out) -> None: + t_captured["out"] = (out[0] if isinstance(out, tuple) else out).detach() + + def student_hook(m: nn.Module, inp: tuple, out) -> None: + s_captured["out"] = out[0] if isinstance(out, tuple) else out + + h_t = teacher_layer.register_forward_hook(teacher_post) + h_s = student_layer.register_forward_hook(student_hook) + + self.student.eval() + batch_losses: list[float] = [] + try: + with torch.no_grad(): + for x, _ in val_dataloader: + if self.device is not None: + x = x.to(self.device) + self.teacher(x) + self.student(x) + batch_losses.append(self.loss_fn(s_captured["out"], t_captured["out"]).item()) + finally: + h_t.remove() + h_s.remove() + self.student.train() + + return sum(batch_losses) / len(batch_losses) + + def distill_layer( + self, + teacher_layer_name: str, + student_layer_name: str, + dataloader: Iterable, + optimizer_factory: Callable[[Iterable], torch.optim.Optimizer], + n_epochs: int, + val_dataloader: Iterable | None = None, + epoch_callback: Callable[[int, float, float | None, str], None] | None = None, + ) -> list[float]: + """ + Distill a single layer. + + Args: + teacher_layer_name: Named path to the module in the teacher, e.g. ``"layers.0"``. + student_layer_name: Named path to the corresponding module in the student. + dataloader: Iterable of ``(inputs, targets)`` batches. + optimizer_factory: Callable that accepts an iterable of parameters and + returns a fresh optimizer. Called once per layer with only + the target layer's parameters, so state never carries over + between layers. Example: + ``lambda params: torch.optim.Adam(params, lr=1e-4)``. + n_epochs: Number of full passes through the dataloader. + ``post_epoch_functions`` is called on the layer after + each pass. + val_dataloader: Optional validation dataloader. When provided, a + no-grad eval pass is run after every training epoch and + the resulting mean loss is passed to ``epoch_callback``. + epoch_callback: Called after each epoch with + ``(epoch, train_loss, val_loss, student_layer_name)`` + where ``val_loss`` is ``None`` when no ``val_dataloader`` + is given. Use this to drive per-epoch side effects such + as calling ``student.increment_alpha()``, stepping a + scheduler, or logging metrics. + + Returns: + List of per-epoch mean training losses. + """ + teacher_layer = get_module(self.teacher, teacher_layer_name) + student_layer = get_module(self.student, student_layer_name) + + # Freeze everything except the target layer (skip non-float params + # like bool pruning masks which cannot require gradients) + frozen: dict[str, bool] = {} + for name, param in self.student.named_parameters(): + if not param.is_floating_point(): + continue + frozen[name] = param.requires_grad + param.requires_grad_(name.startswith(student_layer_name)) + + optimizer = optimizer_factory(student_layer.parameters()) + + tmpdir = None + hooks = [] + epoch_losses: list[float] = [] + self.student.train() + + try: + if self.precompute_layer_inputs: + tmpdir = tempfile.TemporaryDirectory(prefix="ldistil_", dir=self.cache_dir or os.getcwd()) + dataloader = self.precompute_layer_io(teacher_layer, dataloader, self.device, tmpdir.name) + else: + teacher_out: dict[str, torch.Tensor] = {} + student_out: dict[str, torch.Tensor] = {} + + def _teacher_hook(m: nn.Module, inp: tuple, out) -> None: + t = out[0] if isinstance(out, tuple) else out + teacher_out["out"] = t.detach() + + def student_hook(m: nn.Module, inp: tuple, out) -> None: + t = out[0] if isinstance(out, tuple) else out + student_out["out"] = t + + h_t = teacher_layer.register_forward_hook(_teacher_hook) + h_s = student_layer.register_forward_hook(student_hook) + hooks = [h_t, h_s] + + for epoch in range(n_epochs): + if getattr(student_layer, "enable_pruning", False): + student_layer.pruning_layer.pre_epoch_function(epoch, n_epochs) + batch_losses: list[float] = [] + + if self.precompute_layer_inputs: + for t_in, t_out in dataloader: + if self.device is not None: + t_in, t_out = t_in.to(self.device), t_out.to(self.device) + raw = student_layer(t_in) + s_out = raw[0] if isinstance(raw, tuple) else raw + loss = self.loss_fn(s_out, t_out) + optimizer.zero_grad() + loss.backward() + optimizer.step() + batch_losses.append(loss.item()) + else: + for x, _ in dataloader: + if self.device is not None: + x = x.to(self.device) + with torch.no_grad(): + self.teacher(x) + self.student(x) + t_out, s_out = teacher_out["out"], student_out["out"] + loss = self.loss_fn(s_out, t_out) + optimizer.zero_grad() + loss.backward() + optimizer.step() + batch_losses.append(loss.item()) + + mean_loss = sum(batch_losses) / len(batch_losses) + epoch_losses.append(mean_loss) + if getattr(student_layer, "enable_pruning", False): + student_layer.pruning_layer.post_epoch_function(epoch, n_epochs) + val_loss: float | None = None + if val_dataloader is not None: + val_loss = self.val_layer_loss(teacher_layer, student_layer, val_dataloader) + if epoch_callback is not None: + epoch_callback(epoch, mean_loss, val_loss, student_layer_name) + finally: + for h in hooks: + h.remove() + if tmpdir is not None: + tmpdir.cleanup() + for name, param in self.student.named_parameters(): + if name in frozen: + param.requires_grad_(frozen[name]) + + return epoch_losses + + def distill_all( + self, + dataloader: Iterable, + optimizer_factory: Callable[[Iterable], torch.optim.Optimizer], + n_epochs: int, + layer_names: list[tuple[str, str]] | None = None, + val_dataloader: Iterable | None = None, + epoch_callback: Callable[[int, float, float | None, str], None] | None = None, + layer_callback: Callable[[str, str, list[float]], None] | None = None, + ) -> dict[tuple[str, str], list[float]]: + """ + Distill all PQ layers sequentially front-to-back. + + Args: + dataloader: Iterable of ``(inputs, targets)`` batches. + optimizer_factory: Callable ``(params) -> Optimizer``. Called fresh + for each layer so state never carries over. + n_epochs: Number of epochs (full dataloader passes) per layer. + layer_names: List of ``(teacher_layer_name, student_layer_name)`` + tuples to distill. Defaults to pairing each + ``PQWeightBiasBase`` submodule name with itself. + val_dataloader: Optional validation dataloader passed through to + each ``distill_layer`` call. + epoch_callback: Passed through to ``distill_layer``. Called after + each epoch with + ``(epoch, train_loss, val_loss, layer_name)`` + where ``val_loss`` is ``None`` when no + ``val_dataloader`` is given. + layer_callback: Called after each layer completes with + ``(teacher_layer_name, student_layer_name, losses)``. + + Returns: + Dict mapping ``(teacher_layer_name, student_layer_name)`` to list + of per-epoch mean losses. + """ + if layer_names is not None: + pairs = layer_names + else: + names = pq_layer_names(self.student) + pairs = [(n, n) for n in names] + history: dict[tuple[str, str], list[float]] = {} + + for t_name, s_name in pairs: + losses = self.distill_layer( + t_name, + s_name, + dataloader, + optimizer_factory, + n_epochs, + val_dataloader=val_dataloader, + epoch_callback=epoch_callback, + ) + history[(t_name, s_name)] = losses + if layer_callback is not None: + layer_callback(t_name, s_name, losses) + + return history + + +class ModelDistiller: + """ + Knowledge distillation at the model output level. + + The teacher's logits are used as soft targets for the student via KL + divergence loss. Because the teacher (16 raw classes) and student (10 + merged classes) have different output sizes, a raw-to-merged mapping + tensor is used to collapse the teacher logits to 10 classes before + computing the loss. + + Args: + teacher: Full-precision reference model (16-class output). Set to + eval() with gradients disabled. + student: Student model (10-class output). + teacher_transform: Optional callable applied to teacher logits before + the distillation loss is computed. Use this to map + teacher outputs to the student's output space, e.g. to + collapse classes. If ``None``, teacher logits are used + as-is (teacher and student must then have matching output + shapes). + loss_fn: Which distillation loss to use. One of: + ``"kl_ce"`` (default) — ``alpha * KL(student || teacher) + + (1 - alpha) * CE(student, labels)`` with temperature + scaling. Classic Hinton et al. formulation. + ``"kl"`` — KL divergence only, no hard label term. + Useful when ground-truth labels are unavailable or noisy. + ``"mse"`` — mean squared error directly on the logits, + no temperature scaling. Simpler baseline; equivalent to + KL under a uniform teacher distribution. + temperature: Softmax temperature for soft targets (default 4.0). + Only used by ``"kl_ce"`` and ``"kl"``. + alpha: Weight of the KL distillation loss (default 0.7). + The remaining ``(1 - alpha)`` weight is given to the + hard cross-entropy loss. Only used by ``"kl_ce"``. + precompute_teacher_outputs: If True (default), run a single inference + pass of the teacher over the entire dataset before + distillation begins, caching transformed teacher logits + to disk. All epoch loops then + use this cached loader and never call the teacher again. + Set to False only when the dataloader uses stochastic + augmentations that must remain live during training. + prefetch_workers: Number of DataLoader worker processes used to + prefetch cached batches from disk when + ``precompute_teacher_outputs=True``. Defaults to 2. + cache_dir: Directory under which the temporary cache is created + when ``precompute_teacher_outputs=True``. Defaults to + the current working directory. Avoid ``/tmp`` on Linux + as it is typically a RAM-backed ``tmpfs``. + """ + + LOSS_FN_OPTIONS = ("kl_ce", "kl", "mse") + + def __init__( + self, + teacher: nn.Module, + student: nn.Module, + pq_config, + device: torch.device | None, + teacher_transform: Callable[[torch.Tensor], torch.Tensor] | None = None, + loss_fn: str = "kl_ce", + temperature: float = 4.0, + alpha: float = 0.7, + precompute_teacher_outputs: bool = True, + prefetch_workers: int = 2, + cache_dir: str | None = None, + ): + if loss_fn not in self.LOSS_FN_OPTIONS: + raise ValueError(f"loss_fn must be one of {self.LOSS_FN_OPTIONS}, got '{loss_fn}'") + + self.teacher = teacher + self.student = student + self.pq_config = pq_config + self.teacher_transform = teacher_transform + self.loss_fn = loss_fn + self.T = temperature + self.alpha = alpha + self._precompute_teacher_outputs = precompute_teacher_outputs + self.prefetch_workers = prefetch_workers + self.cache_dir = cache_dir + self.device = device + self.teacher.eval() + for p in self.teacher.parameters(): + p.requires_grad_(False) + + def precompute_teacher_outputs( + self, + dataloader: Iterable, + cache_dir: str, + shuffle: bool = True, + ) -> torch.utils.data.DataLoader: + """Run one teacher inference pass and cache ``(x, teacher_output, y)`` to disk.""" + n_batches = 0 + self.teacher.eval() + with torch.no_grad(): + for x, y in dataloader: + if self.device is not None: + x = x.to(self.device) + teacher_logits = self.teacher(x) + teacher_out = self.teacher_transform(teacher_logits) if self.teacher_transform else teacher_logits + torch.save( + (x.cpu(), teacher_out.cpu(), y.cpu()), + os.path.join(cache_dir, f"{n_batches:08d}.pt"), + ) + n_batches += 1 + + dataset = CachedDataset(cache_dir, n_batches) + return torch.utils.data.DataLoader( + dataset, + batch_size=1, + shuffle=shuffle, + num_workers=self.prefetch_workers, + collate_fn=lambda b: b[0], + ) + + def kl_divergence(self, student_logits: torch.Tensor, teacher_merged: torch.Tensor) -> torch.Tensor: + _, C, _, _ = student_logits.shape + flat_s = student_logits.permute(0, 2, 3, 1).reshape(-1, C) + flat_t = teacher_merged.permute(0, 2, 3, 1).reshape(-1, C) + log_p_s = F.log_softmax(flat_s / self.T, dim=1) + p_t = F.softmax(flat_t / self.T, dim=1) + return F.kl_div(log_p_s, p_t, reduction="batchmean") * (self.T**2) + + def loss_kl_ce( + self, + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + labels: torch.Tensor, + ) -> torch.Tensor: + """alpha * KL(student || teacher) + (1 - alpha) * CE(student, labels).""" + kl_loss = self.kl_divergence(student_logits, teacher_logits) + ce_loss = F.cross_entropy(student_logits, labels, ignore_index=-1) + return self.alpha * kl_loss + (1 - self.alpha) * ce_loss + + def loss_kl( + self, + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + ) -> torch.Tensor: + """KL divergence only, no hard label term.""" + return self.kl_divergence(student_logits, teacher_logits) + + def loss_mse( + self, + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + ) -> torch.Tensor: + """MSE directly on logits, no temperature scaling.""" + return F.mse_loss(student_logits, teacher_logits) + + def compute_loss( + self, + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + labels: torch.Tensor, + ) -> torch.Tensor: + if self.loss_fn == "kl_ce": + return self.loss_kl_ce(student_logits, teacher_logits, labels) + elif self.loss_fn == "kl": + return self.loss_kl(student_logits, teacher_logits) + else: # mse + return self.loss_mse(student_logits, teacher_logits) + + def run_val_epoch( + self, + val_dataloader: Iterable, + ) -> float: + """Compute mean validation loss over one pass of val_dataloader (no grad, eval mode). + + When ``precompute_teacher_outputs=True`` the dataloader yields + ``(x, teacher_logits, y)`` and the cached teacher output is used directly. + Otherwise it yields ``(x, y)`` and the teacher is run live. + """ + self.student.eval() + batch_losses: list[float] = [] + with torch.no_grad(): + for batch in val_dataloader: + if self._precompute_teacher_outputs: + x, teacher_logits, y = batch + if self.device is not None: + x, teacher_logits, y = x.to(self.device), teacher_logits.to(self.device), y.to(self.device) + else: + x, y = batch + if self.device is not None: + x, y = x.to(self.device), y.to(self.device) + teacher_logits = self.teacher(x) + teacher_logits = self.teacher_transform(teacher_logits) if self.teacher_transform else teacher_logits + student_logits = self.student(x) + loss = self.compute_loss(student_logits, teacher_logits, y) + batch_losses.append(loss.item()) + self.student.train() + return sum(batch_losses) / len(batch_losses) + + def run_epoch( + self, + dataloader: Iterable, + optimizer: torch.optim.Optimizer, + ) -> float: + self.student.train() + batch_losses: list[float] = [] + + for batch in dataloader: + if self._precompute_teacher_outputs: + x, teacher_logits, y = batch + if self.device is not None: + x, teacher_logits, y = x.to(self.device), teacher_logits.to(self.device), y.to(self.device) + else: + x, y = batch + if self.device is not None: + x, y = x.to(self.device), y.to(self.device) + with torch.no_grad(): + teacher_logits = self.teacher(x) + teacher_logits = self.teacher_transform(teacher_logits) if self.teacher_transform else teacher_logits + + student_logits = self.student(x) + loss = self.compute_loss(student_logits, teacher_logits, y) + loss = get_model_losses(self.student, loss) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + batch_losses.append(loss.item()) + + return sum(batch_losses) / len(batch_losses) + + def distill( + self, + dataloader: Iterable, + optimizer: torch.optim.Optimizer, + val_dataloader: Iterable | None = None, + epoch_callback: Callable[[int, float, float | None], None] | None = None, + ) -> list[float]: + """ + Run model distillation following the PQuantML training pipeline. + + Args: + dataloader: Dataloader provided by user. + optimizer: Optimizer for student parameters. + val_dataloader: Optional validation dataloader. When provided, a + no-grad eval pass is run after every training epoch + and the resulting mean loss is passed to + ``epoch_callback``. + epoch_callback: Called after each epoch with + ``(epoch, train_loss, val_loss)`` where + ``val_loss`` is ``None`` when no ``val_dataloader`` + is given. + + Returns: + List of per-epoch mean training losses. + """ + training_parameters = self.pq_config.training_parameters + + tmpdir = None + tmpdir_val = None + epoch_losses: list[float] = [] + global_epoch = 0 + + try: + if self._precompute_teacher_outputs: + tmpdir = tempfile.TemporaryDirectory(prefix="mdistil_", dir=self.cache_dir or os.getcwd()) + dataloader = self.precompute_teacher_outputs(dataloader, tmpdir.name) + if val_dataloader is not None: + tmpdir_val = tempfile.TemporaryDirectory(prefix="mdistil_val_", dir=self.cache_dir or os.getcwd()) + val_dataloader = self.precompute_teacher_outputs(val_dataloader, tmpdir_val.name, shuffle=False) + + for e in range(training_parameters.pretraining_epochs): + pre_epoch_functions(self.student, e, training_parameters.pretraining_epochs) + mean_loss = self.run_epoch(dataloader, optimizer) + epoch_losses.append(mean_loss) + post_epoch_functions(self.student, e, training_parameters.pretraining_epochs) + val_loss: float | None = None + if val_dataloader is not None: + val_loss = self.run_val_epoch(val_dataloader) + if epoch_callback is not None: + epoch_callback(global_epoch, mean_loss, val_loss) + global_epoch += 1 + + post_pretrain_functions(self.student, self.pq_config, train_loader=dataloader) + + for _ in range(training_parameters.rounds): + for e in range(training_parameters.epochs): + pre_epoch_functions(self.student, e, training_parameters.epochs) + mean_loss = self.run_epoch(dataloader, optimizer) + epoch_losses.append(mean_loss) + post_epoch_functions(self.student, e, training_parameters.epochs) + val_loss = None + if val_dataloader is not None: + val_loss = self.run_val_epoch(val_dataloader) + if epoch_callback is not None: + epoch_callback(global_epoch, mean_loss, val_loss) + global_epoch += 1 + post_round_functions(self.student) + + pre_finetune_functions(self.student) + for e in range(training_parameters.fine_tuning_epochs): + pre_epoch_functions(self.student, e, training_parameters.fine_tuning_epochs) + mean_loss = self.run_epoch(dataloader, optimizer) + epoch_losses.append(mean_loss) + val_loss = None + if val_dataloader is not None: + val_loss = self.run_val_epoch(val_dataloader) + if epoch_callback is not None: + epoch_callback(global_epoch, mean_loss, val_loss) + post_epoch_functions(self.student, e, training_parameters.fine_tuning_epochs) + global_epoch += 1 + finally: + if tmpdir is not None: + tmpdir.cleanup() + if tmpdir_val is not None: + tmpdir_val.cleanup() + + return epoch_losses diff --git a/src/pquant/core/torch/fit_compress.py b/src/pquant/core/torch/fit_compress.py index d99e4fb..fb67528 100644 --- a/src/pquant/core/torch/fit_compress.py +++ b/src/pquant/core/torch/fit_compress.py @@ -174,7 +174,7 @@ def print_info_bits(model): config, ) # Now add the layer specific configuration to the model - # add_layer_specific_quantization_to_model(trained_uncompressed_model, config) + # _add_layer_specific_quantization_to_model(trained_uncompressed_model, config) logging.info("Layerwise quantization bits after FITcompress : ", config.quantization_parameters.layer_specific) @@ -184,7 +184,6 @@ def print_info_bits(model): class node: - def __init__( self, matrices_params_layerwise, @@ -232,7 +231,7 @@ def __init__( self.unquantized_weights = unquantized_weights self.int_bits = int_bits self.frac_bits = frac_bits - self.key = ''.join(random.choices(string.ascii_uppercase + string.digits, k=20)) + self.key = "".join(random.choices(string.ascii_uppercase + string.digits, k=20)) def extract_config_from_node(self, layer_names): """ @@ -250,13 +249,12 @@ def extract_config_from_node(self, layer_names): layer_name: [i_bits, f_bits] for layer_name, i_bits, f_bits in zip(layer_names, self.int_bits, self.frac_bits) } - config = {'quant_config': quant_config, 'pruning_metrics': self.pruning_metrics} + config = {"quant_config": quant_config, "pruning_metrics": self.pruning_metrics} return config class FITcompress: - def __init__(self, model, device, dataloader, criterion, config, layerwise_pruning=False, input_shape=None): """ Calculate initial EF of the uncompressed model and set up quantization & @@ -320,7 +318,7 @@ def __init__(self, model, device, dataloader, criterion, config, layerwise_pruni ) # for N:M pruning in Wanda, use 50% pruning cap during FITcompress - if self.config.pruning_parameters.pruning_method == 'wanda' and type(self.config.pruning_parameters.N) is int: + if self.config.pruning_parameters.pruning_method == "wanda" and type(self.config.pruning_parameters.N) is int: self.pruning_schedule = 0.5 * ( 1 - np.logspace( @@ -333,15 +331,15 @@ def __init__(self, model, device, dataloader, criterion, config, layerwise_pruni # Dictionary structure allows us to possibly iterate over multiple different pruning metrics # but currently only one as in FITcompress, the target pruning sparsity, i.e. percentage - pruning_metrics = {'percentage': 0} + pruning_metrics = {"percentage": 0} # If we want to find sparsity targets per layer (not part of FITcompress paper) if layerwise_pruning: self.pruning_schedulers_layerwise = self.get_pruning_schedulers_layer_specific( - matrices_params_layerwise, None, mode='fit' + matrices_params_layerwise, None, mode="fit" ) # Add the layer-specific starting pruning percentages to the current metric - pruning_metrics = pruning_metrics | {f'{self.layer_names[i]}_percentage': 0 for i in range(self.n_layers)} + pruning_metrics = pruning_metrics | {f"{self.layer_names[i]}_percentage": 0 for i in range(self.n_layers)} # Initialize the first node in the compression space self.initial_node = node( @@ -368,7 +366,7 @@ def __init__(self, model, device, dataloader, criterion, config, layerwise_pruni # Intialize a list to store nodes that can be traversed during the path finding process self.potential_nodes = [self.initial_node] - def get_pruning_schedulers_layer_specific(self, matrices_params_layerwise, global_sparsity_scheduler, mode='fit'): + def get_pruning_schedulers_layer_specific(self, matrices_params_layerwise, global_sparsity_scheduler, mode="fit"): """ Calculates layer-specific pruning schedulers. The idea is that layers with weights that are not that much affected by pertubation should be pruned more/faster than layers with weights @@ -383,7 +381,7 @@ def get_pruning_schedulers_layer_specific(self, matrices_params_layerwise, globa """ schedulers = {} - if mode == 'fit': + if mode == "fit": # Get the layer-wise FIT scores of the initial model _, FIT_layerwise = self.fit_computer.get_FIT_old( FeM=self.FeM, params_after=matrices_params_layerwise, same_theta=True @@ -396,7 +394,6 @@ def get_pruning_schedulers_layer_specific(self, matrices_params_layerwise, globa max_importance = max(FIT_layerwise_summed) for layer_idx, importance in enumerate(FIT_layerwise_summed): - # Scale importance between 0 and 1 importance_ratio = (importance - min_importance) / (max_importance - min_importance) @@ -452,7 +449,7 @@ def assign_parameters(self, model, params): for _, module in model.named_modules(): if isinstance(module, (PQDense, PQConv2d)): for name_param, matrix_param in list(module.named_parameters()): - if name_param.endswith('_weight'): + if name_param.endswith("_weight"): matrix_param.data = nn.parameter.Parameter(params[i].to(self.device)) matrix_param.collect = True i += 1 @@ -643,7 +640,6 @@ def add_pruning_layer_specific(self, current_node, pruning_metrics, layer_idx=No current_node_matrices_params_layerwise = [] # Now iterate through all layers for idx, curr_pruning_percentage in enumerate(pruning_metrics.values()): - if idx == 0: # Global pruning percentage continue @@ -768,7 +764,7 @@ def astar(self, config): """ iterations = 0 while len(self.potential_nodes) > 0 and iterations < 1000: - logging.info(f'Iteration : {iterations} ') + logging.info(f"Iteration : {iterations} ") next_best_node = None @@ -777,7 +773,6 @@ def astar(self, config): for p_node in self.potential_nodes: # If we find a node with wanted compression rate, we can return it and stop the A* algorithm if p_node.curr_compression_rate < self.compression_goal: - logging.info( f"Optimal node found with full distance {p_node.full_dist}, " f"compression rate {p_node.curr_compression_rate}, " @@ -799,10 +794,6 @@ def astar(self, config): self.assign_parameters(self.model, params_quantized_unpruned) - self.post_fitcompress_calibration( - p_node.extract_config_from_node(self.layer_names)['quant_config'], config - ) - return ( p_node, p_node.extract_config_from_node(self.layer_names), @@ -866,7 +857,6 @@ def create_neighbours(self, current_node): """ if self.config.fitcompress_parameters.approximate: - # Update FeM for the best node and use it when creating the neighbours for quantization. # This leads to num_layers less FIT calculations, as we do not need to calculate the FeM again, # which reduces runtime @@ -889,7 +879,6 @@ def create_neighbours(self, current_node): logging.info("Current node states for quantization & pruning: ", current_node_state) if self.config.fitcompress_parameters.optimize_quantization: for layer_idx in range(self.n_layers): - # Set neighbour state to current state neighbour_node_state = current_node_state.copy() @@ -934,7 +923,6 @@ def create_neighbours(self, current_node): self.potential_nodes.append(neighbour_node) if self.config.fitcompress_parameters.optimize_pruning: - # Set neighbour state to current state neighbour_node_state = current_node_state.copy() @@ -973,7 +961,6 @@ def create_neighbours(self, current_node): current_node=current_node, pruning_metrics=neighbour_node_pruning_metrics ) else: - neighbour_node_parameters_layerwise, neighbour_node_unquantized_parameters_layerwise = self.add_pruning( current_node=current_node, params=current_node.parameters.copy(), @@ -1169,7 +1156,6 @@ def calculate_current_compression_rate(self, params_layerwise, quant_config): class FIT: - def __init__(self, model, device, input_spec): """ Initialize the FIT class, which is used to compute the FIT values for quantization and pruning. @@ -1217,12 +1203,11 @@ def get_model_weights(self, model): layer_names = [] # Iterate through all modules in the model for name, module in model.named_modules(): - if isinstance(module, (PQDense, PQConv2d)): layer_names.append(name) for name_param, matrix_param in list(module.named_parameters()): # Search for the weights - if name_param.endswith('_weight'): + if name_param.endswith("_weight"): matrices_params_layerwise.append(matrix_param) # Set their collect flag to True (later on we can then access them easily like this) matrix_param.collect = True @@ -1276,7 +1261,7 @@ def hook_removal(self): self.hooks.clear() assert len(self.hooks) == 0, "Hooks were not removed properly!" - def get_loss(self, model, data_batch, target_batch, loss_func, mode='mini-batch'): + def get_loss(self, model, data_batch, target_batch, loss_func, mode="mini-batch"): """ This function triggers the loss calcuation of a model. We use it such that we can then calculate gradients which are @@ -1302,16 +1287,16 @@ def get_loss(self, model, data_batch, target_batch, loss_func, mode='mini-batch' output = model(data_batch) - if mode == 'mini-batch': + if mode == "mini-batch": # Check which loss_func instance is active if isinstance(loss_func, torch.nn.CrossEntropyLoss): # Calculate loss based on mini-batch and averaged over it loss_func = torch.nn.CrossEntropyLoss() - if mode == 'sample': + if mode == "sample": if isinstance(loss_func, torch.nn.CrossEntropyLoss): # Calculate loss for each sample - loss_func = torch.nn.CrossEntropyLoss(reduce=False, reduction='none') + loss_func = torch.nn.CrossEntropyLoss(reduce=False, reduction="none") loss = loss_func(output, target_batch) @@ -1405,7 +1390,7 @@ def get_EF(self, model, data_loader, loss_func, tolerance=1e-3, min_iterations=1 if data_batch.size(0) != batch_size: continue # Uneven batches break loop - loss = self.get_loss(model, data_batch, target_batch, loss_func, mode='mini-batch') + loss = self.get_loss(model, data_batch, target_batch, loss_func, mode="mini-batch") curr_batch_matrices_params_layerwise = [] curr_batch_minmax_range_params_layerwise = [] for weights in model.parameters(): @@ -1639,7 +1624,6 @@ def get_FIT_real_values(self, params_before, EF_trace_params_layerwise, params_a curr_FIT += EF_trace * delta_theta else: - for theta, EF_trace in zip(params_before, EF_trace_params_layerwise): # Calculate the squared difference between the parameters before and after delta_theta = torch.sum(theta.detach().cpu() ** 2) @@ -1668,10 +1652,8 @@ def get_FIT_old(self, params_before=None, FeM=None, params_after=None, same_thet curr_FIT = 0 if not same_theta: - # Taken from compute_fake_FIT_params() for theta_before, theta_after, layer_FeM in zip(params_before, params_after, FeM): - curr_FIT_layer = torch.sum( layer_FeM * (theta_before.detach().cpu() - theta_after.detach().cpu()) ** 2 ).numpy() @@ -1684,7 +1666,6 @@ def get_FIT_old(self, params_before=None, FeM=None, params_after=None, same_thet # Taken from generate_FIT_pruning_importance() for theta_after, layer_FeM in zip(params_after, FeM): - curr_FIT_layer = layer_FeM * (theta_after.detach().cpu() ** 2) FIT_layerwise.append(curr_FIT_layer) diff --git a/src/pquant/core/torch/fixed_point_quantizer.py b/src/pquant/core/torch/fixed_point_quantizer.py index e3ecc4e..b4f9bc5 100644 --- a/src/pquant/core/torch/fixed_point_quantizer.py +++ b/src/pquant/core/torch/fixed_point_quantizer.py @@ -10,7 +10,7 @@ round_mode_registry: dict[str, Callable[[Any], Any]] = {} saturation_mode_registry: dict[str, Callable[[Any, Any, Any, Any], Any]] = {} -T = TypeVar('T', bound=ArrayLike) +T = TypeVar("T", bound=ArrayLike) def _clip(x, min_value, max_value): @@ -34,44 +34,44 @@ def wrapper(x): return inner -@rnd_mode('TRN') +@rnd_mode("TRN") def floor(x): return torch.floor(x) -@rnd_mode('RND') +@rnd_mode("RND") def round(x): # Round to nearest, ties positive infinity. return torch.floor(x + 0.5) -@rnd_mode('RND_CONV') +@rnd_mode("RND_CONV") def round_conv(x): # Round to nearest, ties to even. return torch.round(x) -@rnd_mode('TRN_ZERO') +@rnd_mode("TRN_ZERO") def floor_zero(x): # Truncate towards zero. sign = torch.sign(x) return torch.floor(torch.abs(x)) * sign # type: ignore -@rnd_mode('RND_ZERO') +@rnd_mode("RND_ZERO") def round_zero(x): # Round to nearest, ties towards zero. sign = torch.sign(x) return -torch.floor(-torch.abs(x) + 0.5) * sign # type:ignore -@rnd_mode('RND_MIN_INF') +@rnd_mode("RND_MIN_INF") def round_min_inf(x): # Round to nearest, ties towards negative infinity. return -torch.floor(-x + 0.5) # type:ignore -@rnd_mode('RND_INF') +@rnd_mode("RND_INF") def round_inf(x): # Round to nearest, ties away from zero. sign = torch.sign(x) @@ -91,7 +91,7 @@ def inner(func): return inner -@sat_mode('WRAP') +@sat_mode("WRAP") def wrap(x, k, i, f): xs = x bk = i + k @@ -99,7 +99,7 @@ def wrap(x, k, i, f): return (xs + bias) % (2.0**bk) - bias -@sat_mode('SAT') +@sat_mode("SAT") def sat(x, k, i, f): f_eps = 2.0 ** (-f) __max = 2.0**i @@ -109,7 +109,7 @@ def sat(x, k, i, f): return r -@sat_mode('SAT_SYM') +@sat_mode("SAT_SYM") def sat_sym(x, k, i, f): f_eps = 2.0 ** (-f) _max = 2.0**i - f_eps @@ -118,7 +118,7 @@ def sat_sym(x, k, i, f): return r -@sat_mode('WRAP_SM') +@sat_mode("WRAP_SM") def wrap_sm_fn(x, k, i, f, training=None, quant_fn: Callable = lambda x: x): # x=ops.round(x*2.**f) # High and low bounds are reflective. When overflows, can be less trash than WARP but still more trash than SAT. @@ -149,20 +149,20 @@ def round(self, x, f: Any = 1.0): def saturate(self, x, k, i, f): return self.sat_fn(x, k, i, f) - def __init__(self, round_mode: str = 'TRN', overflow_mode: str = 'WRAP'): + def __init__(self, round_mode: str = "TRN", overflow_mode: str = "WRAP"): round_mode = round_mode.upper() overflow_mode = overflow_mode.upper() self.stochastic = False - if round_mode.startswith('S_'): + if round_mode.startswith("S_"): round_mode = round_mode[2:] self.stochastic = True - if overflow_mode == 'WRAP_SM': + if overflow_mode == "WRAP_SM": assert round_mode in ( - 'RND', - 'RND_CONV', - ), 'WRAP_SM only supports RND and RND_CONV rounding modes in this implementation.' + "RND", + "RND_CONV", + ), "WRAP_SM only supports RND and RND_CONV rounding modes in this implementation." self.round_mode = round_mode self.overflow_mode = overflow_mode @@ -177,16 +177,16 @@ def forward(self, x, k, i, f, training=False): # will be clipped off anyway. Thus have saturation before rounding, except for # wrap mode, which doesn't round during training. - if self.overflow_mode != 'WRAP': + if self.overflow_mode != "WRAP": x = self.saturate(x, k, i, f) x = self.round(x, f) - if self.overflow_mode == 'WRAP' and not training: + if self.overflow_mode == "WRAP" and not training: x = self.saturate(x, k, i, f) return x def forward_wrap_sm(self, x, k, i, f, training=False): def quant_fn(x): - return self.round(x, f, training and self.stochastic) + return self.round(x, f) x = wrap_sm_fn(x, k, i, f, training, quant_fn) return x @@ -194,14 +194,14 @@ def quant_fn(x): def __call__(self, x, k, i, f, training=False, seed_gen=None): i = torch.maximum(i, -f).detach() + (i - i.detach()) # type: ignore if self.stochastic and training: - assert seed_gen is not None, 'Seed generator must be provided for stochastic rounding.' - if self.overflow_mode != 'WRAP_SM': + assert seed_gen is not None, "Seed generator must be provided for stochastic rounding." + if self.overflow_mode != "WRAP_SM": return self.forward(x, k, i, f, training) else: return self.forward_wrap_sm(x, k, i, f, training) -def get_fixed_quantizer(round_mode: str = 'TRN', overflow_mode: str = 'WRAP'): +def get_fixed_quantizer(round_mode: str = "TRN", overflow_mode: str = "WRAP"): """Get a stateless fixed-point quantizer given the round and overflow mode. The quantizer is differentiable w.r.t. to the input and f, also i if using saturation overflow mode. diff --git a/src/pquant/core/torch/hgq_quantizer.py b/src/pquant/core/torch/hgq_quantizer.py new file mode 100644 index 0000000..e62c9a0 --- /dev/null +++ b/src/pquant/core/torch/hgq_quantizer.py @@ -0,0 +1,422 @@ +""" +Pure-PyTorch implementation of HGQ FixedPointQuantizerKIF + Quantizer wrapper. + +Replaces the Keras hgq2 library dependency for the PyTorch backend. +Combines DefaultBitwidthMapper + FixedPointQuantizerKIF + Quantizer into one +nn.Module with no inheritance chain. + +Reference: "HGQ: High Granularity Quantization for Real-time Neural Networks on FPGAs" + (Sun et al., FPGA '26) +""" + +import logging +import math + +import torch +import torch.nn as nn + +from pquant.core.constants import QuantizationGranularity +from pquant.core.torch.fixed_point_quantizer import get_fixed_quantizer, round_conv + +logger = logging.getLogger(__name__) + + +def _minimal_i_given_xf(absmax: torch.Tensor, f: torch.Tensor) -> torch.Tensor: + """Minimum integer bits needed to represent absmax with f fractional bits.""" + eps = 2.0 ** (-f) + return torch.ceil(torch.log2(absmax + eps + 1e-10)) + + +class HGQQuantizer(nn.Module): + """ + HGQ fixed-point quantizer parameterized by (k, i, f) — keep_negative, integer + bits, fractional bits. + + Parameters + ---------- + k0 : float + Initial sign bit (0 = unsigned, 1 = signed). Non-trainable. + i0 : float + Initial integer bits. Trainable for SAT; tracked buffer for WRAP. + f0 : float + Initial fractional bits. Always trainable. + overflow_mode : str + One of 'SAT', 'SAT_SYM', 'WRAP', 'WRAP_SM'. + round_mode : str + One of 'RND', 'RND_CONV', 'TRN', etc. + is_data : bool + True → data/activation quantizer (batch axis 0 always homogeneous; + per_channel keys on channel axis 1). + False → weight/bias quantizer (per_channel keys on output-channel axis 0). + granularity : str + One of 'per_tensor' (one shared bit-width) or 'per_weight' (one per element). + Controls which axes are shared. per_channel is not supported. + gamma : float + L1 regularisation coefficient on bit-widths. + i_decay_speed : float + WRAP mode only. Rate at which tracked i can decrease per step. + float('inf') means i is reset each step to the minimum required by data. + i_min, i_max : float + Clamp bounds for the i parameter / buffer. + f_min, f_max : float + Clamp bounds for the f parameter. + scaler : float | None + Optional: inputs are divided by scaler before quantisation and multiplied + after (equivalent to scaling the fixed-point range). + qnoise_factor : float | None + Optional: mix factor for quantisation noise injection during training. + output = input + qnoise_factor * (quantised - input). + affine : tuple[float, float] | None + Optional (scale, shift) applied after quantisation: out = out*scale + shift. + """ + + def __init__( + self, + k0: float, + i0: float, + f0: float, + overflow_mode: str, + round_mode: str, + is_data: bool, + granularity: str = QuantizationGranularity.PER_WEIGHT, + gamma: float = 1e-8, + i_decay_speed: float = float("inf"), + i_min: float = -23.0, + i_max: float = 23.0, + f_min: float = -24.0, + f_max: float = 24.0, + scaler=None, + qnoise_factor: float | None = None, + affine=None, + ): + super().__init__() + assert int(k0) in (0, 1), f"k0 must be 0 or 1, got {k0}" + + self.k0 = float(k0) + self.i0 = float(i0) + self.f0 = float(f0) + self.overflow_mode = overflow_mode.upper() + self.round_mode = round_mode.upper() + self.is_data = is_data + self.granularity = QuantizationGranularity(granularity).value + self.gamma = gamma + self.i_decay_speed = i_decay_speed + self.i_min = i_min + self.i_max = i_max + self.f_min = f_min + self.f_max = f_max + self.scaler = scaler + self.qnoise_factor = qnoise_factor + self.affine = affine + + # Set during build() + self.homogeneous_axis: tuple[int, ...] = () + self._built = False + + self._stateless_quantizer = get_fixed_quantizer(round_mode=round_mode, overflow_mode=overflow_mode) + + # Scalar placeholders — replaced with shaped tensors in build(). + # k and i_raw are non-trainable buffers; f (and i for SAT) are Parameters. + self.register_buffer("_k", torch.tensor(self.k0)) + self._f = nn.Parameter(torch.tensor(self.f0)) + + if self.overflow_mode == "WRAP": + # Integer bits tracked as non-trainable running buffer (not optimised). + self.register_buffer("_i_raw", torch.tensor(self.i0)) + else: + # Integer bits are trainable for SAT / SAT_SYM. + self._i = nn.Parameter(torch.tensor(self.i0)) + + def build(self, input_shape: tuple) -> None: + """ + Initialise shaped parameter / buffer tensors. + + Called automatically on the first forward pass. The optimizer must be + created *after* build() has been called so that it tracks the shaped + parameters, not the scalar placeholders from __init__. + """ + device = self._k.device + self.homogeneous_axis = self._homogeneous_axis(len(input_shape)) + bw_shape = self._infer_bw_shape(input_shape) + + # k: non-trainable sign-bit buffer + self.register_buffer("_k", torch.full(bw_shape, self.k0, device=device)) + + # f: always trainable + self._f = nn.Parameter(torch.full(bw_shape, self.f0, device=device)) + + if self.overflow_mode == "WRAP": + self.register_buffer("_i_raw", torch.full(bw_shape, self.i0, device=device)) + else: + self._i = nn.Parameter(torch.full(bw_shape, self.i0, device=device)) + + self._built = True + + def _homogeneous_axis(self, ndim: int) -> tuple[int, ...]: + """Axes shared (collapsed to 1 in the bit-width tensor), derived from `granularity`. + + HGQ supports only per_tensor and per_weight (data always shares the batch axis 0). + per_channel is intentionally unsupported (the channel axis is layout-dependent). + """ + if self.granularity == QuantizationGranularity.PER_TENSOR: + return tuple(range(ndim)) + if self.granularity == QuantizationGranularity.PER_WEIGHT: + return (0,) if self.is_data else () + if self.granularity == QuantizationGranularity.PER_CHANNEL: + raise ValueError("per_channel granularity is not supported for HGQ. Use 'per_tensor' or 'per_weight'.") + raise ValueError(f"Unsupported granularity: {self.granularity}") + + def _infer_bw_shape(self, input_shape: tuple) -> tuple: + """Shape of bit-width parameter tensors: homogeneous axes collapse to 1.""" + return tuple(1 if ax in self.homogeneous_axis else d for ax, d in enumerate(input_shape)) + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def built(self) -> bool: + return self._built + + @property + def k(self) -> torch.Tensor: + return self._k.float() + + @property + def i(self) -> torch.Tensor: + """Integer bits, rounded with STE.""" + if self.overflow_mode == "WRAP": + # _i_raw is a buffer (no gradient); round_conv is still used for + # consistency — it degenerates to plain rounding when grad=0. + return round_conv(self._i_raw.float()) + return round_conv(self._i.float()) + + @property + def f(self) -> torch.Tensor: + """Fractional bits, rounded with STE.""" + return round_conv(self._f.float()) + + @property + def b(self) -> torch.Tensor: + """Total non-sign bits = relu(i + f). Zero when quantizer is pruned.""" + return torch.relu(self.i + self.f) + + # ------------------------------------------------------------------ + # Bit-width mapping helpers + # ------------------------------------------------------------------ + + def _bw_to_x(self, bw: torch.Tensor, x_shape: tuple) -> torch.Tensor: + return bw.expand(x_shape) + + def _x_to_bw_absmax(self, x: torch.Tensor) -> torch.Tensor: + if len(self.homogeneous_axis) == 0: + return x.abs() + return torch.amax(x.abs(), dim=self.homogeneous_axis, keepdim=True) + + # ------------------------------------------------------------------ + # Forward pass + # ------------------------------------------------------------------ + + def forward(self, x: torch.Tensor, training: bool = False) -> torch.Tensor: + if not self._built: + self.build(tuple(x.shape)) + + if training: + with torch.no_grad(): + self._f.data.clamp_(self.f_min, self.f_max) + if self.overflow_mode != "WRAP": + self._i.data.clamp_(self.i_min, self.i_max) + + x_in = x # kept for qnoise + + if self.scaler is not None: + x = x / self.scaler + + # Round bit-width parameters to integers with STE so gradients flow. + f_bw = round_conv(self._f) # bw-shaped + k = self._k.float() + f_x = f_bw.expand(x.shape) + k_x = k.expand(x.shape) + + if self.overflow_mode == "WRAP": + out = self._stateless_quantizer.round(x, f_x) + + if training: + # Track minimum integer bits needed for the current data. + # Done without gradient so it doesn't affect the f gradient path. + with torch.no_grad(): + absmax = self._x_to_bw_absmax(out.detach()) + min_i = _minimal_i_given_xf(absmax, f_bw.detach()) + if math.isinf(self.i_decay_speed): + new_i = min_i + else: + new_i = torch.maximum(self._i_raw - self.i_decay_speed, min_i) + self._i_raw.copy_(new_i.clamp(self.i_min, self.i_max)) + else: + if self.is_data: + # Data quantizer: apply wrap-modulo after rounding (inference). + i_x = self.i.expand(x.shape) + out = self._stateless_quantizer.saturate(out, k_x, i_x, f_x) + # Weight quantizer: rounded output is final, no saturation. + + # Zero out pruned (total bits == 0) values. + i_x = self.i.expand(x.shape) + out = torch.where(k_x + i_x + f_x > 0, out, torch.zeros_like(out)) + + else: # SAT / SAT_SYM + i_bw = round_conv(self._i) # bw-shaped, STE for gradient + i_x = i_bw.expand(x.shape) + out = self._stateless_quantizer(x, k_x, i_x, f_x, training) + + if not training: + out = torch.where(k_x + i_x + f_x > 0, out, torch.zeros_like(out)) + + if self.scaler is not None: + out = out * self.scaler + + if self.qnoise_factor is not None and training: + out = x_in + self.qnoise_factor * (out - x_in) + + if self.affine is not None: + out = out * self.affine[0] + self.affine[1] + + return out + + # ------------------------------------------------------------------ + # Regularisation / constraint / utility methods + # ------------------------------------------------------------------ + + def regularization_loss(self) -> torch.Tensor: + """ + L1 regularisation on bit-widths. + + Replaces Keras `layer.losses` list. Contributes the γ·Σ(bit-widths) term + from eq. (12) in the paper. For SAT mode both i and f are regularised; + for WRAP mode only f (i is not a learnable parameter). + """ + if not self._built or self.gamma == 0.0: + return torch.tensor(0.0, device=self._k.device) + loss = self.gamma * self.f.sum() + if self.overflow_mode != "WRAP": + loss = loss + self.gamma * self.i.sum() + return loss + + def post_epoch_constraint_apply(self) -> None: + """ + Clamp i and f parameters to [*_min, *_max] in-place. + + Replaces Keras `variable.constraint(variable)` + `variable.assign()`. + Call at the end of each epoch. + """ + with torch.no_grad(): + self._f.data.clamp_(self.f_min, self.f_max) + if self.overflow_mode != "WRAP": + self._i.data.clamp_(self.i_min, self.i_max) + # WRAP: _i_raw is already clamped inside forward(). + + def set_bits(self, i, f) -> None: + """ + Overwrite i and f with scalar or tensor values. + + Replaces Keras `variable.assign(variable * 0 + value)` pattern. + Used by `apply_final_compression` and `reload_from_local` in quantizer.py. + """ + with torch.no_grad(): + i_t = torch.as_tensor(i, dtype=torch.float32) + f_t = torch.as_tensor(f, dtype=torch.float32) + self._f.data.copy_(f_t) if f_t.shape == self._f.shape else self._f.data.fill_(f_t.item()) + if self.overflow_mode == "WRAP": + self._i_raw.copy_(i_t) if i_t.shape == self._i_raw.shape else self._i_raw.fill_(i_t.item()) + else: + self._i.data.copy_(i_t) if i_t.shape == self._i.shape else self._i.data.fill_(i_t.item()) + + def bits_(self, shape: tuple) -> torch.Tensor: + """ + Return total bits (k + relu(i+f)) broadcast to *shape*. + + Used by `Quantizer.get_total_bits()` for EBOPs calculations. + """ + if not self._built: + total = self.k0 + max(self.i0 + self.f0, 0.0) + return torch.full(shape, total, device=self._k.device) + return (self.k + self.b).expand(shape) + + +# --------------------------------------------------------------------------- +# Quick sanity check +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + torch.manual_seed(0) + logging.basicConfig(level=logging.DEBUG, format="%(message)s") + logger.debug("=== HGQQuantizer sanity checks ===\n") + + # ---- SAT, signed, weight quantizer (per-element) ---- + x = torch.randn(4, 3) * 0.5 + q_sat = HGQQuantizer(k0=1, i0=2, f0=4, overflow_mode="SAT", round_mode="RND", is_data=False) + out = q_sat(x, training=False) + logger.debug(f"SAT weight forward: input range [{x.min():.3f}, {x.max():.3f}]") + logger.debug(f" output range [{out.min():.3f}, {out.max():.3f}]") + max_val = 2**2 - 2**-4 + assert out.max() <= max_val + 1e-5, "SAT upper bound violated" + assert out.min() >= -(2**2) - 1e-5, "SAT lower bound violated" + logger.debug(" bounds check OK\n") + + # ---- SAT backward (gradient through f) ---- + q_sat2 = HGQQuantizer(k0=1, i0=2, f0=4, overflow_mode="SAT", round_mode="RND", is_data=False) + x2 = torch.randn(4, 3, requires_grad=True) + out2 = q_sat2(x2, training=True) + out2.sum().backward() + assert q_sat2._f.grad is not None and q_sat2._f.grad.abs().sum() > 0 + assert x2.grad is not None and x2.grad.abs().sum() > 0 + logger.debug("SAT backward: gradients w.r.t. f and input OK") + logger.debug(f" f grad mean: {q_sat2._f.grad.mean():.6f}\n") + + # ---- WRAP, signed, data quantizer (per-batch) ---- + x3 = torch.randn(8, 4) * 3.0 + q_wrap = HGQQuantizer(k0=1, i0=0, f0=4, overflow_mode="WRAP", round_mode="RND", is_data=True) + out3 = q_wrap(x3, training=True) + logger.debug(f"WRAP data training: _i_raw after forward: {q_wrap._i_raw.flatten()[:4]}") + assert q_wrap._i_raw.max() >= 0, "i_raw should track positive range" + logger.debug(" _i_raw tracking OK\n") + + # ---- WRAP backward (gradient through f, not i_raw) ---- + q_wrap2 = HGQQuantizer(k0=1, i0=2, f0=4, overflow_mode="WRAP", round_mode="RND", is_data=False) + x4 = torch.randn(4, 3, requires_grad=True) + out4 = q_wrap2(x4, training=True) + out4.sum().backward() + assert q_wrap2._f.grad is not None and q_wrap2._f.grad.abs().sum() > 0 + assert q_wrap2._i_raw.grad is None, "_i_raw must not accumulate gradient" + logger.debug("WRAP backward: f has gradient, _i_raw has none OK\n") + + # ---- Zero bits → pruned: k=0,i=0,f=0 → k+i+f=0 → all zero ---- + q_pruned = HGQQuantizer(k0=0, i0=0, f0=0, overflow_mode="SAT", round_mode="RND", is_data=False) + x5 = torch.randn(4, 3) + out5 = q_pruned(x5, training=False) + assert torch.all(out5 == 0), "Zero-bit quantizer should output zeros" + logger.debug("Pruning (zero bits): output is all zeros OK\n") + + # ---- regularization_loss ---- + q_reg = HGQQuantizer(k0=1, i0=2.0, f0=4.0, overflow_mode="SAT", round_mode="RND", is_data=False, gamma=1e-3) + q_reg(torch.randn(4, 3)) + loss = q_reg.regularization_loss() + expected = 1e-3 * (q_reg.f.sum() + q_reg.i.sum()) + assert torch.isclose(loss, expected, rtol=1e-4), f"reg loss mismatch: {loss} vs {expected}" + logger.debug(f"regularization_loss OK: {loss.item():.6f}\n") + + # ---- bits_ shape ---- + q_bits = HGQQuantizer(k0=1, i0=2.0, f0=4.0, overflow_mode="SAT", round_mode="RND", is_data=False) + q_bits(torch.randn(4, 3)) + b = q_bits.bits_((8, 4, 3)) + assert b.shape == (8, 4, 3), f"bits_ shape wrong: {b.shape}" + logger.debug(f"bits_ shape OK: {b.shape}\n") + + # ---- set_bits ---- + q_set = HGQQuantizer(k0=1, i0=2.0, f0=4.0, overflow_mode="SAT", round_mode="RND", is_data=False) + q_set(torch.randn(4, 3)) + q_set.set_bits(3.0, 5.0) + assert math.isclose(q_set.i.mean().item(), 3.0, abs_tol=1e-5) + assert math.isclose(q_set.f.mean().item(), 5.0, abs_tol=1e-5) + logger.debug("set_bits OK\n") + + logger.debug("=== All checks passed ===") diff --git a/src/pquant/core/torch/layers.py b/src/pquant/core/torch/layers.py index 87ced7c..7badc30 100644 --- a/src/pquant/core/torch/layers.py +++ b/src/pquant/core/torch/layers.py @@ -1,24 +1,46 @@ +import math import typing -from typing import Optional, Tuple, TypeVar, Union +from typing import TypeVar import torch import torch.nn as nn import torch.nn.functional as F -from torch.fx import symbolic_trace from torch.nn.common_types import _size_1_t, _size_2_t -from pquant.core.torch.activations import PQActivation +from pquant.core.torch.activations import PQActivation, PQSoftmax from pquant.core.torch.quantizer import Quantizer -from pquant.core.utils import get_pruning_layer +from pquant.core.torch.utils import get_pruning_layer if typing.TYPE_CHECKING: - from pquant.core.torch.fit_compress import call_fitcompress # noqa: 401 - -from keras import ops + pass # noqa: 401 T = TypeVar("T") +def _resolve_data_quant_bits(quant_bits, config): + """Return (k, i, f) from an explicit tuple, or the config's data-lane defaults.""" + if quant_bits is not None: + return quant_bits + parameters = config.quantization_parameters + return ( + parameters.default_data_keep_negatives, + parameters.default_data_integer_bits, + parameters.default_data_fractional_bits, + ) + + +def _resolve_weight_quant_bits(quant_bits, config): + """Return (k, i, f) from an explicit tuple, or the config's weight defaults.""" + if quant_bits is not None: + return quant_bits + parameters = config.quantization_parameters + return ( + parameters.default_weight_keep_negatives, + parameters.default_weight_integer_bits, + parameters.default_weight_fractional_bits, + ) + + class PQWeightBiasBase(nn.Module): def __init__( self, @@ -27,41 +49,23 @@ def __init__( quantize_input=True, quantize_output=False, enable_pruning: bool = None, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, *args, **kwargs, ): super().__init__(**kwargs) - if in_quant_bits is not None: - self.k_input, self.i_input, self.f_input = in_quant_bits - else: - self.k_input = config.quantization_parameters.default_data_keep_negatives - self.i_input = config.quantization_parameters.default_data_integer_bits - self.f_input = config.quantization_parameters.default_data_fractional_bits - - if weight_quant_bits is not None: - self.k_weight, self.i_weight, self.f_weight = weight_quant_bits - else: - self.k_weight = config.quantization_parameters.default_weight_keep_negatives - self.i_weight = config.quantization_parameters.default_weight_integer_bits - self.f_weight = config.quantization_parameters.default_weight_fractional_bits - if bias_quant_bits is not None: - self.k_bias, self.i_bias, self.f_bias = bias_quant_bits - else: - self.k_bias = config.quantization_parameters.default_weight_keep_negatives - self.i_bias = config.quantization_parameters.default_weight_integer_bits - self.f_bias = config.quantization_parameters.default_weight_fractional_bits - - if out_quant_bits is not None: - self.k_output, self.i_output, self.f_output = out_quant_bits - else: - self.k_output = config.quantization_parameters.default_data_keep_negatives - self.i_output = config.quantization_parameters.default_data_integer_bits - self.f_output = config.quantization_parameters.default_data_fractional_bits + self.k_input, self.i_input, self.f_input = _resolve_data_quant_bits(in_quant_bits, config) + self.k_weight, self.i_weight, self.f_weight = _resolve_weight_quant_bits(weight_quant_bits, config) + self.k_bias, self.i_bias, self.f_bias = _resolve_weight_quant_bits(bias_quant_bits, config) + self.k_output, self.i_output, self.f_output = _resolve_data_quant_bits(out_quant_bits, config) self.pruning_layer = get_pruning_layer(config=config, layer_type=layer_type) self.pruning_method = config.pruning_parameters.pruning_method @@ -78,7 +82,13 @@ def __init__( self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress self.hgq_gamma = config.quantization_parameters.hgq_gamma self.granularity = config.quantization_parameters.granularity - self.final_compression_done = False + self.weight_quant_granularity = ( + weight_quant_granularity if weight_quant_granularity is not None else self.granularity + ) + self.in_quant_granularity = in_quant_granularity if in_quant_granularity is not None else self.granularity + self.bias_quant_granularity = bias_quant_granularity if bias_quant_granularity is not None else self.granularity + self.out_quant_granularity = out_quant_granularity if out_quant_granularity is not None else self.granularity + self.register_buffer("final_compression_done", torch.tensor(False)) self.built = False self.parallelization_factor = -1 self.hgq_beta = config.quantization_parameters.hgq_beta @@ -87,61 +97,69 @@ def __init__( self.post_fitcompress_calibration = False self.saved_inputs = [] self.saved_outputs = [] + self.config = config - def check_is_built(self, input_shape): + def _check_is_built(self, input_shape): if self.built: return - # Build function to delay quantizer creation until after custom i,f bits have been set + # Quantizer creation is delayed until the first forward so custom i/f bits set + # after __init__ are picked up. if self.quantize_input: self.input_quantizer = Quantizer( - k=torch.tensor(self.k_input), - i=torch.tensor(self.i_input), - f=torch.tensor(self.f_input), + k=self.k_input, + i=self.i_input, + f=self.f_input, overflow=self.overflow_mode_data, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=self.in_quant_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.config.quantization_parameters.dynamic_data_quantization, ) self.weight_quantizer = Quantizer( - k=torch.tensor(self.k_weight), - i=torch.tensor(self.i_weight), - f=torch.tensor(self.f_weight), + k=self.k_weight, + i=self.i_weight, + f=self.f_weight, overflow=self.overflow_mode_parameters, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=False, + granularity=self.weight_quant_granularity, hgq_gamma=self.hgq_gamma, - granularity=self.granularity, place="weight", + shape=self._weight.shape, ) - self.bias_quantizer = Quantizer( - k=torch.tensor(self.k_bias), - i=torch.tensor(self.i_bias), - f=torch.tensor(self.f_bias), + k=self.k_bias, + i=self.i_bias, + f=self.f_bias, overflow=self.overflow_mode_parameters, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=False, + granularity=self.bias_quant_granularity, hgq_gamma=self.hgq_gamma, place="bias", + shape=None if self._bias is None else self._bias.shape, ) if self.quantize_output: self.output_quantizer = Quantizer( - k=torch.tensor(self.k_output), - i=torch.tensor(self.i_output), - f=torch.tensor(self.f_output), + k=self.k_output, + i=self.i_output, + f=self.f_output, overflow=self.overflow_mode_data, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=self.out_quant_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.config.quantization_parameters.dynamic_data_quantization, ) - self.n_parallel = ops.prod(tuple(input_shape)[1:-1]) + self.n_parallel = math.prod(tuple(input_shape)[1:-1]) self.parallelization_factor = self.parallelization_factor if self.parallelization_factor > 0 else self.n_parallel self.built = True self.input_shape = (1,) + input_shape[1:] @@ -161,20 +179,42 @@ def get_output_quantization_bits(self): def apply_final_compression(self): pass + def _register_compressed_parameters(self, bias): + """Store the wrapped layer's weight/bias as `_weight`/`_bias`; the `weight`/`bias` + properties then return their pruned and quantized views.""" + self._weight = nn.Parameter(self.weight.clone()).to(self.weight.device) + self.register_parameter("_weight", self._weight) + if bias: + self._bias = nn.Parameter(self.bias.clone()).to(self.bias.device) + self.register_parameter("_bias", self._bias) + else: + self.register_parameter("_bias", None) + self.pruning_layer.build(self._weight.shape) + def post_pre_train_function(self): self.is_pretraining = False if self.pruning_layer is not None: self.pruning_layer.post_pre_train_function() - def save_weights(self): + def _save_weights(self): self.init_weight = self._weight.clone() - def rewind_weights(self): + def _rewind_weights(self): + if not hasattr(self, "init_weight"): + return self._weight.data = self.init_weight.clone() def ebops(self): return 0.0 + def _masked_weight_bits(self, bw_ker): + """Zero the bit counts of weights that are pruned away or below the quantization step size.""" + bw_ker = bw_ker * self.pruning_layer.get_hard_mask() + _, _, f = self.get_weight_quantization_bits() + quantization_step_size = 2 ** (-f - 1) + step_size_mask = (torch.abs(self._weight) > quantization_step_size).float() + return bw_ker * step_size_mask + def hgq_loss(self): if self.is_pretraining or not self.use_hgq: return 0.0 @@ -189,20 +229,20 @@ def hgq_loss(self): return loss def quantize(self, x, quantizer): - if self.enable_quantization and not self.is_fitcompress_pretraining(): + if self.enable_quantization and not self._is_fitcompress_pretraining(): return quantizer(x) if x is not None else x return x - def prune(self, weight): + def _prune(self, weight): if self.enable_pruning: weight = self.pruning_layer(weight) return weight - def is_fitcompress_pretraining(self): + def _is_fitcompress_pretraining(self): return self.is_pretraining and self.use_fitcompress def pre_forward(self, x): - self.check_is_built(x.shape) + self._check_is_built(x.shape) if self.post_fitcompress_calibration: self.saved_inputs.append(x) return x @@ -212,7 +252,7 @@ def pre_forward(self, x): self.pruning_layer.collect_input(x, self.weight, self.training) return x - def post_forward(self, x): + def _post_forward(self, x): if self.post_fitcompress_calibration: self.saved_outputs.append(x) return x @@ -235,10 +275,14 @@ def __init__( enable_pruning: bool = None, device=None, dtype=None, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, **kwargs, ): super().__init__( @@ -256,52 +300,44 @@ def __init__( weight_quant_bits=weight_quant_bits, bias_quant_bits=bias_quant_bits, out_quant_bits=out_quant_bits, + weight_quant_granularity=weight_quant_granularity, + in_quant_granularity=in_quant_granularity, + bias_quant_granularity=bias_quant_granularity, + out_quant_granularity=out_quant_granularity, **kwargs, ) self.in_features = in_features self.out_features = out_features self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress - self._weight = nn.Parameter(self.weight.clone()).to(self.weight.device) - self.register_parameter("_weight", self._weight) - if bias: - self._bias = nn.Parameter(self.bias.clone()).to(self.bias.device) - self.register_parameter("_bias", self._bias) - else: - self.register_parameter("_bias", None) - self.pruning_layer.build(self._weight.shape) - self.final_compression_done = nn.Parameter(torch.tensor(False), requires_grad=False) + self._register_compressed_parameters(bias) def ebops(self, include_mask=False): bw_inp = self.input_quantizer.get_total_bits(self.input_shape) - bw_ker = self.weight_quantizer.get_total_bits(ops.shape(self._weight)) + bw_ker = self.weight_quantizer.get_total_bits(self._weight.shape) if include_mask: - bw_ker = bw_ker * self.pruning_layer.get_hard_mask() - _, _, f = self.get_weight_quantization_bits() - quantization_step_size = 2 ** (-f - 1) - step_size_mask = (torch.abs(self._weight) >= quantization_step_size).float() - bw_ker = bw_ker * step_size_mask - ebops = ops.sum(F.linear(bw_inp, bw_ker)) + bw_ker = self._masked_weight_bits(bw_ker) + ebops = torch.sum(F.linear(bw_inp, bw_ker)) if self._bias is not None: - bw_bias = self.bias_quantizer.get_total_bits(ops.shape(self._bias)) - size = ops.cast(ops.prod(self.input_shape[:-1]) * self.out_features, self._weight.dtype) - ebops += ops.mean(bw_bias) * size + bw_bias = self.bias_quantizer.get_total_bits(self._bias.shape) + size = float(math.prod(self.input_shape[:-1]) * self.out_features) + ebops += torch.mean(bw_bias) * size ebops = ebops * self.parallelization_factor / self.n_parallel return ebops @property def weight(self): - if self.final_compression_done or self.is_fitcompress_pretraining(): + if self.final_compression_done or self._is_fitcompress_pretraining(): return self._weight if self.pruning_first: - weight = self.prune(self._weight) + weight = self._prune(self._weight) return self.quantize(weight, self.weight_quantizer) else: weight = self.quantize(self._weight, self.weight_quantizer) - return self.prune(weight) + return self._prune(weight) @property def bias(self): - if self.final_compression_done or self.is_fitcompress_pretraining(): + if self.final_compression_done or self._is_fitcompress_pretraining(): return self._bias bias = self.quantize(self._bias, self.bias_quantizer) return bias @@ -310,12 +346,12 @@ def apply_final_compression(self): self._weight.data = self.weight if self._bias is not None: self._bias.data = self.bias - self.final_compression_done.data = torch.tensor(True) + self.final_compression_done.fill_(True) def forward(self, x): x = self.pre_forward(x) x = super().forward(x) - x = self.post_forward(x) + x = self._post_forward(x) return x def extra_repr(self) -> str: @@ -331,85 +367,29 @@ def extra_repr(self) -> str: ) -class PQConv2d(PQWeightBiasBase, nn.Conv2d): - def __init__( - self, - config, - in_channels: int, - out_channels: int, - kernel_size: _size_2_t, - stride: _size_2_t = 1, - padding: Union[str, _size_2_t] = 0, - dilation: _size_2_t = 1, - groups: int = 1, - bias: bool = True, - padding_mode: str = "zeros", # TODO: refine this type - device=None, - dtype=None, - quantize_input=True, - quantize_output=False, - enable_pruning: bool = None, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, - **kwargs, - ): - super().__init__( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - groups=groups, - bias=bias, - padding_mode=padding_mode, - device=device, - dtype=dtype, - config=config, - layer_type="conv", - quantize_input=quantize_input, - quantize_output=quantize_output, - enable_pruning=enable_pruning, - in_quant_bits=in_quant_bits, - weight_quant_bits=weight_quant_bits, - bias_quant_bits=bias_quant_bits, - out_quant_bits=out_quant_bits, - **kwargs, - ) - self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress - self._weight = nn.Parameter(self.weight.clone()).to(self.weight.device) - self.register_parameter("_weight", self._weight) - if bias: - self._bias = nn.Parameter(self.bias.clone()).to(self.bias.device) - self.register_parameter("_bias", self._bias) - else: - self.register_parameter("_bias", None) - self.pruning_layer.build(self._weight.shape) +class PQConvBase(PQWeightBiasBase): + """Pruning/quantization behavior shared by PQConv1d and PQConv2d.""" + + conv_bits_fn = None # F.conv1d / F.conv2d, set by the subclasses def ebops(self, include_mask=False): bw_inp = self.input_quantizer.get_total_bits(self.input_shape) - bw_ker = self.weight_quantizer.get_total_bits(ops.shape(self._weight)) + bw_ker = self.weight_quantizer.get_total_bits(self._weight.shape) if include_mask: - bw_ker = bw_ker * self.pruning_layer.get_hard_mask() - _, _, f = self.get_weight_quantization_bits() - quantization_step_size = 2 ** (-f - 1) - step_size_mask = (torch.abs(self._weight) > quantization_step_size).float() - bw_ker = bw_ker * step_size_mask + bw_ker = self._masked_weight_bits(bw_ker) if self.parallelization_factor < 0: - ebops = ops.sum(F.conv2d(bw_inp, bw_ker, stride=self.stride, padding=self.padding, dilation=self.dilation)) + ebops = torch.sum( + self.conv_bits_fn(bw_inp, bw_ker, stride=self.stride, padding=self.padding, dilation=self.dilation) + ) else: - reduce_axis_kernel = tuple(range(2, 4)) - reduce_axis_input = (0,) + tuple(range(2, 4)) - - bw_inp = ops.max(bw_inp, axis=reduce_axis_input) - bw_ker = ops.sum(bw_ker, axis=reduce_axis_kernel) - ebops = ops.sum(bw_inp[None, :] * bw_ker) + spatial_axes = tuple(range(2, 2 + len(self.kernel_size))) + bw_inp = torch.amax(bw_inp, dim=(0,) + spatial_axes) + bw_ker = torch.sum(bw_ker, dim=spatial_axes) + ebops = torch.sum(bw_inp[None, :] * bw_ker) if self._bias is not None: - size = ops.cast(ops.prod(list(self.input_shape)), self.weight.dtype) - bw_bias = self.bias_quantizer.get_total_bits(ops.shape(self._bias)) - ebops += ops.mean(bw_bias) * size + size = float(math.prod(self.input_shape)) + bw_bias = self.bias_quantizer.get_total_bits(self._bias.shape) + ebops += torch.mean(bw_bias) * size return ebops @property @@ -417,39 +397,27 @@ def weight(self): if self.final_compression_done: return self._weight if self.pruning_first: - weight = self.prune(self._weight) + weight = self._prune(self._weight) return self.quantize(weight, self.weight_quantizer) - else: - weight = self.quantize(self._weight, self.weight_quantizer) - return self.prune(weight) + weight = self.quantize(self._weight, self.weight_quantizer) + return self._prune(weight) @property def bias(self): if self.final_compression_done: return self._bias - bias = self.quantize(self._bias, self.bias_quantizer) - return bias + return self.quantize(self._bias, self.bias_quantizer) def apply_final_compression(self): self._weight.data = self.weight if self._bias is not None: self._bias.data = self.bias - self.final_compression_done = True + self.final_compression_done.fill_(True) def forward(self, x): x = self.pre_forward(x) - weight = self.weight - bias = self.bias - x = F.conv2d( - x, - weight, - bias, - self.stride, - self.padding, - self.dilation, - self.groups, - ) - x = self.post_forward(x) + x = super().forward(x) + x = self._post_forward(x) return x def extra_repr(self): @@ -466,22 +434,21 @@ def extra_repr(self): s += ", bias=False" if self.padding_mode != "zeros": s += ", padding_mode={padding_mode}" - s += ", self.quantize_input={quantize_input} " - s += ", self.quantize_output={quantize_output}" - + s += ", quantize_input={quantize_input}" + s += ", quantize_output={quantize_output}" return s.format(**self.__dict__) -class PQConv1d(PQWeightBiasBase, nn.Conv1d): +class PQConv2d(PQConvBase, nn.Conv2d): def __init__( self, config, in_channels: int, out_channels: int, - kernel_size: _size_1_t, - stride: _size_1_t = 1, - padding: Union[str, _size_1_t] = 0, - dilation: _size_1_t = 1, + kernel_size: _size_2_t, + stride: _size_2_t = 1, + padding: str | _size_2_t = 0, + dilation: _size_2_t = 1, groups: int = 1, bias: bool = True, padding_mode: str = "zeros", # TODO: refine this type @@ -490,10 +457,14 @@ def __init__( quantize_input=True, quantize_output=False, enable_pruning: bool = None, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, **kwargs, ): super().__init__( @@ -517,98 +488,92 @@ def __init__( weight_quant_bits=weight_quant_bits, bias_quant_bits=bias_quant_bits, out_quant_bits=out_quant_bits, + weight_quant_granularity=weight_quant_granularity, + in_quant_granularity=in_quant_granularity, + bias_quant_granularity=bias_quant_granularity, + out_quant_granularity=out_quant_granularity, **kwargs, ) self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress - self._weight = nn.Parameter(self.weight.clone()).to(self.weight.device) - self.register_parameter("_weight", self._weight) - if bias: - self._bias = nn.Parameter(self.bias.clone()).to(self.bias.device) - self.register_parameter("_bias", self._bias) - else: - self.register_parameter("_bias", None) - self.pruning_layer.build(self._weight.shape) + self._register_compressed_parameters(bias) - def ebops(self, include_mask=False): - bw_inp = self.input_quantizer.get_total_bits(self.input_shape) - bw_ker = self.weight_quantizer.get_total_bits(ops.shape(self._weight)) - if include_mask: - bw_ker = bw_ker * self.pruning_layer.get_hard_mask() - _, _, f = self.get_weight_quantization_bits() - quantization_step_size = 2 ** (-f - 1) - step_size_mask = (torch.abs(self._weight) > quantization_step_size).float() - bw_ker = bw_ker * step_size_mask - if self.parallelization_factor < 0: - ebops = ops.sum(F.conv1d(bw_inp, bw_ker, stride=self.stride, padding=self.padding, dilation=self.dilation)) - else: - reduce_axis_kernel = tuple(range(2, 3)) - reduce_axis_input = (0,) + tuple(range(2, 3)) - - bw_inp = ops.max(bw_inp, axis=reduce_axis_input) - bw_ker = ops.sum(bw_ker, axis=reduce_axis_kernel) - ebops = ops.sum(bw_inp[None, :] * bw_ker) - if self.bias is not None: - size = ops.cast(ops.prod(list(self.input_shape)), self.weight.dtype) - bw_bias = self.bias_quantizer.get_total_bits(ops.shape(self._bias)) - ebops += ops.mean(bw_bias) * size - return ebops + conv_bits_fn = staticmethod(F.conv2d) - @property - def weight(self): - if self.final_compression_done: - return self._weight - if self.pruning_first: - weight = self.prune(self._weight) - return self.quantize(weight, self.weight_quantizer) - else: - weight = self.quantize(self._weight, self.weight_quantizer) - return self.prune(weight) - @property - def bias(self): - if self.final_compression_done: - return self._bias - bias = self.quantize(self._bias, self.bias_quantizer) - return bias - - def apply_final_compression(self): - self._weight.data = self.weight - if self._bias is not None: - self._bias.data = self.bias - self.final_compression_done = True +class PQConv1d(PQConvBase, nn.Conv1d): + def __init__( + self, + config, + in_channels: int, + out_channels: int, + kernel_size: _size_1_t, + stride: _size_1_t = 1, + padding: str | _size_1_t = 0, + dilation: _size_1_t = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", # TODO: refine this type + device=None, + dtype=None, + quantize_input=True, + quantize_output=False, + enable_pruning: bool = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_granularity=None, + in_quant_granularity=None, + bias_quant_granularity=None, + out_quant_granularity=None, + **kwargs, + ): + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + config=config, + layer_type="conv", + quantize_input=quantize_input, + quantize_output=quantize_output, + enable_pruning=enable_pruning, + in_quant_bits=in_quant_bits, + weight_quant_bits=weight_quant_bits, + bias_quant_bits=bias_quant_bits, + out_quant_bits=out_quant_bits, + weight_quant_granularity=weight_quant_granularity, + in_quant_granularity=in_quant_granularity, + bias_quant_granularity=bias_quant_granularity, + out_quant_granularity=out_quant_granularity, + **kwargs, + ) + self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress + self._register_compressed_parameters(bias) - def forward(self, x): - x = self.pre_forward(x) - x = super().forward(x) - x = self.post_forward(x) - return x + conv_bits_fn = staticmethod(F.conv1d) - def extra_repr(self): - s = "{in_channels}, {out_channels}, kernel_size={kernel_size}, stride={stride}" - if self.padding != (0,) * len(self.padding): - s += ", padding={padding}" - if self.dilation != (1,) * len(self.dilation): - s += ", dilation={dilation}" - if self.output_padding != (0,) * len(self.output_padding): - s += ", output_padding={output_padding}" - if self.groups != 1: - s += ", groups={groups}" - if self._bias is None: - s += ", bias=False" - if self.padding_mode != "zeros": - s += ", padding_mode={padding_mode}" - s += ", self.quantize_input={quantize_input}" - s += ", self.quantize_output={quantize_output}" - return s.format(**self.__dict__) +def add_compression_layers(model, config, input_shape=None, add_missing_quantizers=False): + device = next((p.device for p in model.parameters()), torch.device("cpu")) + model = _add_quantized_activations_to_model_layer(model, config) + model = _add_pruning_to_model(model, config) + if add_missing_quantizers: + # Imported here (not at module top) to avoid a circular import: tracing.py + # imports the layer classes defined in this module. + from pquant.core.torch.tracing import check_quantization -def add_compression_layers(model, config, input_shape=None): - model = add_quantized_activations_to_model_layer(model, config) - model = add_pruning_to_model(model, config) - model.to("cuda") + model = check_quantization(model, add_missing_quantizers=True, config=config) + model.to(device) if input_shape is not None: - model(torch.rand(input_shape).to("cuda")) - model.to("cuda") + model(torch.rand(input_shape).to(device)) return model @@ -618,24 +583,15 @@ def __init__( config, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, **kwargs, ): super().__init__(**kwargs) - if in_quant_bits is not None: - self.k_input, self.i_input, self.f_input = in_quant_bits - else: - self.k_input = config.quantization_parameters.default_data_keep_negatives - self.i_input = config.quantization_parameters.default_data_integer_bits - self.f_input = config.quantization_parameters.default_data_fractional_bits - - if out_quant_bits is not None: - self.k_output, self.i_output, self.f_output = out_quant_bits - else: - self.k_output = config.quantization_parameters.default_data_keep_negatives - self.i_output = config.quantization_parameters.default_data_integer_bits - self.f_output = config.quantization_parameters.default_data_fractional_bits + self.k_input, self.i_input, self.f_input = _resolve_data_quant_bits(in_quant_bits, config) + self.k_output, self.i_output, self.f_output = _resolve_data_quant_bits(out_quant_bits, config) self.overflow_mode_data = config.quantization_parameters.overflow_mode_data self.config = config self.is_pretraining = True @@ -649,29 +605,37 @@ def __init__( self.saved_inputs = [] self.quantize_input = quantize_input self.quantize_output = quantize_output + # Optional per-quantizer granularity override; None → inherit config granularity. + granularity = config.quantization_parameters.granularity + self.in_quant_granularity = in_quant_granularity if in_quant_granularity is not None else granularity + self.out_quant_granularity = out_quant_granularity if out_quant_granularity is not None else granularity def build(self, input_shape): self.input_quantizer = Quantizer( - k=torch.tensor(self.k_input), - i=torch.tensor(self.i_input), - f=torch.tensor(self.f_input), + k=self.k_input, + i=self.i_input, + f=self.f_input, overflow=self.overflow_mode_data, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=self.in_quant_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.config.quantization_parameters.dynamic_data_quantization, ) self.output_quantizer = Quantizer( - k=torch.tensor(self.k_output), - i=torch.tensor(self.i_output), - f=torch.tensor(self.f_output), + k=self.k_output, + i=self.i_output, + f=self.f_output, overflow=self.overflow_mode_data, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=self.out_quant_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.config.quantization_parameters.dynamic_data_quantization, ) self.input_shape = (1,) + input_shape[1:] @@ -698,13 +662,13 @@ def hgq_loss(self): loss += self.output_quantizer.hgq_loss() return loss - def is_fitcompress_pretraining(self): + def _is_fitcompress_pretraining(self): return self.is_pretraining and self.use_fitcompress - def pre_pooling(self, x): + def _pre_pooling(self, x): if not hasattr(self, "input_quantizer"): self.build(x.shape) - if self.is_fitcompress_pretraining(): + if self._is_fitcompress_pretraining(): if self.post_fitcompress_calibration: # Save inputs self.saved_inputs.append(x) @@ -714,8 +678,8 @@ def pre_pooling(self, x): x = self.input_quantizer(x) return x - def post_pooling(self, x): - if self.quantize_output and self.enable_quantization and not self.is_fitcompress_pretraining(): + def _post_pooling(self, x): + if self.quantize_output and self.enable_quantization and not self._is_fitcompress_pretraining(): x = self.output_quantizer(x) return x @@ -734,8 +698,10 @@ def __init__( count_include_pad: bool = True, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, **kwargs, ): super().__init__( @@ -749,13 +715,15 @@ def __init__( quantize_output=quantize_output, in_quant_bits=in_quant_bits, out_quant_bits=out_quant_bits, + in_quant_granularity=in_quant_granularity, + out_quant_granularity=out_quant_granularity, **kwargs, ) def forward(self, x): - x = self.pre_pooling(x) + x = self._pre_pooling(x) x = super().forward(x) - x = self.post_pooling(x) + x = self._post_pooling(x) return x @@ -768,11 +736,13 @@ def __init__( padding: _size_2_t = 0, ceil_mode: bool = False, count_include_pad: bool = True, - divisor_override: Optional[int] = None, + divisor_override: int | None = None, quantize_input=True, quantize_output=False, - in_quant_bits: Tuple[T, T, T] = None, - out_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, **kwargs, ): super().__init__( @@ -787,52 +757,35 @@ def __init__( quantize_output=quantize_output, in_quant_bits=in_quant_bits, out_quant_bits=out_quant_bits, + in_quant_granularity=in_quant_granularity, + out_quant_granularity=out_quant_granularity, **kwargs, ) def forward(self, x): - x = self.pre_pooling(x) + x = self._pre_pooling(x) x = super().forward(x) - x = self.post_pooling(x) + x = self._post_pooling(x) return x -class PQBatchNorm2d(nn.BatchNorm2d): - def __init__( +class PQBatchNormBase: + """Quantization behavior shared by PQBatchNorm1d and PQBatchNorm2d.""" + + def _init_quantization( self, config, - num_features: int, - eps: float = 1e-5, - momentum: typing.Optional[float] = 0.1, - affine: bool = True, - track_running_stats: bool = True, - device=None, - dtype=None, - quantize_input=True, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, + quantize_input, + in_quant_bits, + weight_quant_bits, + bias_quant_bits, + in_quant_granularity, + weight_quant_granularity, + bias_quant_granularity, ): - super().__init__(num_features, eps, momentum, affine, track_running_stats, device=device, dtype=dtype) - if in_quant_bits is not None: - self.k_input, self.i_input, self.f_input = in_quant_bits - else: - self.k_input = config.quantization_parameters.default_data_keep_negatives - self.i_input = config.quantization_parameters.default_data_integer_bits - self.f_input = config.quantization_parameters.default_data_fractional_bits - - if weight_quant_bits is not None: - self.k_weight, self.i_weight, self.f_weight = weight_quant_bits - else: - self.k_weight = config.quantization_parameters.default_weight_keep_negatives - self.i_weight = config.quantization_parameters.default_weight_integer_bits - self.f_weight = config.quantization_parameters.default_weight_fractional_bits - if bias_quant_bits is not None: - self.k_bias, self.i_bias, self.f_bias = bias_quant_bits - else: - self.k_bias = config.quantization_parameters.default_weight_keep_negatives - self.i_bias = config.quantization_parameters.default_weight_integer_bits - self.f_bias = config.quantization_parameters.default_weight_fractional_bits + self.k_input, self.i_input, self.f_input = _resolve_data_quant_bits(in_quant_bits, config) + self.k_weight, self.i_weight, self.f_weight = _resolve_weight_quant_bits(weight_quant_bits, config) + self.k_bias, self.i_bias, self.f_bias = _resolve_weight_quant_bits(bias_quant_bits, config) self.overflow_mode_parameters = config.quantization_parameters.overflow_mode_parameters self.overflow_mode_data = config.quantization_parameters.overflow_mode_data self.round_mode = config.quantization_parameters.round_mode @@ -843,6 +796,10 @@ def __init__( self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress self.config = config self.quantize_input = quantize_input + granularity = config.quantization_parameters.granularity + self.in_quant_granularity = in_quant_granularity if in_quant_granularity is not None else granularity + self.weight_quant_granularity = weight_quant_granularity if weight_quant_granularity is not None else granularity + self.bias_quant_granularity = bias_quant_granularity if bias_quant_granularity is not None else granularity self._weight = nn.Parameter(self.weight.clone()).to(self.weight.device) self.register_parameter("_weight", self._weight) if self.bias is not None: @@ -851,45 +808,51 @@ def __init__( else: self.register_parameter("_bias", None) self.built = False - self.final_compression_done = False + self.register_buffer("final_compression_done", torch.tensor(False)) self.is_pretraining = True self.post_fitcompress_calibration = False self.saved_inputs = [] - def check_is_built(self, input_shape): + def _check_is_built(self, input_shape): if self.built: return self.built = True self.input_quantizer = Quantizer( - k=torch.tensor(self.k_input), - i=torch.tensor(self.i_input), - f=torch.tensor(self.f_input), + k=self.k_input, + i=self.i_input, + f=self.f_input, overflow=self.overflow_mode_data, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=self.in_quant_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.config.quantization_parameters.dynamic_data_quantization, ) self.weight_quantizer = Quantizer( - k=torch.tensor(self.k_weight), - i=torch.tensor(self.i_weight), - f=torch.tensor(self.f_weight), - round_mode=self.round_mode, + k=self.k_weight, + i=self.i_weight, + f=self.f_weight, overflow=self.overflow_mode_parameters, - is_data=False, + round_mode=self.round_mode, is_heterogeneous=self.use_hgq, + is_data=False, + granularity=self.weight_quant_granularity, place="weight", + shape=self._weight.shape, ) self.bias_quantizer = Quantizer( - k=torch.tensor(self.k_bias), - i=torch.tensor(self.i_bias), - f=torch.tensor(self.f_bias), - round_mode=self.round_mode, + k=self.k_bias, + i=self.i_bias, + f=self.f_bias, overflow=self.overflow_mode_parameters, - is_data=False, + round_mode=self.round_mode, is_heterogeneous=self.use_hgq, + is_data=False, + granularity=self.bias_quant_granularity, place="bias", + shape=None if self._bias is None else self._bias.shape, ) if self.use_hgq: self.input_quantizer.quantizer.build(input_shape) @@ -899,9 +862,9 @@ def check_is_built(self, input_shape): self.input_shape = (1,) + input_shape[1:] def apply_final_compression(self): - self.final_compression_done = True self._weight.data = self.weight self._bias.data = self.bias + self.final_compression_done.fill_(True) def get_input_quantization_bits(self): return self.input_quantizer.get_quantization_bits() @@ -912,32 +875,32 @@ def get_weight_quantization_bits(self): def get_bias_quantization_bits(self): return self.bias_quantizer.get_quantization_bits() - def is_fitcompress_pretraining(self): + def _is_fitcompress_pretraining(self): return self.is_pretraining and self.use_fitcompress @property def weight(self): - if self.enable_quantization and not self.final_compression_done and not self.is_fitcompress_pretraining(): + if self.enable_quantization and not self.final_compression_done and not self._is_fitcompress_pretraining(): return self.weight_quantizer(self._weight) return self._weight @property def bias(self): - if self.enable_quantization and not self.final_compression_done and not self.is_fitcompress_pretraining(): + if self.enable_quantization and not self.final_compression_done and not self._is_fitcompress_pretraining(): return self.bias_quantizer(self._bias) return self._bias def ebops(self): bw_inp = self.input_quantizer.get_total_bits(self.input_shape) - bw_ker = ops.reshape(self.weight_quantizer.get_total_bits(self.running_mean.shape), self._shape) - bw_bias = ops.reshape(self.bias_quantizer.get_total_bits(self.running_mean.shape), self._shape) - size = ops.cast(ops.prod(list(self.input_shape)), self._weight.dtype) - ebops = ops.sum(bw_inp * bw_ker) + ops.mean(bw_bias) * size + bw_ker = torch.reshape(self.weight_quantizer.get_total_bits(self.running_mean.shape), self._shape) + bw_bias = torch.reshape(self.bias_quantizer.get_total_bits(self.running_mean.shape), self._shape) + size = float(math.prod(self.input_shape)) + ebops = torch.sum(bw_inp * bw_ker) + torch.mean(bw_bias) * size return ebops def hgq_loss(self): if self.is_pretraining or not self.use_hgq: - return ops.convert_to_tensor(0.0) + return torch.tensor(0.0) loss = self.hgq_beta * self.ebops() loss += self.weight_quantizer.hgq_loss() loss += self.bias_quantizer.hgq_loss() @@ -949,52 +912,109 @@ def post_pre_train_function(self): self.is_pretraining = False def forward(self, input: torch.Tensor) -> torch.Tensor: - self.check_is_built(input.shape) + self._check_is_built(input.shape) if self.quantize_input and self.enable_quantization: - if not self.is_fitcompress_pretraining(): + if not self._is_fitcompress_pretraining(): input = self.input_quantizer(input) - else: - if self.post_fitcompress_calibration: - self.saved_inputs.append(input) + elif self.post_fitcompress_calibration: + self.saved_inputs.append(input) return super().forward(input) -class PQBatchNorm1d(nn.BatchNorm1d): +class PQBatchNorm2d(PQBatchNormBase, nn.BatchNorm2d): def __init__( self, config, num_features: int, eps: float = 1e-5, - momentum: typing.Optional[float] = 0.1, + momentum: float | None = 0.1, affine: bool = True, track_running_stats: bool = True, device=None, dtype=None, quantize_input=True, - in_quant_bits: Tuple[T, T, T] = None, - weight_quant_bits: Tuple[T, T, T] = None, - bias_quant_bits: Tuple[T, T, T] = None, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + weight_quant_granularity=None, + bias_quant_granularity=None, ): super().__init__(num_features, eps, momentum, affine, track_running_stats, device=device, dtype=dtype) - if in_quant_bits is not None: - self.k_input, self.i_input, self.f_input = in_quant_bits - else: - self.k_input = config.quantization_parameters.default_data_keep_negatives - self.i_input = config.quantization_parameters.default_data_integer_bits - self.f_input = config.quantization_parameters.default_data_fractional_bits + self._init_quantization( + config, + quantize_input, + in_quant_bits, + weight_quant_bits, + bias_quant_bits, + in_quant_granularity, + weight_quant_granularity, + bias_quant_granularity, + ) - if weight_quant_bits is not None: - self.k_weight, self.i_weight, self.f_weight = weight_quant_bits - else: - self.k_weight = config.quantization_parameters.default_weight_keep_negatives - self.i_weight = config.quantization_parameters.default_weight_integer_bits - self.f_weight = config.quantization_parameters.default_weight_fractional_bits - if bias_quant_bits is not None: - self.k_bias, self.i_bias, self.f_bias = bias_quant_bits - else: - self.k_bias = config.quantization_parameters.default_weight_keep_negatives - self.i_bias = config.quantization_parameters.default_weight_integer_bits - self.f_bias = config.quantization_parameters.default_weight_fractional_bits + +class PQBatchNorm1d(PQBatchNormBase, nn.BatchNorm1d): + def __init__( + self, + config, + num_features: int, + eps: float = 1e-5, + momentum: float | None = 0.1, + affine: bool = True, + track_running_stats: bool = True, + device=None, + dtype=None, + quantize_input=True, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + weight_quant_granularity=None, + bias_quant_granularity=None, + ): + super().__init__(num_features, eps, momentum, affine, track_running_stats, device=device, dtype=dtype) + self._init_quantization( + config, + quantize_input, + in_quant_bits, + weight_quant_bits, + bias_quant_bits, + in_quant_granularity, + weight_quant_granularity, + bias_quant_granularity, + ) + + +class PQLayerNorm(nn.LayerNorm): + def __init__( + self, + config, + normalized_shape: int | tuple[int, ...] | torch.Size, + eps: float = 1e-5, + elementwise_affine: bool = True, + bias: bool = True, + device=None, + dtype=None, + quantize_input=True, + quantize_output=False, + in_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, + weight_quant_granularity=None, + bias_quant_granularity=None, + ): + try: + super().__init__(normalized_shape, eps, elementwise_affine, bias, device=device, dtype=dtype) + except TypeError: + # Older torch versions don't accept the bias kwarg + super().__init__(normalized_shape, eps, elementwise_affine, device=device, dtype=dtype) + self.k_input, self.i_input, self.f_input = _resolve_data_quant_bits(in_quant_bits, config) + self.k_output, self.i_output, self.f_output = _resolve_data_quant_bits(out_quant_bits, config) + self.k_weight, self.i_weight, self.f_weight = _resolve_weight_quant_bits(weight_quant_bits, config) + self.k_bias, self.i_bias, self.f_bias = _resolve_weight_quant_bits(bias_quant_bits, config) self.overflow_mode_parameters = config.quantization_parameters.overflow_mode_parameters self.overflow_mode_data = config.quantization_parameters.overflow_mode_data self.round_mode = config.quantization_parameters.round_mode @@ -1005,242 +1025,413 @@ def __init__( self.use_fitcompress = config.fitcompress_parameters.enable_fitcompress self.config = config self.quantize_input = quantize_input - self._weight = nn.Parameter(self.weight.clone()).to(self.weight.device) - self.register_parameter("_weight", self._weight) + self.quantize_output = quantize_output + granularity = config.quantization_parameters.granularity + self.in_quant_granularity = in_quant_granularity if in_quant_granularity is not None else granularity + self.out_quant_granularity = out_quant_granularity if out_quant_granularity is not None else granularity + self.weight_quant_granularity = weight_quant_granularity if weight_quant_granularity is not None else granularity + self.bias_quant_granularity = bias_quant_granularity if bias_quant_granularity is not None else granularity + if self.weight is not None: + self._weight = nn.Parameter(self.weight.clone()).to(self.weight.device) + self.register_parameter("_weight", self._weight) + else: + self.register_parameter("_weight", None) if self.bias is not None: self._bias = nn.Parameter(self.bias.clone()).to(self.bias.device) self.register_parameter("_bias", self._bias) else: self.register_parameter("_bias", None) - self.register_parameter("_weight", self._weight) self.built = False - self.final_compression_done = False + self.register_buffer("final_compression_done", torch.tensor(False)) self.is_pretraining = True self.post_fitcompress_calibration = False self.saved_inputs = [] - def check_is_built(self, input_shape): + def _check_is_built(self, input_shape): if self.built: return self.built = True self.input_quantizer = Quantizer( - k=torch.tensor(self.k_input), - i=torch.tensor(self.i_input), - f=torch.tensor(self.f_input), + k=self.k_input, + i=self.i_input, + f=self.f_input, overflow=self.overflow_mode_data, round_mode=self.round_mode, is_heterogeneous=self.use_hgq, is_data=True, + granularity=self.in_quant_granularity, hgq_gamma=self.hgq_gamma, place="datalane", + dynamic_data=self.config.quantization_parameters.dynamic_data_quantization, ) - self.weight_quantizer = Quantizer( - k=torch.tensor(self.k_weight), - i=torch.tensor(self.i_weight), - f=torch.tensor(self.f_weight), + self.output_quantizer = Quantizer( + k=self.k_output, + i=self.i_output, + f=self.f_output, + overflow=self.overflow_mode_data, round_mode=self.round_mode, + is_heterogeneous=self.use_hgq, + is_data=True, + granularity=self.out_quant_granularity, + hgq_gamma=self.hgq_gamma, + place="datalane", + dynamic_data=self.config.quantization_parameters.dynamic_data_quantization, + ) + self.weight_quantizer = Quantizer( + k=self.k_weight, + i=self.i_weight, + f=self.f_weight, overflow=self.overflow_mode_parameters, - is_data=False, + round_mode=self.round_mode, is_heterogeneous=self.use_hgq, + is_data=False, + granularity=self.weight_quant_granularity, place="weight", + shape=self._weight.shape, ) self.bias_quantizer = Quantizer( - k=torch.tensor(self.k_bias), - i=torch.tensor(self.i_bias), - f=torch.tensor(self.f_bias), - round_mode=self.round_mode, + k=self.k_bias, + i=self.i_bias, + f=self.f_bias, overflow=self.overflow_mode_parameters, - is_data=False, + round_mode=self.round_mode, is_heterogeneous=self.use_hgq, + is_data=False, + granularity=self.bias_quant_granularity, place="bias", + shape=None if self._bias is None else self._bias.shape, ) if self.use_hgq: self.input_quantizer.quantizer.build(input_shape) - shape = [1] * len(input_shape) - shape[1] = input_shape[1] - self._shape = tuple(shape) - self.input_shape = (1,) + input_shape[1:] + self.output_quantizer.quantizer.build(input_shape) + self.input_shape = (1,) + tuple(input_shape[1:]) def apply_final_compression(self): - self.final_compression_done = True - self._weight.data = self.weight - self._bias.data = self.bias + if self._weight is not None: + self._weight.data = self.weight + if self._bias is not None: + self._bias.data = self.bias + self.final_compression_done.fill_(True) def get_input_quantization_bits(self): return self.input_quantizer.get_quantization_bits() + def get_output_quantization_bits(self): + return self.output_quantizer.get_quantization_bits() + def get_weight_quantization_bits(self): return self.weight_quantizer.get_quantization_bits() def get_bias_quantization_bits(self): return self.bias_quantizer.get_quantization_bits() - def is_fitcompress_pretraining(self): + def _is_fitcompress_pretraining(self): return self.is_pretraining and self.use_fitcompress @property def weight(self): - if self.enable_quantization and not self.final_compression_done and not self.is_fitcompress_pretraining(): + if self._weight is None: + return None + if self.enable_quantization and not self.final_compression_done and not self._is_fitcompress_pretraining(): return self.weight_quantizer(self._weight) return self._weight @property def bias(self): - if self.enable_quantization and not self.final_compression_done and not self.is_fitcompress_pretraining(): + if self._bias is None: + return None + if self.enable_quantization and not self.final_compression_done and not self._is_fitcompress_pretraining(): return self.bias_quantizer(self._bias) return self._bias def ebops(self): - bw_inp = self.input_quantizer.get_total_bits(self.input_shape) - bw_ker = ops.reshape(self.weight_quantizer.get_total_bits(self.running_mean.shape), self._shape) - bw_bias = ops.reshape(self.bias_quantizer.get_total_bits(self.running_mean.shape), self._shape) - size = ops.cast(ops.prod(list(self.input_shape)), self._weight.dtype) - ebops = ops.sum(bw_inp * bw_ker) + ops.mean(bw_bias) * size - return ebops + return 0.0 def hgq_loss(self): if self.is_pretraining or not self.use_hgq: - return ops.convert_to_tensor(0.0) + return torch.tensor(0.0) loss = self.hgq_beta * self.ebops() - loss += self.weight_quantizer.hgq_loss() - loss += self.bias_quantizer.hgq_loss() + if self._weight is not None: + loss += self.weight_quantizer.hgq_loss() + if self._bias is not None: + loss += self.bias_quantizer.hgq_loss() if self.quantize_input: loss += self.input_quantizer.hgq_loss() + if self.quantize_output: + loss += self.output_quantizer.hgq_loss() return loss - def post_pretrain_function(self): + def post_pre_train_function(self): self.is_pretraining = False def forward(self, input: torch.Tensor) -> torch.Tensor: - self.check_is_built(input.shape) + self._check_is_built(input.shape) if self.quantize_input and self.enable_quantization: - if not self.is_fitcompress_pretraining(): + if not self._is_fitcompress_pretraining(): input = self.input_quantizer(input) - else: - if self.post_fitcompress_calibration: - self.saved_inputs.append(input) - return super().forward(input) + elif self.post_fitcompress_calibration: + self.saved_inputs.append(input) + out = F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps) + if self.quantize_output and self.enable_quantization and not self._is_fitcompress_pretraining(): + out = self.output_quantizer(out) + return out + + def extra_repr(self) -> str: + return ( + f"normalized_shape={tuple(self.normalized_shape)}, eps={self.eps}, " + f"elementwise_affine={self.elementwise_affine}, " + f"quantize_input={self.quantize_input}, quantize_output={self.quantize_output}" + ) + + +class PQMultiheadAttention(nn.Module): + """Multi-head attention with quantization support, implemented without F.multihead_attention. + + Uses separate PQDense projections for Q, K, V, and output, and computes + scaled dot-product attention manually. + + Args: + config: PQuant configuration object. + embed_dim: Total embedding dimension. + num_heads: Number of attention heads. + dropout: Dropout probability on attention weights. + bias: Whether to add bias to projection layers. + kdim: Key feature dimension (defaults to embed_dim). + vdim: Value feature dimension (defaults to embed_dim). + batch_first: If True, input/output tensors are (batch, seq, feature). + If False (default), tensors are (seq, batch, feature). + quantize_input: Whether to quantize the Q/K/V projection inputs (the MHA inputs). + quantize_output: Whether to quantize the output projection's output (the MHA output). + The q/k/v projection outputs and the out_proj input (the context) are always + quantized, mirroring HGQ's QMultiHeadAttention. + in_quant_bits: (k, i, f) bits for input quantization. + weight_quant_bits: (k, i, f) bits for weight quantization. + bias_quant_bits: (k, i, f) bits for bias quantization. + out_quant_bits: (k, i, f) bits for output quantization. + attn_quant_bits: (k, i, f) bits for the softmax output quantizer (the attention + weights). The scores and context need no dedicated quantizers: the softmax's + input quantizer and the output projection's input quantizer cover them. + """ + + def __init__( + self, + config, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + bias: bool = True, + kdim: int = None, + vdim: int = None, + batch_first: bool = False, + quantize_input: bool = True, + quantize_output: bool = False, + approximate_softmax: bool = False, + in_quant_bits: tuple[T, T, T] = None, + weight_quant_bits: tuple[T, T, T] = None, + bias_quant_bits: tuple[T, T, T] = None, + out_quant_bits: tuple[T, T, T] = None, + attn_quant_bits: tuple[T, T, T] = None, + in_quant_granularity=None, + out_quant_granularity=None, + param_quant_granularity=None, + **kwargs, + ): + super().__init__(**kwargs) + assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads" + + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.batch_first = batch_first + self.dropout = dropout + self.approximate_softmax = approximate_softmax + self.scale = float(torch.tensor(self.head_dim**-0.5, dtype=torch.float32).item()) + self.softmax = PQSoftmax(config, -1, quantize_input=True, quantize_output=True, out_quant_bits=attn_quant_bits) + + kdim = kdim if kdim is not None else embed_dim + vdim = vdim if vdim is not None else embed_dim + + self.in_quant_granularity = in_quant_granularity + self.out_quant_granularity = out_quant_granularity + self.param_quant_granularity = param_quant_granularity + proj_kwargs = dict( + bias=bias, + in_quant_bits=in_quant_bits, + weight_quant_bits=weight_quant_bits, + bias_quant_bits=bias_quant_bits, + out_quant_bits=out_quant_bits, + weight_quant_granularity=param_quant_granularity, + bias_quant_granularity=param_quant_granularity, + ) + qkv_kwargs = dict( + quantize_input=quantize_input, quantize_output=True, in_quant_granularity=in_quant_granularity, **proj_kwargs + ) + self.q_proj = PQDense(config, embed_dim, embed_dim, enable_pruning=False, **qkv_kwargs) + self.k_proj = PQDense(config, kdim, embed_dim, enable_pruning=False, **qkv_kwargs) + self.v_proj = PQDense(config, vdim, embed_dim, enable_pruning=False, **qkv_kwargs) + self.out_proj = PQDense( + config, + embed_dim, + embed_dim, + quantize_input=True, + quantize_output=quantize_output, + out_quant_granularity=out_quant_granularity, + **proj_kwargs, + ) -def add_layer_specific_quantization_to_model(name, layer, config): - if isinstance(layer, PQWeightBiasBase): - if name in config.quantization_parameters.layer_specific: - layer_config = config.quantization_parameters.layer_specific[name] - if "weight" in layer_config: - if "keep_negatives" in layer_config["weight"]: - layer.k_weight = torch.tensor(layer_config["weight"]["keep_negatives"]) - if "integer_bits" in layer_config["weight"]: - layer.i_weight = torch.tensor(layer_config["weight"]["integer_bits"]) - if "fractional_bits" in layer_config["weight"]: - layer.f_weight = torch.tensor(layer_config["weight"]["fractional_bits"]) - if "bias" in layer_config: - if "keep_negatives" in layer_config["bias"]: - layer.k_bias = torch.tensor(layer_config["bias"]["keep_negatives"]) - if "integer_bits" in layer_config["bias"]: - layer.i_bias = torch.tensor(layer_config["bias"]["integer_bits"]) - if "fractional_bits" in layer_config["bias"]: - layer.f_bias = torch.tensor(layer_config["bias"]["fractional_bits"]) - if "input" in layer_config: - if "keep_negatives" in layer_config["input"]: - input_keep_negatives = torch.tensor(layer_config["input"]["keep_negatives"]) - layer.k_input = input_keep_negatives - if "integer_bits" in layer_config["input"]: - input_int_bits = torch.tensor(layer_config["input"]["integer_bits"]) - layer.i_input = input_int_bits - if "fractional_bits" in layer_config["input"]: - input_fractional_bits = torch.tensor(layer_config["input"]["fractional_bits"]) - layer.f_input = input_fractional_bits - if "quantize" in layer_config["input"]: - quantize = layer_config["input"]["quantize"] - layer.quantize_input = quantize - if "output" in layer_config: - if "keep_negatives" in layer_config["input"]: - output_keep_negatives = torch.tensor(layer_config["output"]["keep_negatives"]) - layer.k_output = output_keep_negatives - if "integer_bits" in layer_config["output"]: - output_int_bits = torch.tensor(layer_config["output"]["integer_bits"]) - layer.i_output = input_int_bits - if "fractional_bits" in layer_config["output"]: - input_fractional_bits = torch.tensor(layer_config["output"]["fractional_bits"]) - layer.f_output = input_fractional_bits - if "quantize" in layer_config["output"]: - quantize = layer_config["output"]["quantize"] - layer.quantize_output = quantize - - elif layer.__class__ in [PQBatchNorm2d, PQBatchNorm1d]: - if name in config.quantization_parameters.layer_specific: - layer_config = config.quantization_parameters.layer_specific[name] - if "weight" in layer_config: - i = torch.tensor(layer_config["weight"]["integer_bits"]) - f = torch.tensor(layer_config["weight"]["fractional_bits"]) - layer.i_weight = i - layer.f_weight = f - if "bias" in layer_config: - i = torch.tensor(layer_config["bias"]["integer_bits"]) - f = torch.tensor(layer_config["bias"]["fractional_bits"]) - layer.i_bias = i - layer.f_bias = f - if "input" in layer_config: - if "integer_bits" in layer_config["input"]: - input_int_bits = torch.tensor(layer_config["input"]["integer_bits"]) - layer.i_input = input_int_bits - if "fractional_bits" in layer_config["input"]: - input_fractional_bits = torch.tensor(layer_config["input"]["fractional_bits"]) - layer.f_input = input_fractional_bits - if "quantize" in layer_config["input"]: - quantize = layer_config["input"]["quantize"] - layer.quantize_input = quantize - elif layer.__class__ in [PQAvgPool1d, PQAvgPool2d]: - if name in config.quantization_parameters.layer_specific: - layer_config = config.quantization_parameters.layer_specific[name] - if "input" in layer_config: - if "integer_bits" in layer_config["input"]: - input_int_bits = torch.tensor(layer_config["input"]["integer_bits"]) - layer.i_input = input_int_bits - if "fractional_bits" in layer_config["input"]: - input_fractional_bits = torch.tensor(layer_config["input"]["fractional_bits"]) - layer.f_input = input_fractional_bits - if "quantize" in layer_config["input"]: - quantize = layer_config["input"]["quantize"] - layer.quantize_input = quantize - if "output" in layer_config: - if "integer_bits" in layer_config["output"]: - output_int_bits = torch.tensor(layer_config["output"]["integer_bits"]) - layer.i_output = output_int_bits - if "fractional_bits" in layer_config["output"]: - output_fractional_bits = torch.tensor(layer_config["output"]["fractional_bits"]) - layer.f_output = output_fractional_bits - if "quantize" in layer_config["output"]: - quantize = layer_config["output"]["quantize"] - layer.quantize_output = quantize - - elif layer.__class__ == PQActivation: - if name in config.quantization_parameters.layer_specific: - layer_config = config.quantization_parameters.layer_specific[name] - if "input" in layer_config: - if "integer_bits" in layer_config["input"]: - input_int_bits = torch.tensor(layer_config["input"]["integer_bits"]) - layer.i_input = input_int_bits - if "fractional_bits" in layer_config["input"]: - input_fractional_bits = torch.tensor(layer_config["input"]["fractional_bits"]) - layer.f_input = input_fractional_bits - if "quantize" in layer_config["input"]: - quantize = layer_config["input"]["quantize"] - layer.quantize_input = quantize - if "output" in layer_config: - if "integer_bits" in layer_config["output"]: - output_int_bits = torch.tensor(layer_config["output"]["integer_bits"]) - layer.i_output = output_int_bits - if "fractional_bits" in layer_config["output"]: - output_fractional_bits = torch.tensor(layer_config["output"]["fractional_bits"]) - layer.f_output = output_fractional_bits - if "quantize" in layer_config["output"]: - quantize = layer_config["output"]["quantize"] - layer.quantize_output = quantize + self.attn_dropout = None + self.enable_quantization = config.quantization_parameters.enable_quantization + self.use_hgq = config.quantization_parameters.use_high_granularity_quantization + self.hgq_beta = config.quantization_parameters.hgq_beta + self.is_pretraining = True + + def post_pre_train_function(self): + # The projections, softmax and data quantizers are handled separately by the + # recursive modules() walk in post_pretrain_functions. + self.is_pretraining = False + + def _head_bits(self, proj, seq_len): + """Bitwidths of a projection's output, in per-head layout (1, H, seq, head_dim).""" + bw = proj.output_quantizer.get_total_bits((1, seq_len, self.embed_dim)) + return bw.reshape(1, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + + def attention_ebops(self): + """EBOPs of the q @ k^T and attn @ v einsums (mirrors HGQ's QMultiHeadAttention).""" + attn_shape = self.softmax.input_shape # (1, H, T, S), stored when the softmax was built + query_len, key_len = attn_shape[2], attn_shape[3] + bw_q = self._head_bits(self.q_proj, query_len) + bw_k = self._head_bits(self.k_proj, key_len) + bw_v = self._head_bits(self.v_proj, key_len) + bw_attn = self.softmax.output_quantizer.get_total_bits(attn_shape) + ebops_qk = torch.einsum("bhtd,bhsd->", bw_q, bw_k) + ebops_av = torch.einsum("bhts,bhsd->", bw_attn, bw_v) + return ebops_qk + ebops_av + + def ebops(self): + # Only the attention einsum costs are this module's own: the projections, + # softmax and lookup tables are counted by get_ebops's recursive modules() walk. + return self.attention_ebops() + + def hgq_loss(self): + if self.is_pretraining or not self.use_hgq: + return torch.tensor(0.0) + return self.hgq_beta * self.attention_ebops() + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_padding_mask: torch.Tensor | None = None, + attn_mask: torch.Tensor | None = None, + need_weights: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not self.batch_first: + # (T, B, E) -> (B, T, E) + query = query.transpose(0, 1) + key = key.transpose(0, 1) + value = value.transpose(0, 1) + + B, T = query.shape[0], query.shape[1] + S = key.shape[1] + + q = self.q_proj(query) # (B, T, E) + k = self.k_proj(key) # (B, S, E) + v = self.v_proj(value) # (B, S, E) + + q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) + + attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale + + if attn_mask is not None: + if attn_mask.dim() == 2: + attn_mask = attn_mask.unsqueeze(0).unsqueeze(0) + elif attn_mask.dim() == 3: + attn_mask = attn_mask.view(B, self.num_heads, T, S) + attn_scores = attn_scores + attn_mask + + mask = None + if key_padding_mask is not None: + mask = ~key_padding_mask.unsqueeze(1).unsqueeze(2) # (B, 1, 1, S) + + # The softmax's own input/output quantizers handle the scores and the attention weights; + attn_weights = self.softmax(attn_scores, mask=mask) + + if self.attn_dropout is not None and self.training: + attn_weights = self.attn_dropout(attn_weights) + + out = torch.matmul(attn_weights, v) + + # Merge heads: (B, T, E) + out = out.transpose(1, 2).contiguous().view(B, T, self.embed_dim) + out = self.out_proj(out) + + if not self.batch_first: + out = out.transpose(0, 1) + + if need_weights: + # Average attention weights over heads: (B, T, S) + return out, attn_weights.mean(dim=1) + return out, None + + def extra_repr(self) -> str: + return ( + f"embed_dim={self.embed_dim}, num_heads={self.num_heads}, dropout={self.dropout}, batch_first={self.batch_first}" + ) + + +# Prunable leaf layers, larger layers like MHA not included as they consist of these layers +LAYERS_WITH_PRUNING_LAYER = (PQConv2d, PQConv1d, PQDense) +PQ_MODULES = ( + PQConv2d, + PQConv1d, + PQDense, + PQActivation, + PQBatchNorm2d, + PQBatchNorm1d, + PQLayerNorm, + PQAvgPoolBase, + PQSoftmax, + PQMultiheadAttention, + Quantizer, +) + + +def _apply_quant_bits(layer, section_config, suffix): + """Copy keep_negatives/integer/fractional bits from one config section onto layer.{k,i,f}_{suffix}.""" + if "keep_negatives" in section_config: + setattr(layer, f"k_{suffix}", torch.tensor(section_config["keep_negatives"])) + if "integer_bits" in section_config: + setattr(layer, f"i_{suffix}", torch.tensor(section_config["integer_bits"])) + if "fractional_bits" in section_config: + setattr(layer, f"f_{suffix}", torch.tensor(section_config["fractional_bits"])) + + +def _add_layer_specific_quantization_to_model(name, layer, config): + layer_config = config.quantization_parameters.layer_specific.get(name) + if layer_config is None: + return layer + if isinstance(layer, (PQWeightBiasBase, PQLayerNorm)): + sections = ("weight", "bias", "input", "output") + elif isinstance(layer, (PQBatchNormBase)): + sections = ("weight", "bias", "input") + elif isinstance(layer, (PQAvgPoolBase, PQActivation)): + sections = ("input", "output") + else: + return layer + for section in sections: + if section not in layer_config: + continue + _apply_quant_bits(layer, layer_config[section], section) + if section in ("input", "output") and "quantize" in layer_config[section]: + setattr(layer, f"quantize_{section}", layer_config[section]["quantize"]) return layer -def add_quantized_activations_to_model_layer(module, config, prefix=""): +def _add_quantized_activations_to_model_layer(module, config, prefix=""): if not config.quantization_parameters.enable_quantization: return module quantize_input = config.quantization_parameters.quantize_input @@ -1262,7 +1453,7 @@ def add_quantized_activations_to_model_layer(module, config, prefix=""): quantize_input=quantize_input, quantize_output=quantize_output, ) - relu = add_layer_specific_quantization_to_model(full_name, relu, config) + relu = _add_layer_specific_quantization_to_model(full_name, relu, config) setattr(module, name, relu) elif layer.__class__ in [nn.Tanh]: type_of_tanh = "tanh" if config.quantization_parameters.use_real_tanh else "hard_tanh" @@ -1274,7 +1465,7 @@ def add_quantized_activations_to_model_layer(module, config, prefix=""): quantize_input=quantize_input, quantize_output=quantize_output, ) - tanh = add_layer_specific_quantization_to_model(full_name, tanh, config) + tanh = _add_layer_specific_quantization_to_model(full_name, tanh, config) setattr(module, name, tanh) elif layer.__class__ == nn.AvgPool1d: new_layer = PQAvgPool1d( @@ -1287,7 +1478,7 @@ def add_quantized_activations_to_model_layer(module, config, prefix=""): quantize_input, quantize_output, ) - new_layer = add_layer_specific_quantization_to_model(full_name, new_layer, config) + new_layer = _add_layer_specific_quantization_to_model(full_name, new_layer, config) setattr(module, name, new_layer) elif layer.__class__ == nn.AvgPool2d: new_layer = PQAvgPool2d( @@ -1301,10 +1492,11 @@ def add_quantized_activations_to_model_layer(module, config, prefix=""): quantize_input, quantize_output, ) - new_layer = add_layer_specific_quantization_to_model(full_name, new_layer, config) + new_layer = _add_layer_specific_quantization_to_model(full_name, new_layer, config) setattr(module, name, new_layer) - elif layer.__class__ == nn.BatchNorm2d: - new_layer = PQBatchNorm2d( + elif layer.__class__ in (nn.BatchNorm1d, nn.BatchNorm2d): + pq_batchnorm = PQBatchNorm1d if layer.__class__ is nn.BatchNorm1d else PQBatchNorm2d + new_layer = pq_batchnorm( config, num_features=layer.num_features, eps=layer.eps, @@ -1313,66 +1505,45 @@ def add_quantized_activations_to_model_layer(module, config, prefix=""): track_running_stats=layer.track_running_stats, quantize_input=quantize_input, ) - new_layer = add_layer_specific_quantization_to_model(full_name, new_layer, config) + new_layer = _add_layer_specific_quantization_to_model(full_name, new_layer, config) setattr(module, name, new_layer) - elif layer.__class__ == nn.BatchNorm1d: - new_layer = PQBatchNorm1d( - config, - num_features=layer.num_features, + elif layer.__class__ == nn.LayerNorm: + ln_kwargs = dict( + normalized_shape=layer.normalized_shape, eps=layer.eps, - momentum=layer.momentum, - affine=layer.affine, - track_running_stats=layer.track_running_stats, + elementwise_affine=layer.elementwise_affine, quantize_input=quantize_input, + quantize_output=quantize_output, ) - new_layer = add_layer_specific_quantization_to_model(full_name, new_layer, config) + new_layer = PQLayerNorm(config, **ln_kwargs) + if layer.elementwise_affine: + if layer.weight is not None and new_layer._weight is not None: + new_layer._weight.data.copy_(layer.weight.data) + if layer.bias is not None and new_layer._bias is not None: + new_layer._bias.data.copy_(layer.bias.data) + new_layer = _add_layer_specific_quantization_to_model(full_name, new_layer, config) setattr(module, name, new_layer) else: - layer = add_quantized_activations_to_model_layer(layer, config, full_name) + layer = _add_quantized_activations_to_model_layer(layer, config, full_name) return module -def add_quantized_activations_to_model_functional(module, config): - # Currently not in use. TODO: Fix this - if config.quantization_parameters.use_high_granularity_quantization: - return module - # Replaces functional activation calls with quantized versions - traced_model = symbolic_trace(module) - for node in traced_model.graph.nodes: - if node.op in ["call_method", "call_function"] and (node.target == "tanh" or "function relu" in str(node.target)): - with traced_model.graph.inserting_after(node): - if node.name in config.quantization_parameters.layer_specific: - bits = config.quantization_parameters.layer_specific[node.name]["bits"] - else: - bits = ( - config.quantization_parameters.default_integer_bits - + config.quantization_parameters.default_fractional_bits - + 1 - ) # 1 sign bit - kwargs = {"bits": bits} - if node.target == "tanh": - kwargs["use_real_tanh"] = config.quantization_parameters.use_real_tanh - kwargs["use_symmetric"] = config.quantization_parameters.use_symmetric_quantization - # new_node = traced_model.graph.call_function(quantized_tanh, node.args, kwargs) - else: - kwargs = {"integer_bits": config.quantization_parameters.default_integer_bits, "bits": bits} - # new_node = traced_model.graph.call_function(quantized_relu, node.args, kwargs) - # node.replace_all_uses_with(new_node) - traced_model.graph.erase_node(node) - - traced_model.graph.lint() - traced_model.recompile() - return traced_model - - -def disable_pruning_from_layers(name, layer, config): - enable_pruning = name not in config.pruning_parameters.disable_pruning_for_layers - if layer.__class__ in [PQDense, PQConv2d, PQConv1d] and not enable_pruning: - layer.enable_pruning = enable_pruning +def _disable_pruning_from_layers(name, layer, config): + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER) and name in config.pruning_parameters.disable_pruning_for_layers: + layer.enable_pruning = False return layer -def add_pruning_to_model(module, config, prefix=""): +def _replace_layer_with_pq_layer(module, name, full_name, layer, sparse_layer, config): + sparse_layer._weight.data = layer.weight.data + if layer.bias is not None: + sparse_layer._bias.data = layer.bias.data + sparse_layer = _add_layer_specific_quantization_to_model(full_name, sparse_layer, config) + sparse_layer = _disable_pruning_from_layers(full_name, sparse_layer, config) + setattr(module, name, sparse_layer) + + +def _add_pruning_to_model(module, config, prefix=""): quantize_input = config.quantization_parameters.quantize_input quantize_output = config.quantization_parameters.quantize_output for name, layer in module.named_children(): @@ -1381,15 +1552,10 @@ def add_pruning_to_model(module, config, prefix=""): sparse_layer = PQDense( config, layer.in_features, layer.out_features, layer.bias is not None, quantize_input, quantize_output ) - sparse_layer._weight.data = layer.weight.data - if layer.bias is not None: - sparse_layer._bias.data = layer.bias.data - - sparse_layer = add_layer_specific_quantization_to_model(full_name, sparse_layer, config) - sparse_layer = disable_pruning_from_layers(full_name, sparse_layer, config) - setattr(module, name, sparse_layer) - elif layer.__class__ is nn.Conv2d: - sparse_layer = PQConv2d( + _replace_layer_with_pq_layer(module, name, full_name, layer, sparse_layer, config) + elif layer.__class__ in (nn.Conv1d, nn.Conv2d): + pq_conv = PQConv1d if layer.__class__ is nn.Conv1d else PQConv2d + sparse_layer = pq_conv( config, layer.in_channels, layer.out_channels, @@ -1405,43 +1571,45 @@ def add_pruning_to_model(module, config, prefix=""): quantize_input, quantize_output, ) - sparse_layer._weight.data = layer.weight.data - if layer.bias is not None: - sparse_layer._bias.data = layer.bias.data - sparse_layer = add_layer_specific_quantization_to_model(full_name, sparse_layer, config) - sparse_layer = disable_pruning_from_layers(full_name, sparse_layer, config) - setattr(module, name, sparse_layer) - elif layer.__class__ is nn.Conv1d: - sparse_layer = PQConv1d( + _replace_layer_with_pq_layer(module, name, full_name, layer, sparse_layer, config) + elif layer.__class__ is nn.MultiheadAttention: + if layer.bias_k is not None or layer.add_zero_attn: + raise ValueError(f"add_bias_kv/add_zero_attn are not supported by PQMultiheadAttention ({full_name})") + sparse_layer = PQMultiheadAttention( config, - layer.in_channels, - layer.out_channels, - layer.kernel_size, - layer.stride, - layer.padding, - layer.dilation, - layer.groups, - layer.bias is not None, - layer.padding_mode, - layer.weight.device, - layer.weight.dtype, - quantize_input, - quantize_output, + embed_dim=layer.embed_dim, + num_heads=layer.num_heads, + dropout=layer.dropout, + bias=layer.in_proj_bias is not None, + kdim=layer.kdim, + vdim=layer.vdim, + batch_first=layer.batch_first, + quantize_input=quantize_input, + quantize_output=quantize_output, ) - sparse_layer._weight.data = layer.weight.data - if layer.bias is not None: - sparse_layer._bias.data = layer.bias.data - sparse_layer = add_layer_specific_quantization_to_model(full_name, sparse_layer, config) - sparse_layer = disable_pruning_from_layers(full_name, sparse_layer, config) + if layer._qkv_same_embed_dim: + q_w, k_w, v_w = layer.in_proj_weight.chunk(3) + else: + q_w, k_w, v_w = layer.q_proj_weight, layer.k_proj_weight, layer.v_proj_weight + if layer.in_proj_bias is not None: + biases = (*layer.in_proj_bias.chunk(3), layer.out_proj.bias) + else: + biases = (None, None, None, None) + projs = (sparse_layer.q_proj, sparse_layer.k_proj, sparse_layer.v_proj, sparse_layer.out_proj) + for proj, weight, bias in zip(projs, (q_w, k_w, v_w, layer.out_proj.weight), biases): + proj._weight.data = weight.data.clone() + if bias is not None: + proj._bias.data = bias.data.clone() + sparse_layer = _add_layer_specific_quantization_to_model(full_name, sparse_layer, config) setattr(module, name, sparse_layer) else: - add_pruning_to_model(layer, config, full_name) + _add_pruning_to_model(layer, config, full_name) return module def apply_final_compression(module): for layer in module.modules(): - if isinstance(layer, (PQWeightBiasBase, PQBatchNorm2d, PQBatchNorm1d, Quantizer)): + if isinstance(layer, (PQWeightBiasBase, PQBatchNorm2d, PQBatchNorm1d, PQLayerNorm, Quantizer)): layer.apply_final_compression() return module @@ -1456,100 +1624,92 @@ def call_post_round_functions(model, rewind, rounds, r): post_round_functions(model) +def _update_pruning_mask(layer): + if layer.enable_pruning and hasattr(layer.pruning_layer, "update_mask"): + layer.pruning_layer.update_mask(layer._weight) + + def post_epoch_functions(model, epoch, total_epochs, **kwargs): for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): layer.pruning_layer.post_epoch_function(epoch, total_epochs, **kwargs) + _update_pruning_mask(layer) elif isinstance(layer, Quantizer): layer.post_epoch_function() def pre_epoch_functions(model, epoch, total_epochs): for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): layer.pruning_layer.pre_epoch_function(epoch, total_epochs) def post_round_functions(model): for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): layer.pruning_layer.post_round_function() def save_weights_functions(model): for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): - layer.save_weights() + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + layer._save_weights() def rewind_weights_functions(model): for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): - layer.rewind_weights() + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + layer._rewind_weights() def pre_finetune_functions(model): for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): layer.pruning_layer.pre_finetune_function() def post_pretrain_functions(model, config, train_loader=None, loss_function=None, input_shape=None): + for layer in model.modules(): + if isinstance(layer, PQ_MODULES): + # For FITCompress this must happen before the compression path search, + # so quantization is already enabled during it. + layer.post_pre_train_function() if config.fitcompress_parameters.enable_fitcompress: from pquant.core.torch.fit_compress import call_fitcompress # noqa: 811 - for layer in model.modules(): - if isinstance( - layer, (PQConv2d, PQConv1d, PQDense, PQActivation, PQBatchNorm2d, PQBatchNorm1d, PQAvgPoolBase, Quantizer) - ): - # Trigger it here to enable quantization before FITCompress - layer.post_pre_train_function() config, pruning_mask_importance_scores = call_fitcompress( config, model, train_loader, loss_function, input_shape=input_shape ) idx = 0 for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): - # layer.post_pre_train_function() - # set_data_quantization_bits(model) - layer.pruning_layer.mask.assign(pruning_mask_importance_scores[idx]) + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + with torch.no_grad(): + layer.pruning_layer.mask.data = pruning_mask_importance_scores[idx] layer.pruning_layer.pre_finetune_function() # So mask is not updated during training anymore idx += 1 return - else: - for layer in model.modules(): - if isinstance( - layer, (PQConv2d, PQConv1d, PQDense, PQActivation, PQBatchNorm2d, PQBatchNorm1d, PQAvgPoolBase, Quantizer) - ): - layer.post_pre_train_function() if config.pruning_parameters.pruning_method == "pdp" or ( config.pruning_parameters.pruning_method == "wanda" and config.pruning_parameters.calculate_pruning_budget ): - # pass - pdp_setup(model, config) + _pdp_setup(model, config) -def pdp_setup(model, config): +def _pdp_setup(model, config): """ Calculates a global sparsity threshold. Initializes target sparsity for each layer, which depends on how large percentage of weights in the layer is smaller than the global threshold """ - global_weights = None - for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): - if global_weights is None: - global_weights = layer._weight.flatten() - else: - global_weights = torch.concat((global_weights, layer._weight.flatten())) - + global_weights = torch.concat( + [layer._weight.flatten() for layer in model.modules() if isinstance(layer, LAYERS_WITH_PRUNING_LAYER)] + ) abs_global_weights = torch.abs(global_weights) global_weight_topk, _ = torch.topk(abs_global_weights, abs_global_weights.numel()) threshold = global_weight_topk[int((1 - config.pruning_parameters.sparsity) * global_weight_topk.numel())] global_weights_below_threshold = torch.where(abs_global_weights < threshold, 1, 0) idx = 0 for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): weight_size = layer._weight.numel() w = torch.sum(global_weights_below_threshold[idx : idx + weight_size]) layer.pruning_layer.init_r = w / weight_size @@ -1562,11 +1722,10 @@ def get_layer_keep_ratio(model): total_w = 0 remaining_weights = 0 for layer in model.modules(): - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): weight = layer.weight - total_w += ops.size(weight) - rem = ops.count_nonzero(weight) - remaining_weights += rem + total_w += weight.numel() + remaining_weights += torch.count_nonzero(weight) elif layer.__class__ in (nn.Conv2d, nn.Conv1d, nn.Linear): total_w += layer.weight.numel() remaining_weights += torch.count_nonzero(layer.weight) @@ -1575,150 +1734,112 @@ def get_layer_keep_ratio(model): return 0.0 -def is_training_stage(layer): - return False if layer.pruning_layer._is_finetuning or layer.pruning_layer._is_pretraining else True +def _is_training_stage(layer): + return not (layer.pruning_layer._is_finetuning or layer.pruning_layer._is_pretraining) def get_model_losses(model, losses): for layer in model.modules(): - loss = 0.0 - if isinstance(layer, (PQConv2d, PQConv1d, PQDense)): - if layer.enable_pruning and is_training_stage(layer) and not layer.use_fitcompress: - loss += layer.pruning_layer.calculate_additional_loss() + if isinstance(layer, LAYERS_WITH_PRUNING_LAYER): + if layer.enable_pruning and _is_training_stage(layer) and not layer.use_fitcompress: + losses += layer.pruning_layer.calculate_additional_loss() if layer.use_hgq: - loss += layer.hgq_loss() - losses += loss - elif isinstance(layer, (PQAvgPool1d, PQAvgPool2d, PQBatchNorm2d, PQBatchNorm1d, PQActivation)): + losses += layer.hgq_loss() + elif isinstance( + layer, + ( + PQAvgPool1d, + PQAvgPool2d, + PQBatchNorm2d, + PQBatchNorm1d, + PQLayerNorm, + PQActivation, + PQSoftmax, + PQMultiheadAttention, + ), + ): if layer.use_hgq: losses += layer.hgq_loss() return losses -def create_default_layer_quantization_pruning_config(model, config): - # subconfig = {"layer_specific": {}, "disable_pruning_for_layers": []} +def _create_default_layer_quantization_pruning_config(model, config): + quant_params = config.quantization_parameters + + def data_section(quantize, integer_bits, fractional_bits): + return { + "quantize": quantize, + "keep_negatives": quant_params.default_data_keep_negatives, + "integer_bits": integer_bits, + "fractional_bits": fractional_bits, + } + + def param_section(integer_bits, fractional_bits): + return { + "keep_negatives": quant_params.default_weight_keep_negatives, + "integer_bits": integer_bits, + "fractional_bits": fractional_bits, + } + for name, layer in model.named_modules(): if layer.__class__ in [nn.Linear, nn.Conv1d, nn.Conv2d]: - if layer.bias is None: - config.quantization_parameters.layer_specific[name] = { - "input": { - "keep_negatives": config.quantization_parameters.default_data_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7, - "quantize": config.quantization_parameters.quantize_input, - }, - "weight": { - "keep_negatives": config.quantization_parameters.default_weight_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7, - }, - "output": { - "keep_negatives": config.quantization_parameters.default_data_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7, - "quantize": config.quantization_parameters.quantize_output, - }, - } - else: - config.quantization_parameters.layer_specific[name] = { - "input": { - "keep_negatives": config.quantization_parameters.default_data_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7, - "quantize": config.quantization_parameters.quantize_input, - }, - "weight": { - "keep_negatives": config.quantization_parameters.default_weight_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7, - }, - "bias": { - "keep_negatives": config.quantization_parameters.default_weight_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7, - }, - "output": { - "keep_negatives": config.quantization_parameters.default_data_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7, - "quantize": config.quantization_parameters.quantize_output, - }, - } + layer_config = { + "input": data_section(quant_params.quantize_input, 0, 7), + "weight": param_section(0, 7), + "output": data_section(quant_params.quantize_output, 0, 7), + } + if layer.bias is not None: + layer_config["bias"] = param_section(0, 7) + quant_params.layer_specific[name] = layer_config config.pruning_parameters.disable_pruning_for_layers.append(name) elif layer.__class__ in [nn.Tanh, nn.ReLU, nn.AvgPool1d, nn.AvgPool2d, nn.AvgPool3d]: - config.quantization_parameters.layer_specific[name] = { - "input": { - "quantize": config.quantization_parameters.quantize_input, - "keep_negatives": config.quantization_parameters.default_data_keep_negatives, - "integer_bits": 0.0, - "fractional_bits": 7.0, - }, - "output": { - "quantize": config.quantization_parameters.quantize_output, - "keep_negatives": config.quantization_parameters.default_data_keep_negatives, - "integer_bits": 0.0, - "fractional_bits": 7.0, - }, + quant_params.layer_specific[name] = { + "input": data_section(quant_params.quantize_input, 0.0, 7.0), + "output": data_section(quant_params.quantize_output, 0.0, 7.0), } elif layer.__class__ in [nn.BatchNorm2d]: - config.quantization_parameters.layer_specific[name] = { - "input": { - "quantize": config.quantization_parameters.quantize_input, - "keep_negatives": config.quantization_parameters.default_data_keep_negatives, - "integer_bits": 0.0, - "fractional_bits": 7.0, - }, - "weight": { - "keep_negatives": config.quantization_parameters.default_weight_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7.0, - }, - "bias": { - "keep_negatives": config.quantization_parameters.default_weight_keep_negatives, - "integer_bits": 0, - "fractional_bits": 7.0, - }, + quant_params.layer_specific[name] = { + "input": data_section(quant_params.quantize_input, 0.0, 7.0), + "weight": param_section(0, 7.0), + "bias": param_section(0, 7.0), } return config def populate_config_with_all_layers(model, config): - return create_default_layer_quantization_pruning_config(model, config) + return _create_default_layer_quantization_pruning_config(model, config) -def remove_compression_layers(module, config): +def _remove_compression_layers(module, config): for name, layer in module.named_children(): if isinstance(layer, PQDense): - out_features = layer.out_features - in_features = layer.in_features - bias = True if layer.bias is not None else False - setattr(module, name, nn.Linear(in_features=in_features, out_features=out_features, bias=bias)) - getattr(module, name).weight.data.copy_(layer.weight) - if getattr(module, name).bias is not None: - getattr(module, name).bias.data.copy_(layer.bias) + new_layer = nn.Linear( + in_features=layer.in_features, out_features=layer.out_features, bias=layer.bias is not None + ) + new_layer.weight.data.copy_(layer.weight) + if new_layer.bias is not None: + new_layer.bias.data.copy_(layer.bias) + setattr(module, name, new_layer) elif isinstance(layer, (PQConv1d, PQConv2d)): - bias_values = layer.bias if layer.bias is not None else None - bias = True if bias_values is not None else False + bias_values = layer.bias conv = nn.Conv2d if isinstance(layer, PQConv2d) else nn.Conv1d - setattr( - module, - name, - conv( - layer.in_channels, - layer.out_channels, - layer.kernel_size, - layer.stride, - layer.padding, - layer.dilation, - layer.groups, - bias, - layer.padding_mode, - ), + new_layer = conv( + layer.in_channels, + layer.out_channels, + layer.kernel_size, + layer.stride, + layer.padding, + layer.dilation, + layer.groups, + bias_values is not None, + layer.padding_mode, ) - getattr(module, name).weight.data.copy_(layer.weight) - if getattr(module, name).bias is not None: - getattr(module, name).bias.data.copy_(bias_values.data) + new_layer.weight.data.copy_(layer.weight) + if new_layer.bias is not None: + new_layer.bias.data.copy_(bias_values.data) + setattr(module, name, new_layer) else: - remove_compression_layers(layer, config) + _remove_compression_layers(layer, config) return module @@ -1731,23 +1852,25 @@ def post_training_prune(model, config, calibration_data): model = add_compression_layers(model, config, inputs.shape) post_pretrain_functions(model, config) model(inputs) - return remove_compression_layers(model, config) + return _remove_compression_layers(model, config) def get_ebops(model, **kwargs): ebops = 0 for m in model.modules(): - if isinstance(m, (PQWeightBiasBase)): + if isinstance(m, PQWeightBiasBase): ebops += m.ebops(include_mask=m.enable_pruning) - elif isinstance(m, (PQAvgPoolBase, PQBatchNorm1d, PQBatchNorm2d, PQActivation)): + elif isinstance( + m, + ( + PQAvgPoolBase, + PQBatchNorm1d, + PQBatchNorm2d, + PQLayerNorm, + PQActivation, + PQSoftmax, + PQMultiheadAttention, + ), + ): ebops += m.ebops() return ebops - - -def load_torch_hgq_model(model, path_to_checkpoint): - model.load_state_dict(torch.load(path_to_checkpoint), strict=False) - for m in model.modules(): - if isinstance(m, Quantizer) and m.quantizer.built: - # Populate HGQ quantizer bit values from PQuantML quantizer - m.reload_from_local() - return model diff --git a/src/pquant/core/torch/optimizers.py b/src/pquant/core/torch/optimizers.py index a85022e..8914454 100644 --- a/src/pquant/core/torch/optimizers.py +++ b/src/pquant/core/torch/optimizers.py @@ -20,20 +20,20 @@ def step(self, closure=None): # Store the old params old_params = [] for group in self.param_groups: - old_params.append({param: param.data.clone() for param in group['params'] if param.grad is not None}) + old_params.append({param: param.data.clone() for param in group["params"] if param.grad is not None}) # Perform the standard AdamW step loss = super().step(closure) # Perform the pWD step for group, old_group in zip(self.param_groups, old_params): - lambda_p_group = group.get('lambda_p', self.lambda_p) # support prams groups + lambda_p_group = group.get("lambda_p", self.lambda_p) # support prams groups if lambda_p_group > 0: # Apply regularization only for lambda_p > 0 - for param in group['params']: + for param in group["params"]: if param.grad is None: continue # Use old parameters in the decay factor param_old = old_group[param] X = param_old.abs() ** (2 - self.p_norm) - update_term = X / (X + self.p_norm * group['lr'] * lambda_p_group) + update_term = X / (X + self.p_norm * group["lr"] * lambda_p_group) # pWD step param.data.mul_(update_term) return loss @@ -50,20 +50,20 @@ def step(self, closure=None): # Store the old params old_params = [] for group in self.param_groups: - old_params.append({param: param.data.clone() for param in group['params'] if param.grad is not None}) + old_params.append({param: param.data.clone() for param in group["params"] if param.grad is not None}) # Perform the standard SGD step loss = super().step(closure) # Perform the pWD step for group, old_group in zip(self.param_groups, old_params): - lambda_p_group = group.get('lambda_p', self.lambda_p) # support prams groups + lambda_p_group = group.get("lambda_p", self.lambda_p) # support prams groups if lambda_p_group > 0: # Apply regularization only for lambda_p > 0 - for param in group['params']: + for param in group["params"]: if param.grad is None: continue # Use old parameters in the decay factor param_old = old_group[param] X = param_old.abs() ** (2 - self.p_norm) - update_term = X / (X + self.p_norm * group['lr'] * lambda_p_group) + update_term = X / (X + self.p_norm * group["lr"] * lambda_p_group) # pWD step param.data.mul_(update_term) return loss diff --git a/src/pquant/core/torch/pruning_methods/__init__.py b/src/pquant/core/torch/pruning_methods/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/pquant/core/torch/pruning_methods/activation_pruning.py b/src/pquant/core/torch/pruning_methods/activation_pruning.py new file mode 100644 index 0000000..554a4b5 --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/activation_pruning.py @@ -0,0 +1,104 @@ +import torch +import torch.nn as nn + + +class ActivationPruning(nn.Module): + def __init__(self, config, layer_type, *args, **kwargs): + super().__init__() + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + self.act_type = "relu" + self.layer_type = layer_type + self._is_pretraining = True + self._is_finetuning = False + self.is_pretraining = True + self.is_finetuning = False + self.threshold = float(config.pruning_parameters.threshold) + self.t_start_collecting_batch = int(self.config.pruning_parameters.t_start_collecting_batch) + self.built = False + + def build(self, input_shape): + if self.built: + return + if self.layer_type in ("conv", "depthwise_conv"): + if len(input_shape) == 3: + shape = (input_shape[0], 1, 1) + else: + shape = (input_shape[0], 1, 1, 1) + else: + shape = (input_shape[0], 1) + self.shape = shape + n_channels = input_shape[0] + self.register_buffer("mask", torch.ones(shape)) + self.register_buffer("mask_placeholder", torch.ones(shape)) + self.register_buffer("activations", torch.zeros(n_channels)) + self.register_buffer("batches_collected", torch.zeros((), dtype=torch.int32)) + self.register_buffer("t", torch.zeros((), dtype=torch.int32)) + self.built = True + + @torch.no_grad() + def collect_output(self, output, training): + if not training: + return + if self._is_pretraining or self._is_finetuning: + return + if int(self.t.item()) < self.t_start_collecting_batch: + return + + t_delta = int(self.config.pruning_parameters.t_delta) + + gt_zero = (output > 0).to(output.dtype) + if self.layer_type == "linear": + per_channel = gt_zero.mean(dim=0) + else: + axes = (0,) + tuple(range(2, output.dim())) + per_channel = gt_zero.mean(dim=axes) + + self.activations.add_(per_channel.to(self.activations.dtype)) + self.batches_collected.add_(1) + + if int(self.batches_collected.item()) % t_delta == 0: + denom = max(int(self.batches_collected.item()), 1) + pct_active = self.activations / denom + new_mask = (pct_active > self.threshold).to(self.mask_placeholder.dtype).reshape(self.shape) + self.mask_placeholder.copy_(new_mask) + self.activations.zero_() + self.batches_collected.zero_() + self.t.zero_() + + def forward(self, weight): + if self._is_pretraining: + return weight + return self.mask.to(weight.dtype) * weight + + def get_hard_mask(self, weight=None): + return self.mask + + def post_pre_train_function(self): + self._is_pretraining = False + self.is_pretraining = False + + def pre_epoch_function(self, epoch, total_epochs, **kwargs): + pass + + def post_round_function(self): + pass + + def pre_finetune_function(self): + self._is_finetuning = True + self.is_finetuning = True + + def calculate_additional_loss(self): + return 0.0 + + def get_layer_sparsity(self, weight): + pass + + @torch.no_grad() + def post_epoch_function(self, epoch, total_epochs, **kwargs): + if not self._is_pretraining: + self.t.add_(1) + self.mask.copy_(self.mask_placeholder) diff --git a/src/pquant/core/torch/pruning_methods/autosparse.py b/src/pquant/core/torch/pruning_methods/autosparse.py new file mode 100644 index 0000000..81b723b --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/autosparse.py @@ -0,0 +1,143 @@ +import math + +import numpy as np +import torch +import torch.nn as nn + +_PI = math.pi +_L0 = -6.0 +_L1 = 6.0 + + +def cosine_decay(i, T): + return (1 + math.cos(_PI * i / T)) / 2 + + +def sigmoid_decay(i, T): + x = _L0 + (_L1 - _L0) * i / T + return 1.0 - 1.0 / (1.0 + math.exp(-x)) + + +def cosine_sigmoid_decay(i, T): + return max(cosine_decay(i, T), sigmoid_decay(i, T)) + + +def get_threshold_size(config, weight_shape): + if config.pruning_parameters.threshold_type == "layerwise": + return (1, 1) + elif config.pruning_parameters.threshold_type == "channelwise": + return (weight_shape[0], 1) + elif config.pruning_parameters.threshold_type == "weightwise": + return (weight_shape[0], int(np.prod(weight_shape[1:]))) + + +class _AutoSparsePrune(torch.autograd.Function): + @staticmethod + def forward(ctx, x, alpha, backward_sparsity_flag, backward_sparsity): + mask = torch.relu(x) + kth_value = None + if backward_sparsity_flag: + flat = x.reshape(-1) + k = max(int(flat.numel() * backward_sparsity), 1) + topk_vals, _ = torch.topk(flat, k) + kth_value = topk_vals[-1] + ctx.save_for_backward(x, alpha, kth_value if kth_value is not None else torch.zeros((), device=x.device)) + ctx.backward_sparsity_flag = backward_sparsity_flag + return mask + + @staticmethod + def backward(ctx, upstream): + x, alpha, kth_value = ctx.saved_tensors + grads = torch.where(x <= 0, alpha.to(x.dtype), torch.ones_like(x)) + if ctx.backward_sparsity_flag: + grads = torch.where(x < kth_value, torch.zeros_like(grads), grads) + return grads * upstream, None, None, None + + +def autosparse_prune(x, alpha, backward_sparsity_flag, backward_sparsity): + return _AutoSparsePrune.apply(x, alpha, backward_sparsity_flag, backward_sparsity) + + +class AutoSparse(nn.Module): + def __init__(self, config, layer_type, *args, **kwargs): + super().__init__() + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + self.layer_type = layer_type + self._alpha_init = float(config.pruning_parameters.alpha) + self._backward_sparsity_flag = bool(config.pruning_parameters.backward_sparsity) + self._backward_sparsity = 0.5 + self._is_pretraining = True + self._is_finetuning = False + self.is_pretraining = True + self.is_finetuning = False + self.built = False + + def build(self, input_shape): + if self.built: + return + threshold_size = get_threshold_size(self.config, input_shape) + self.threshold = nn.Parameter(torch.full(threshold_size, float(self.config.pruning_parameters.threshold_init))) + self.register_buffer("mask", torch.ones(tuple(input_shape))) + self.register_buffer("alpha", torch.tensor(self._alpha_init)) + self.built = True + + def _g(self, x): + return torch.sigmoid(x) + + def forward(self, weight): + if self._is_pretraining: + return weight + if self._is_finetuning: + return self.mask.to(weight.dtype) * weight + + weight_reshaped = weight.reshape(weight.shape[0], -1) + w_t = weight_reshaped.abs() - torch.sigmoid(self.threshold) + + new_binary_mask = (w_t > 0).to(weight.dtype).reshape(weight.shape) + with torch.no_grad(): + self.mask.copy_(new_binary_mask) + + sparse = torch.sign(weight) * autosparse_prune( + w_t, self.alpha, self._backward_sparsity_flag, self._backward_sparsity + ).reshape(weight.shape) + return sparse + + def get_hard_mask(self, weight=None): + return self.mask + + def get_mask(self, weight): + weight_reshaped = weight.reshape(weight.shape[0], -1) + w_t = weight_reshaped.abs() - torch.sigmoid(self.threshold) + return (w_t > 0).to(weight.dtype).reshape(weight.shape) + + def get_layer_sparsity(self, weight): + m = self.get_mask(weight) + return m.count_nonzero() / m.numel() + + def pre_epoch_function(self, epoch, total_epochs): + pass + + def calculate_additional_loss(self): + return 0.0 + + def pre_finetune_function(self): + self._is_finetuning = True + self.is_finetuning = True + + def post_round_function(self): + pass + + def post_pre_train_function(self): + self._is_pretraining = False + self.is_pretraining = False + + @torch.no_grad() + def post_epoch_function(self, epoch, total_epochs): + decay = cosine_sigmoid_decay(epoch, total_epochs) + self.alpha.fill_(self._alpha_init * decay) + if epoch >= self.config.pruning_parameters.alpha_reset_epoch: + self.alpha.zero_() diff --git a/src/pquant/core/torch/pruning_methods/constraint_functions.py b/src/pquant/core/torch/pruning_methods/constraint_functions.py new file mode 100644 index 0000000..565c960 --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/constraint_functions.py @@ -0,0 +1,99 @@ +import abc + +import torch +import torch.nn as nn + + +class _FlipGradient(torch.autograd.Function): + @staticmethod + def forward(ctx, x, scale): + ctx.scale = float(scale) + return x + + @staticmethod + def backward(ctx, upstream): + return upstream * ctx.scale, None + + +def flip_gradient(x, scale=-1.0): + return _FlipGradient.apply(x, scale) + + +class Constraint(nn.Module): + def __init__(self, lmbda_init=1.0, scale=1.0, damping=1.0, use_grad=True, lr=0.0, **kwargs): + super().__init__() + self.use_grad_ = bool(use_grad) + self.lr_ = float(lr) + self.register_buffer("scale", torch.tensor(float(scale))) + self.register_buffer("damping", torch.tensor(float(damping))) + if self.use_grad_: + self.lmbda = nn.Parameter(torch.tensor(float(lmbda_init))) + else: + self.register_buffer("lmbda", torch.tensor(float(lmbda_init))) + self.register_buffer("prev_infs", torch.tensor(0.0)) + + def build(self, input_shape): + # Exists for API parity with the keras version — no-op in torch. + pass + + def forward(self, weight, training=None): + raw_infeasibility = self.get_infeasibility(weight) + infeasibility = self.pipe_infeasibility(raw_infeasibility) + + if self.use_grad_: + ascent_lmbda = flip_gradient(self.lmbda) + else: + lmbda_step = self.lr_ * self.scale * self.prev_infs + ascent_lmbda = self.lmbda + lmbda_step + if training: + with torch.no_grad(): + self.lmbda.add_(lmbda_step) + self.prev_infs.copy_(infeasibility.detach()) + + l_term = ascent_lmbda * infeasibility + damp_term = self.damping * infeasibility.square() / 2 + return self.scale * (l_term + damp_term) + + @abc.abstractmethod + def get_infeasibility(self, weight): + raise NotImplementedError + + def pipe_infeasibility(self, infeasibility): + return infeasibility + + @torch.no_grad() + def turn_off(self): + if not self.use_grad_: + self.lr_ = 0.0 + self.scale.zero_() + self.lmbda.data.zero_() if isinstance(self.lmbda, nn.Parameter) else self.lmbda.zero_() + + +class EqualityConstraint(Constraint): + def __init__(self, metric_fn, target_value=0.0, **kwargs): + super().__init__(**kwargs) + self.metric_fn = metric_fn + self.target_value = float(target_value) + + def get_infeasibility(self, weight): + return (self.metric_fn(weight) - self.target_value).abs() + + +class LessThanOrEqualConstraint(Constraint): + def __init__(self, metric_fn, target_value=0.0, **kwargs): + super().__init__(**kwargs) + self.metric_fn = metric_fn + self.target_value = float(target_value) + + def get_infeasibility(self, weight): + return torch.clamp(self.metric_fn(weight) - self.target_value, min=0.0) + + +class GreaterThanOrEqualConstraint(Constraint): + def __init__(self, metric_fn, target_value=0.0, **kwargs): + super().__init__(**kwargs) + self.metric_fn = metric_fn + self.target_value = float(target_value) + + def get_infeasibility(self, weight): + return torch.clamp(self.target_value - self.metric_fn(weight), min=0.0) diff --git a/src/pquant/core/torch/pruning_methods/cs.py b/src/pquant/core/torch/pruning_methods/cs.py new file mode 100644 index 0000000..15f9f53 --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/cs.py @@ -0,0 +1,81 @@ +import torch +import torch.nn as nn + + +class ContinuousSparsification(nn.Module): + def __init__(self, config, layer_type, *args, **kwargs): + super().__init__() + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + self.final_temp = float(config.pruning_parameters.final_temp) + self._is_finetuning = False + self._is_pretraining = True + self.is_pretraining = True + self.is_finetuning = False + self.layer_type = layer_type + self.built = False + + def build(self, input_shape): + if self.built: + return + init_val = float(self.config.pruning_parameters.threshold_init) + s_init = torch.full(tuple(input_shape), init_val) + self.s = nn.Parameter(s_init.clone()) + self.register_buffer("s_init", s_init.clone()) + self.register_buffer("scaling", 1.0 / torch.sigmoid(s_init)) + self.register_buffer("beta", torch.tensor(1.0)) + self.register_buffer("mask", torch.ones(tuple(input_shape))) + self.built = True + + def forward(self, weight): + if self._is_pretraining or self._is_finetuning: + return self.mask.to(weight.dtype) * weight + new_mask = self.get_mask() + with torch.no_grad(): + self.mask.copy_(new_mask.detach()) + return new_mask * weight + + def pre_finetune_function(self): + self._is_finetuning = True + self.is_finetuning = True + with torch.no_grad(): + self.mask.copy_(self.get_hard_mask().to(self.mask.dtype)) + + def get_mask(self): + return torch.sigmoid(self.beta * self.s) * self.scaling + + def post_pre_train_function(self): + self._is_pretraining = False + self.is_pretraining = False + + def pre_epoch_function(self, epoch, total_epochs): + pass + + @torch.no_grad() + def post_epoch_function(self, epoch, total_epochs): + if total_epochs <= 1: + self.beta.mul_(self.final_temp) + else: + self.beta.mul_(self.final_temp ** (1 / (total_epochs - 1))) + + def get_hard_mask(self, weight=None): + if self.config.pruning_parameters.enable_pruning: + return (self.s > 0).to(self.s.dtype) + return torch.tensor(1.0, device=self.s.device, dtype=self.s.dtype) + + @torch.no_grad() + def post_round_function(self): + new_s = torch.minimum(self.beta * self.s, self.s_init) + self.s.data.copy_(new_s) + self.beta.fill_(1.0) + + def calculate_additional_loss(self): + if self._is_pretraining or self._is_finetuning: + return torch.zeros((), dtype=self.s.dtype, device=self.s.device) + return self.config.pruning_parameters.threshold_decay * torch.norm(self.get_mask().reshape(-1), p=1) + + def get_layer_sparsity(self, weight): + return self.get_hard_mask().sum() / weight.numel() diff --git a/src/pquant/core/torch/pruning_methods/dst.py b/src/pquant/core/torch/pruning_methods/dst.py new file mode 100644 index 0000000..4cb1c5e --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/dst.py @@ -0,0 +1,105 @@ +import numpy as np +import torch +import torch.nn as nn + + +def get_threshold_size(config, weight_shape): + if config.pruning_parameters.threshold_type == "layerwise": + return (1, 1) + elif config.pruning_parameters.threshold_type == "channelwise": + return (weight_shape[0], 1) + elif config.pruning_parameters.threshold_type == "weightwise": + return (weight_shape[0], int(np.prod(weight_shape[1:]))) + + +class _BinaryStep(torch.autograd.Function): + @staticmethod + def forward(ctx, weight): + ctx.save_for_backward(weight) + return (weight > 0).to(weight.dtype) + + @staticmethod + def backward(ctx, upstream): + (weight,) = ctx.saved_tensors + abs_w = weight.abs() + idx_lt04 = torch.where(abs_w <= 0.4, 2 - 4 * abs_w, torch.zeros_like(weight)) + idx_04to1 = torch.where((abs_w > 0.4) & (abs_w <= 1.0), torch.full_like(weight, 0.4), torch.zeros_like(weight)) + grads = idx_lt04 + idx_04to1 + return grads * upstream + + +def binary_step(weight): + return _BinaryStep.apply(weight) + + +class DST(nn.Module): + def __init__(self, config, layer_type, *args, **kwargs): + super().__init__() + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + self.layer_type = layer_type + self._is_pretraining = True + self._is_finetuning = False + self.is_pretraining = True + self.is_finetuning = False + self.built = False + + def build(self, input_shape): + if self.built: + return + threshold_size = get_threshold_size(self.config, input_shape) + self.threshold = nn.Parameter(torch.zeros(threshold_size)) + self.register_buffer("mask", torch.ones(tuple(input_shape))) + self.built = True + + def forward(self, weight): + if self._is_pretraining or self._is_finetuning: + return weight * self.mask.to(weight.dtype) + + mask = self.get_mask(weight) + ratio = 1.0 - mask.sum() / mask.numel() + if float(ratio.detach()) >= self.config.pruning_parameters.max_pruning_pct: + with torch.no_grad(): + self.threshold.data.zero_() + mask = self.get_mask(weight) + with torch.no_grad(): + self.mask.copy_(mask.detach()) + return weight * mask + + def get_hard_mask(self, weight=None): + return self.mask + + def get_mask(self, weight): + weight_orig_shape = weight.shape + weights_reshaped = weight.reshape(weight.shape[0], -1) + pre_binarystep = weights_reshaped.abs() - self.threshold + mask = binary_step(pre_binarystep) + return mask.reshape(weight_orig_shape) + + def pre_epoch_function(self, epoch, total_epochs): + pass + + def get_layer_sparsity(self, weight): + return self.get_mask(weight).sum() / weight.numel() + + def calculate_additional_loss(self): + if self._is_pretraining or self._is_finetuning: + return torch.zeros((), dtype=self.threshold.dtype, device=self.threshold.device) + return self.config.pruning_parameters.alpha * torch.sum(torch.exp(-self.threshold)) + + def pre_finetune_function(self): + self._is_finetuning = True + self.is_finetuning = True + + def post_epoch_function(self, epoch, total_epochs): + pass + + def post_pre_train_function(self): + self._is_pretraining = False + self.is_pretraining = False + + def post_round_function(self): + pass diff --git a/src/pquant/core/torch/pruning_methods/fitcompress.py b/src/pquant/core/torch/pruning_methods/fitcompress.py new file mode 100644 index 0000000..2c119fe --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/fitcompress.py @@ -0,0 +1,45 @@ +import torch +import torch.nn as nn + + +class FITCompress(nn.Module): + def __init__(self, config, *args, **kwargs): + super().__init__() + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + self._is_pretraining = True + self._is_finetuning = False + self.built = False + + def build(self, input_shape): + if self.built: + return + self.register_buffer("mask", torch.ones(tuple(input_shape))) + self.built = True + + def forward(self, weight): + return self.mask.to(weight.dtype) * weight + + def get_hard_mask(self, weight=None): + return self.mask + + def pre_epoch_function(self, epoch, total_epochs, **kwargs): + pass + + def calculate_additional_loss(self): + return 0.0 + + def pre_finetune_function(self): + self._is_finetuning = True + + def post_round_function(self): + pass + + def post_pre_train_function(self): + self._is_pretraining = False + + def post_epoch_function(self, epoch, total_epochs, **kwargs): + pass diff --git a/src/pquant/core/torch/pruning_methods/mdmm.py b/src/pquant/core/torch/pruning_methods/mdmm.py new file mode 100644 index 0000000..8afe615 --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/mdmm.py @@ -0,0 +1,141 @@ +import inspect + +import torch +import torch.nn as nn + +from pquant.core.torch.pruning_methods.constraint_functions import ( + EqualityConstraint, + GreaterThanOrEqualConstraint, + LessThanOrEqualConstraint, +) +from pquant.core.torch.pruning_methods.metric_functions import ( + StructuredSparsityMetric, + UnstructuredSparsityMetric, +) + +_METRIC_REGISTRY = { + "UnstructuredSparsity": UnstructuredSparsityMetric, + "StructuredSparsity": StructuredSparsityMetric, +} + +_CONSTRAINT_REGISTRY = { + "Equality": EqualityConstraint, + "LessThanOrEqual": LessThanOrEqualConstraint, + "GreaterThanOrEqual": GreaterThanOrEqualConstraint, +} + + +class MDMM(nn.Module): + def __init__(self, config, layer_type, *args, **kwargs): + super().__init__() + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + self.layer_type = layer_type + self.constraint_layer = None + self._is_finetuning = False + self._is_pretraining = True + self.is_pretraining = True + self.is_finetuning = False + self._last_penalty = None + self.built = False + + def build(self, input_shape): + if self.built: + return + pruning_parameters = self.config.pruning_parameters + metric_type = pruning_parameters.metric_type + constraint_type = pruning_parameters.constraint_type + target_value = pruning_parameters.target_value + target_sparsity = pruning_parameters.target_sparsity + l0_mode = pruning_parameters.l0_mode + scale_mode = pruning_parameters.scale_mode + + candidate_kwargs = { + "epsilon": pruning_parameters.epsilon, + "target_sparsity": target_sparsity, + "l0_mode": l0_mode, + "scale_mode": scale_mode, + "rf": pruning_parameters.rf, + } + + metric_cls = _METRIC_REGISTRY.get(metric_type) + if metric_cls is None: + raise ValueError(f"Unknown metric_type: {metric_type}") + sig = inspect.signature(getattr(metric_cls, "__init__", metric_cls)) + metric_kwargs = {k: v for k, v in candidate_kwargs.items() if v is not None and k in sig.parameters} + metric_fn = metric_cls(**metric_kwargs) + + common_args = { + "metric_fn": metric_fn, + "target_value": target_value, + "scale": pruning_parameters.scale, + "damping": pruning_parameters.damping, + "use_grad": pruning_parameters.use_grad, + "lr": pruning_parameters.constraint_lr, + } + + constraint_type_cls = _CONSTRAINT_REGISTRY.get(constraint_type) + if constraint_type_cls is None: + raise ValueError(f"Unknown constraint_type: {constraint_type}") + self.constraint_layer = constraint_type_cls(**common_args) + + self.register_buffer("mask", torch.ones(tuple(input_shape))) + self.built = True + + def forward(self, weight): + epsilon = self.config.pruning_parameters.epsilon + hard_mask = (weight.abs() > epsilon).to(weight.dtype) + not_active = self._is_pretraining or self._is_finetuning + + if not not_active: + with torch.no_grad(): + self.mask.copy_(hard_mask.detach()) + + penalty = self.constraint_layer(weight, training=self.training).sum() + + if not_active: + self._last_penalty = torch.zeros((), device=weight.device, dtype=weight.dtype) + else: + self._last_penalty = penalty + + if self._is_finetuning: + return weight * hard_mask + return weight + + def get_hard_mask(self, weight=None): + if weight is None: + return self.mask + epsilon = self.config.pruning_parameters.epsilon + return (weight.abs() > epsilon).to(weight.dtype) + + def get_layer_sparsity(self, weight): + return self.get_hard_mask(weight).sum() / weight.numel() + + def calculate_additional_loss(self): + if self._last_penalty is None: + return 0.0 + return self._last_penalty + + def pre_epoch_function(self, epoch, total_epochs): + pass + + def pre_finetune_function(self): + self._is_finetuning = True + self.is_finetuning = True + if hasattr(self.constraint_layer, "module"): + self.constraint_layer.module.turn_off() + else: + self.constraint_layer.turn_off() + + def post_epoch_function(self, epoch, total_epochs): + pass + + def post_pre_train_function(self): + self._is_pretraining = False + self.is_pretraining = False + + def post_round_function(self): + pass diff --git a/src/pquant/core/torch/pruning_methods/metric_functions.py b/src/pquant/core/torch/pruning_methods/metric_functions.py new file mode 100644 index 0000000..054f338 --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/metric_functions.py @@ -0,0 +1,57 @@ +import torch + + +class UnstructuredSparsityMetric: + """L0-L1 based metric — torch port of the keras version.""" + + def __init__(self, l0_mode="coarse", scale_mode="mean", epsilon=1e-3, target_sparsity=0.8, alpha=100.0): + assert l0_mode in ["coarse", "smooth"], "Mode must be 'coarse' or 'smooth'" + assert scale_mode in ["sum", "mean"], "Scale mode must be 'sum' or 'mean'" + assert 0 <= target_sparsity <= 1, "target_sparsity must be between 0 and 1" + self.l0_mode = l0_mode + self.scale_mode = scale_mode + self.target_sparsity = float(target_sparsity) + self.epsilon = float(epsilon) + self.alpha = float(alpha) + self.l0_fn = self._coarse_l0 if l0_mode == "coarse" else self._smooth_l0 + self._scaling = self._mean_scaling if scale_mode == "mean" else self._sum_scaling + + def _sum_scaling(self, fn_value, num): + return fn_value + + def _mean_scaling(self, fn_value, num): + return fn_value / num + + def _coarse_l0(self, weight_vector): + return (weight_vector.abs() <= self.epsilon).to(torch.float32).mean() + + def _smooth_l0(self, weight_vector): + return torch.exp(-self.alpha * weight_vector.square()).mean() + + def __call__(self, weight): + num_weights = torch.tensor(float(weight.numel()), dtype=weight.dtype, device=weight.device) + flat = weight.reshape(-1) + l0_term = self.l0_fn(flat) + l1_term = flat.abs().sum() + factor = (self.target_sparsity**2) - l0_term.square() + fn_value = factor * l1_term + return self._scaling(fn_value, num_weights) + + +class StructuredSparsityMetric: + def __init__(self, rf=1, epsilon=1e-3): + self.rf = int(rf) + self.epsilon = float(epsilon) + + def __call__(self, weight): + w_reshaped = weight.reshape(weight.shape[0], -1) + num_weights = w_reshaped.shape[1] + padding = (self.rf - num_weights % self.rf) % self.rf + if padding: + w_padded = torch.nn.functional.pad(w_reshaped, (0, padding)) + else: + w_padded = w_reshaped + groups = w_padded.reshape(w_padded.shape[0], -1, self.rf) + group_norms = torch.sqrt((groups.square()).sum(dim=-1)) + zero_groups = (group_norms <= self.epsilon).to(torch.float32) + return zero_groups.sum() / float(group_norms.numel()) diff --git a/src/pquant/core/torch/pruning_methods/pdp.py b/src/pquant/core/torch/pruning_methods/pdp.py new file mode 100644 index 0000000..3fa097a --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/pdp.py @@ -0,0 +1,144 @@ +import math + +import torch +import torch.nn as nn + + +class PDP(nn.Module): + def __init__(self, config, layer_type, *args, **kwargs): + super().__init__() + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self._init_r = float(config.pruning_parameters.sparsity) + self._epsilon = float(config.pruning_parameters.epsilon) + self.temp = float(config.pruning_parameters.temperature) + self.config = config + self.layer_type = layer_type + self._is_pretraining = True + self._is_finetuning = False + self.is_pretraining = True + self.is_finetuning = False + self.built = False + + # Wanda/PDP setup code externally assigns to .init_r / .sparsity — keep as properties so + # assignments propagate to the stored Python float used by pre_epoch_function. + @property + def init_r(self): + return self._init_r + + @init_r.setter + def init_r(self, value): + self._init_r = float(value.detach().item()) if torch.is_tensor(value) else float(value) + + def build(self, input_shape): + if self.built: + return + structured = self.config.pruning_parameters.structured_pruning + if structured: + if self.layer_type == "linear": + mask_shape = (input_shape[0], 1) + elif len(input_shape) == 3: + mask_shape = (input_shape[0], 1, 1) + else: + mask_shape = (input_shape[0], 1, 1, 1) + else: + mask_shape = tuple(input_shape) + + self.softmax_shape = list(input_shape) + [1] + self._mask_numel = math.prod(mask_shape) + self.flat_weight_size = float(self._mask_numel) + + self.register_buffer("mask", torch.ones(mask_shape)) + self.register_buffer("r", torch.tensor(self._init_r)) + + if structured: + self._compute_mask = self._mask_structured_channel if self.layer_type == "conv" else self._mask_structured_linear + else: + self._compute_mask = self._mask_unstructured + self.built = True + + def post_pre_train_function(self): + self._is_pretraining = False + self.is_pretraining = False + + @torch.no_grad() + def pre_epoch_function(self, epoch, _): + if hasattr(self, "r"): + val = min(1.0, self._epsilon * (epoch + 1)) * self._init_r + self.r.fill_(val) + + def post_round_function(self): + pass + + def pre_finetune_function(self): + self._is_finetuning = True + self.is_finetuning = True + with torch.no_grad(): + self.mask.copy_((self.mask >= 0.5).to(self.mask.dtype)) + + def post_epoch_function(self, epoch, total_epochs): + pass + + def _mask_unstructured(self, weight): + weight_reshaped = weight.reshape(self.softmax_shape) + abs_flat = weight.abs().reshape(-1) + all_vals, _ = torch.topk(abs_flat, self._mask_numel) + ind = int((1 - float(self.r.item())) * self.flat_weight_size) - 1 + lim = max(0, min(ind, int(self.flat_weight_size) - 2)) + Wh, Wt = all_vals[lim], all_vals[lim + 1] + t = torch.ones_like(weight_reshaped) * (0.5 * (Wh + Wt)) + soft_input = torch.cat((t**2, weight_reshaped**2), dim=-1) / self.temp + mw = torch.softmax(soft_input, dim=-1)[..., 1] + return mw.reshape(weight.shape) + + def _mask_structured_linear(self, weight): + norm = torch.norm(weight, dim=1, p=2, keepdim=True) + norm_flat = norm.reshape(-1) + W_all, _ = torch.topk(norm_flat, self._mask_numel) + ind = int((1 - float(self.r.item())) * self.flat_weight_size) - 1 + lim = max(0, min(ind, self._mask_numel - 2)) + Wh, Wt = W_all[lim], W_all[lim + 1] + t = torch.ones_like(norm) * 0.5 * (Wh + Wt) + soft_input = torch.cat((t**2, norm**2), dim=1) / self.temp + mw = torch.softmax(soft_input, dim=1)[..., 1] + return mw.unsqueeze(-1) + + def _mask_structured_channel(self, weight): + weight_reshaped = weight.reshape(weight.shape[0], -1) + norm = torch.norm(weight_reshaped, dim=1, p=2) + norm_flat = norm.reshape(-1) + W_all, _ = torch.topk(norm_flat, self._mask_numel) + ind = int((1 - float(self.r.item())) * self.flat_weight_size) - 1 + lim = max(0, min(ind, self._mask_numel - 2)) + Wh, Wt = W_all[lim], W_all[lim + 1] + norm = norm.unsqueeze(-1) + t = torch.ones_like(norm) * 0.5 * (Wh + Wt) + soft_input = torch.cat((t**2, norm**2), dim=-1) / self.temp + mw = torch.softmax(soft_input, dim=-1)[..., 1] + while mw.dim() < weight.dim(): + mw = mw.unsqueeze(-1) + return mw + + def forward(self, weight): + if self._is_pretraining or self._is_finetuning: + return self.mask.to(weight.dtype) * weight + return self._compute_mask(weight) * weight + + def update_mask(self, weight): + """Update stored mask from current weights. Called once per epoch from post_epoch_functions.""" + if not self._is_pretraining and not self._is_finetuning: + with torch.no_grad(): + self.mask.copy_(self._compute_mask(weight)) + + def get_hard_mask(self, weight=None): + return (self.mask >= 0.5).to(self.mask.dtype) + + def calculate_additional_loss(self): + return 0.0 + + def get_layer_sparsity(self, weight): + hard_mask = (self.mask >= 0.5).to(self.mask.dtype) + masked_weight = hard_mask * weight + return masked_weight.count_nonzero() / masked_weight.numel() diff --git a/src/pquant/core/torch/pruning_methods/wanda.py b/src/pquant/core/torch/pruning_methods/wanda.py new file mode 100644 index 0000000..331c65c --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/wanda.py @@ -0,0 +1,160 @@ +import torch +import torch.nn as nn + + +class Wanda(nn.Module): + def __init__(self, config, layer_type, *args, **kwargs): + super().__init__() + if isinstance(config, dict): + from pquant.core.hyperparameter_optimization import PQConfig + + config = PQConfig.load_from_config(config) + self.config = config + self.act_type = "relu" + self.layer_type = layer_type + self._is_pretraining = True + self._is_finetuning = False + self.is_pretraining = True + self.is_finetuning = False + self._sparsity = float(self.config.pruning_parameters.sparsity) + self.N = self.config.pruning_parameters.N + self.M = self.config.pruning_parameters.M + self.t_start_collecting_batch = int(self.config.pruning_parameters.t_start_collecting_batch) + self.built = False + + # Expose sparsity as a tensor so it supports torch tensor API (.cpu(), etc.) + # while keeping the Python float for cheap internal arithmetic. + @property + def sparsity(self): + return torch.tensor(self._sparsity) + + @sparsity.setter + def sparsity(self, v): + self._sparsity = float(v.detach().item()) if torch.is_tensor(v) else float(v) + + def build(self, input_shape): + if self.built: + return + n_in = input_shape[0] if self.layer_type == "depthwise_conv" else input_shape[1] + self.register_buffer("mask", torch.ones(tuple(input_shape))) + self.register_buffer("inputs_sq_sum", torch.zeros(n_in)) + self.register_buffer("batches_collected", torch.zeros((), dtype=torch.int32)) + self.register_buffer("t", torch.zeros((), dtype=torch.int32)) + self.register_buffer("done", torch.zeros((), dtype=torch.bool)) + self.built = True + + @torch.no_grad() + def collect_input(self, x, weight, training): + if not training: + return + if self._is_pretraining or self._is_finetuning: + return + if bool(self.done.item()): + return + if int(self.t.item()) < self.t_start_collecting_batch: + return + + t_delta = int(self.config.pruning_parameters.t_delta) + + if self.layer_type == "linear": + per_batch_sq = (x * x).sum(dim=0) + else: + axes = (0,) + tuple(range(2, x.dim())) + per_batch_sq = (x * x).sum(dim=axes) + + self.inputs_sq_sum.add_(per_batch_sq.to(self.inputs_sq_sum.dtype)) + self.batches_collected.add_(1) + + if int(self.batches_collected.item()) == t_delta: + norm = torch.sqrt(self.inputs_sq_sum) + new_mask = self._compute_prune_mask(norm, weight) + self.mask.copy_(new_mask.to(self.mask.dtype)) + self.done.fill_(True) + self.inputs_sq_sum.zero_() + self.batches_collected.zero_() + + def _compute_prune_mask(self, norm, weight): + if self.layer_type == "linear": + return self._handle_linear(norm, weight) + if self.layer_type == "depthwise_conv": + return self._handle_depthwise_conv(norm, weight) + return self._handle_conv(norm, weight) + + def _handle_linear(self, norm, weight): + metric = weight.abs() * norm + if self.N is not None and self.M is not None: + metric_reshaped = metric.reshape(-1, self.M) + weight_reshaped = weight.reshape(-1, self.M) + mask = self.get_mask(weight_reshaped, metric_reshaped, sparsity=self.N / self.M) + return mask.reshape(weight.shape) + metric_reshaped = metric.reshape(1, -1) + weight_reshaped = weight.reshape(1, -1) + mask = self.get_mask(weight_reshaped, metric_reshaped, sparsity=self._sparsity) + return mask.reshape(weight.shape) + + def _handle_conv(self, norm, weight): + if weight.dim() == 3: + norm_reshaped = norm.reshape(1, norm.shape[0], 1) + else: + norm_reshaped = norm.reshape(1, norm.shape[0], 1, 1) + metric = weight.abs() * norm_reshaped + if self.N is not None and self.M is not None: + metric_reshaped = metric.reshape(-1, self.M) + weight_reshaped = weight.reshape(-1, self.M) + mask = self.get_mask(weight_reshaped, metric_reshaped, sparsity=self.N / self.M) + return mask.reshape(weight.shape) + metric_reshaped = metric.reshape(metric.shape[0], -1) + weight_reshaped = weight.reshape(weight.shape[0], -1) + mask = self.get_mask(weight_reshaped, metric_reshaped, sparsity=self._sparsity) + return mask.reshape(weight.shape) + + def _handle_depthwise_conv(self, norm, weight): + norm_reshaped = norm.reshape(norm.shape[0], 1, 1, 1) + metric = weight.abs() * norm_reshaped + metric_reshaped = metric.reshape(metric.shape[0], -1) + weight_reshaped = weight.reshape(weight.shape[0], -1) + mask = self.get_mask(weight_reshaped, metric_reshaped, sparsity=self._sparsity) + return mask.reshape(weight.shape) + + def get_mask(self, weight, metric, sparsity): + d0, d1 = metric.shape + keep_idxs = ( + torch.argsort(metric, dim=1, stable=True)[:, int(d1 * sparsity) :] + + torch.arange(d0, device=metric.device)[:, None] * d1 + ) + keep_idxs = keep_idxs.flatten() + kept_values = torch.zeros(weight.numel(), dtype=weight.dtype, device=weight.device) + kept_values[keep_idxs] = weight.flatten()[keep_idxs] + kept_values = kept_values.reshape(weight.shape) + return (kept_values != 0).to(weight.dtype) + + def forward(self, weight): + return self.mask.to(weight.dtype) * weight + + def get_hard_mask(self, weight=None): + return self.mask + + def post_pre_train_function(self): + self._is_pretraining = False + self.is_pretraining = False + + def pre_epoch_function(self, epoch, total_epochs, **kwargs): + pass + + def post_round_function(self): + pass + + def pre_finetune_function(self): + self._is_finetuning = True + self.is_finetuning = True + + def calculate_additional_loss(self): + return 0.0 + + def get_layer_sparsity(self, weight): + pass + + @torch.no_grad() + def post_epoch_function(self, epoch, total_epochs, **kwargs): + if not self._is_pretraining: + self.t.add_(1) diff --git a/src/pquant/core/torch/quantizer.py b/src/pquant/core/torch/quantizer.py index 7fd57f8..51f7e38 100644 --- a/src/pquant/core/torch/quantizer.py +++ b/src/pquant/core/torch/quantizer.py @@ -1,9 +1,9 @@ -from enum import Enum - import torch import torch.nn as nn -from pquant.core.quantizer_functions import create_quantizer +from pquant.core.constants import QuantizationGranularity +from pquant.core.torch.fixed_point_quantizer import get_fixed_quantizer +from pquant.core.torch.hgq_quantizer import HGQQuantizer class Quantizer(nn.Module): @@ -16,36 +16,45 @@ def __init__( round_mode, is_heterogeneous, is_data=False, - granularity='per_tensor', + granularity=QuantizationGranularity.PER_TENSOR, hgq_gamma=0, place="datalane", + dynamic_data=True, + shape=None, ): super().__init__() - self.k = torch.nn.Parameter(torch.tensor(k), requires_grad=False) + self.overflow = overflow - self.b_init = k + i + f self.round_mode = round_mode self.use_hgq = is_heterogeneous self.is_data = is_data - self.i_init = i - self.f_init = f - self.i = torch.nn.Parameter(torch.tensor(i), requires_grad=False) - self.f = torch.nn.Parameter(torch.tensor(f), requires_grad=False) - self.b = torch.nn.Parameter(torch.tensor(i + k + f), requires_grad=False) - self.quantizer = create_quantizer(self.k, i, f, self.overflow, self.round_mode, self.use_hgq, self.is_data, place) + self.dynamic_data = dynamic_data + self.granularity = QuantizationGranularity(granularity).value + + # Params even when using HGQ, as they are used during hls4ml conversion + param_shape = () if (is_data or self.use_hgq) else self.compute_weight_param_shape(shape) + self.k = torch.nn.Parameter(torch.full(param_shape, float(k)), requires_grad=False) + self.i = torch.nn.Parameter(torch.full(param_shape, float(i)), requires_grad=False) + self.f = torch.nn.Parameter(torch.full(param_shape, float(f)), requires_grad=False) + self.b = torch.nn.Parameter(torch.full(param_shape, float(i + k + f)), requires_grad=False) + self.quantizer = create_quantizer( + k, + i, + f, + self.overflow, + self.round_mode, + self.use_hgq, + self.is_data, + granularity=self.granularity, + gamma=hgq_gamma, + ) self.is_pretraining = True self.hgq_gamma = hgq_gamma - if isinstance(granularity, Enum): - self.granularity = granularity.value - else: - self.granularity = granularity - self.final_compression_done = nn.Parameter(torch.tensor(False), requires_grad=False) - if self.granularity == 'per_tensor': - self.initialize_quantization_parameters(self.i_init, self.f_init) + self.register_buffer("final_compression_done", torch.tensor(False)) def get_quantization_bits(self): if self.use_hgq: - return self.quantizer.quantizer.k, self.quantizer.quantizer.i, self.quantizer.quantizer.f + return self.quantizer.k, self.quantizer.i, self.quantizer.f else: return self.k, self.i, self.f @@ -56,89 +65,121 @@ def get_total_bits(self, shape): b = self.i + self.f + self.k return torch.ones(shape).to(b.device) * b + def _sync_hgq_mirror_bits(self): + if not self.quantizer.built: + return + with torch.no_grad(): + k, i, f = self.quantizer.k.detach(), self.quantizer.i.detach(), self.quantizer.f.detach() + self.k.data = k.clone() + self.i.data = i.clone() + self.f.data = f.clone() + self.b.data = k + i + f + def set_quantization_bits(self, i, f): if self.use_hgq: - self.quantizer.quantizer._i.assign(self.quantizer.quantizer._i * 0.0 + i) - self.quantizer.quantizer._f.assign(self.quantizer.quantizer._f * 0.0 + f) - self.i.data = torch.tensor(i) - self.f.data = torch.tensor(f) + self.quantizer.set_bits(i, f) + else: + self.i.data = torch.as_tensor(i, dtype=self.i.dtype, device=self.i.device).broadcast_to(self.i.shape).clone() + self.f.data = torch.as_tensor(f, dtype=self.f.dtype, device=self.f.device).broadcast_to(self.f.shape).clone() def post_pre_train_function(self): self.is_pretraining = False - def compute_dynamic_bits(self, x): - if self.granularity == "per_channel": + def calculate_bits_from_abs(self, abs_x): + m = torch.ceil(torch.log2(abs_x + 1e-6)) + int_bits = torch.clamp(m, min=0).clamp(max=self.b - self.k.to(m.device)) + frac_bits = torch.clamp(self.b - int_bits - self.k, min=0) + return int_bits, frac_bits + + def compute_data_dynamic_bits(self, x): + if not (self.training and self.dynamic_data): + _, i, f = self.get_quantization_bits() + return i, f + abs_x = torch.amax(torch.abs(x)) + return self.calculate_bits_from_abs(abs_x) + + def compute_weight_param_shape(self, shape): + if shape is None or self.granularity == QuantizationGranularity.PER_TENSOR or len(shape) == 1: + return () + elif self.granularity == QuantizationGranularity.PER_CHANNEL: + return (shape[0],) + (1,) * (len(shape) - 1) # Channels first + else: + return shape + + def compute_weight_dynamic_bits(self, x): + if self.granularity == QuantizationGranularity.PER_TENSOR or x.ndim == 1 or not self.training: + _, i, f = self.get_quantization_bits() + return i, f + if self.granularity == QuantizationGranularity.PER_CHANNEL: if x.ndim == 2: abs_x = torch.amax(torch.abs(x), dim=1, keepdim=True) elif x.ndim == 3: abs_x = torch.amax(torch.abs(x), dim=(1, 2), keepdim=True) elif x.ndim == 4: abs_x = torch.amax(torch.abs(x), dim=(1, 2, 3), keepdim=True) - elif self.granularity == "per_weight": + elif self.granularity == QuantizationGranularity.PER_WEIGHT: abs_x = torch.abs(x) else: raise ValueError("The selected granularity is not supported.") + return self.calculate_bits_from_abs(abs_x) - m = torch.ceil(torch.log2(abs_x + 1e-6)) - int_bits = torch.clamp(m, min=0) - b = self.b if hasattr(self, "b") else self.k + self.i_init + self.f_init - frac_bits = torch.clamp(b - int_bits - self.k, min=0) - return int_bits, frac_bits + def compute_dynamic_bits(self, x): + if self.is_data: + return self.compute_data_dynamic_bits(x) + return self.compute_weight_dynamic_bits(x) def forward(self, x): if self.use_hgq: - x = self.quantizer(x, training=self.training) - _, i, f = self.get_quantization_bits() - self.initialize_quantization_parameters(i, f) - return x + return self.quantizer(x, training=self.training) + elif self.final_compression_done: + return self.quantizer(x, k=self.k, i=self.i, f=self.f, training=False) else: - if self.granularity == 'per_tensor': - self.initialize_quantization_parameters(self.i_init, self.f_init) - _, i, f = self.get_quantization_bits() - return self.quantizer(x, k=self.k, i=i, f=f, training=self.training) - else: - i, f = self.compute_dynamic_bits(x) - self.initialize_quantization_parameters(i, f) + i, f = self.compute_dynamic_bits(x) self.i.data = i self.f.data = f - _, i, f = self.get_quantization_bits() x = self.quantizer(x, k=self.k, i=i, f=f, training=self.training) return x def hgq_loss(self): if self.is_pretraining or not self.use_hgq: return 0.0 - loss = 0 - for layer_loss in self.quantizer.quantizer.losses: - loss += layer_loss - return loss + return self.quantizer.regularization_loss() def post_epoch_function(self): - if self.use_hgq and self.quantizer.quantizer.built: - constrained_i = self.quantizer.quantizer._i.constraint(self.quantizer.quantizer._i) - self.quantizer.quantizer._i.assign(constrained_i) - constrained_f = self.quantizer.quantizer._f.constraint(self.quantizer.quantizer._f) - self.quantizer.quantizer._f.assign(constrained_f) + if self.use_hgq and self.quantizer.built: + self.quantizer.post_epoch_constraint_apply() def apply_final_compression(self): if self.use_hgq and not self.quantizer.built: return + if self.use_hgq: + with torch.no_grad(): + self.quantizer._f.data.clamp_(self.quantizer.f_min, self.quantizer.f_max) + if self.quantizer.overflow_mode != "WRAP": + self.quantizer._i.data.clamp_(self.quantizer.i_min, self.quantizer.i_max) + self._sync_hgq_mirror_bits() + self.final_compression_done.fill_(True) + return _, i, f = self.get_quantization_bits() self.i.data = i self.f.data = f self.b.data = i + f - self.final_compression_done.data = torch.tensor(True) - - def initialize_quantization_parameters(self, i, f): - if hasattr(self, "f"): - return - # Lazy initialization - self.i = torch.nn.Parameter(torch.tensor(i), requires_grad=False) - self.f = torch.nn.Parameter(torch.tensor(f), requires_grad=False) - self.b = torch.nn.Parameter(torch.tensor(self.k.detach().clone() + i + f), requires_grad=False) - - def reload_from_local(self): - if not self.use_hgq: - return - self.quantizer.quantizer._i.assign(self.i) - self.quantizer.quantizer._f.assign(self.f) + self.final_compression_done.fill_(True) + + +def create_quantizer( + k, i, f, overflow, round_mode, is_heterogeneous, is_data, granularity=QuantizationGranularity.PER_WEIGHT, gamma=1e-8 +): + if is_heterogeneous: + return HGQQuantizer( + k0=k, + i0=i, + f0=f, + overflow_mode=overflow, + round_mode=round_mode, + is_data=is_data, + granularity=granularity, + gamma=gamma, + ) + else: + return get_fixed_quantizer(round_mode=round_mode, overflow_mode=overflow) diff --git a/src/pquant/core/torch/tracing.py b/src/pquant/core/torch/tracing.py new file mode 100644 index 0000000..a3dfdf5 --- /dev/null +++ b/src/pquant/core/torch/tracing.py @@ -0,0 +1,410 @@ +import logging + +import torch +import torch.nn.functional as F + +from pquant.core.constants import QuantizationGranularity +from pquant.core.torch.activations import PQActivation +from pquant.core.torch.layers import ( + PQAvgPoolBase, + PQBatchNorm1d, + PQBatchNorm2d, + PQWeightBiasBase, +) +from pquant.core.torch.quantizer import Quantizer + +_PQUANTML_LAYER_TYPES = (PQWeightBiasBase, PQAvgPoolBase, PQBatchNorm1d, PQBatchNorm2d, PQActivation) + + +def _analyze_quantization(model): + """Trace model with torch.fx and return (traced, node_issues, edges_to_quantize). + + node_issues maps each problematic Node to a list of human-readable + descriptions of what quantization is missing at that node. + + edges_to_quantize is a set of (producer_node, consumer_node) tuples + identifying graph edges on which a data quantizer should be inserted to + fix the missing quantization. + """ + import operator + from collections import defaultdict + + from torch.fx import GraphModule, Tracer + + class _PQTracer(Tracer): + def is_leaf_module(self, m, module_qualified_name): + if isinstance(m, _PQUANTML_LAYER_TYPES) or isinstance(m, Quantizer): + return True + return super().is_leaf_module(m, module_qualified_name) + + tracer = _PQTracer() + graph = tracer.trace(model) + traced = GraphModule(tracer.root, graph) + modules = dict(traced.named_modules()) + + arith_functions = { + operator.add, + operator.iadd, + operator.sub, + operator.isub, + operator.mul, + operator.imul, + operator.matmul, + operator.imatmul, + operator.truediv, + operator.itruediv, + operator.floordiv, + operator.ifloordiv, + operator.pow, + operator.ipow, + torch.add, + torch.sub, + torch.mul, + torch.div, + torch.divide, + torch.true_divide, + torch.floor_divide, + torch.matmul, + torch.bmm, + torch.mm, + torch.einsum, + torch.pow, + torch.cat, + torch.stack, + } + for _name in ("concat", "concatenate"): + _fn = getattr(torch, _name, None) + if _fn is not None: + arith_functions.add(_fn) + + _nonlin_torch_names = ( + "sigmoid", + "tanh", + "exp", + "log", + "log2", + "log10", + "sqrt", + "rsqrt", + "reciprocal", + "softmax", + "log_softmax", + "sin", + "cos", + "tan", + ) + _nonlin_F_names = ( + "sigmoid", + "tanh", + "softmax", + "log_softmax", + "gelu", + "silu", + "elu", + "selu", + "softplus", + "mish", + "hardsigmoid", + "hardswish", + "leaky_relu", + ) + nonlin_functions = set() + for _name in _nonlin_torch_names: + _fn = getattr(torch, _name, None) + if _fn is not None: + nonlin_functions.add(_fn) + for _name in _nonlin_F_names: + _fn = getattr(F, _name, None) + if _fn is not None: + nonlin_functions.add(_fn) + + quant_sensitive_functions = arith_functions | nonlin_functions + quant_sensitive_methods = { + "add", + "add_", + "__add__", + "__radd__", + "__iadd__", + "sub", + "sub_", + "__sub__", + "__rsub__", + "__isub__", + "mul", + "mul_", + "__mul__", + "__rmul__", + "__imul__", + "matmul", + "__matmul__", + "__rmatmul__", + "__imatmul__", + "div", + "div_", + "__truediv__", + "__rtruediv__", + "__itruediv__", + "floor_divide", + "floor_divide_", + "__floordiv__", + "__rfloordiv__", + "__ifloordiv__", + "true_divide", + "true_divide_", + "pow", + "pow_", + "__pow__", + "__rpow__", + "__ipow__", + "bmm", + "mm", + "einsum", + "sigmoid", + "sigmoid_", + "tanh", + "tanh_", + "exp", + "exp_", + "log", + "log_", + "log2", + "log2_", + "log10", + "log10_", + "sqrt", + "sqrt_", + "rsqrt", + "rsqrt_", + "reciprocal", + "reciprocal_", + "softmax", + "log_softmax", + "sin", + "sin_", + "cos", + "cos_", + "tan", + "tan_", + } + + def is_quant_sensitive(node): + if node.op == "call_function": + return node.target in quant_sensitive_functions + if node.op == "call_method": + return node.target in quant_sensitive_methods + return False + + def get_module(node): + if node.op == "call_module": + return modules.get(node.target) + return None + + quantized = {} + node_issues = defaultdict(list) + edges_to_quantize = set() + + for node in traced.graph.nodes: + input_nodes = node.all_input_nodes + all_inputs_quantized = bool(input_nodes) and all(quantized.get(n, False) for n in input_nodes) + + if node.op in ("placeholder", "get_attr"): + quantized[node] = False + elif node.op == "call_module": + mod = get_module(node) + if isinstance(mod, _PQUANTML_LAYER_TYPES): + if not getattr(mod, "quantize_input", False) and not all_inputs_quantized: + node_issues[node].append( + f"PQuantML layer '{node.target}' has quantize_input=False but receives unquantized input" + ) + for n in input_nodes: + if not quantized.get(n, False): + edges_to_quantize.add((n, node)) + # A relu activation is a grid-preserving clip (its optional multiplier is a + # power-of-two scale), so if the value reaching it is already quantized the + # output stays on-grid and needs no output quantizer. + grid_preserving = isinstance(mod, PQActivation) and mod.activation_name == "relu" + input_quantized = bool(getattr(mod, "quantize_input", False)) or all_inputs_quantized + quantized[node] = bool(getattr(mod, "quantize_output", False)) or (grid_preserving and input_quantized) + elif isinstance(mod, Quantizer): + quantized[node] = True + else: + quantized[node] = all_inputs_quantized + elif node.op in ("call_function", "call_method"): + if is_quant_sensitive(node): + for n in input_nodes: + if not quantized.get(n, False): + node_issues[node].append(f"input '{n.name}' is not quantized") + edges_to_quantize.add((n, node)) + quantized[node] = False + else: + quantized[node] = all_inputs_quantized + elif node.op == "output": + for n in input_nodes: + if not quantized.get(n, False): + node_issues[node].append(f"model output '{n.name}' is not quantized") + edges_to_quantize.add((n, node)) + quantized[node] = all_inputs_quantized + else: + quantized[node] = False + + return traced, node_issues, edges_to_quantize + + +def _insert_missing_quantizers(traced, edges_to_quantize, config): + from collections import defaultdict + + qp = config.quantization_parameters + + def _make_quantizer(): + return Quantizer( + k=qp.default_data_keep_negatives, + i=qp.default_data_integer_bits, + f=qp.default_data_fractional_bits, + overflow=qp.overflow_mode_data, + round_mode=qp.round_mode, + is_heterogeneous=False, + is_data=True, + granularity=QuantizationGranularity.PER_TENSOR, + hgq_gamma=qp.hgq_gamma, + ) + + def _enable_pquantml_output_quantization(layer): + layer.quantize_output = True + if isinstance(layer, PQWeightBiasBase) and getattr(layer, "built", False) and not hasattr(layer, "output_quantizer"): + device = next(layer.parameters()).device + layer.output_quantizer = Quantizer( + k=torch.tensor(layer.k_output), + i=torch.tensor(layer.i_output), + f=torch.tensor(layer.f_output), + overflow=layer.overflow_mode_data, + round_mode=layer.round_mode, + is_heterogeneous=layer.use_hgq, + is_data=True, + hgq_gamma=layer.hgq_gamma, + place="datalane", + ).to(device) + + modules = dict(traced.named_modules()) + + by_producer = defaultdict(list) + for producer, consumer in edges_to_quantize: + by_producer[producer].append(consumer) + + graph = traced.graph + idx = 0 + for producer, consumers in by_producer.items(): + pqml_layer = None + if producer.op == "call_module": + mod = modules.get(producer.target) + if isinstance(mod, _PQUANTML_LAYER_TYPES): + pqml_layer = mod + if pqml_layer is not None: + _enable_pquantml_output_quantization(pqml_layer) + continue + q_name = f"_auto_missing_quantizer_{idx}" + idx += 1 + traced.add_module(q_name, _make_quantizer()) + with graph.inserting_after(producer): + qnode = graph.call_module(q_name, (producer,)) + for consumer in consumers: + consumer.replace_input_with(producer, qnode) + + graph.lint() + traced.recompile() + return traced + + +def check_quantization(model, add_missing_quantizers=False, config=None): + """Verify quantization is present everywhere in the model's forward graph. + + The model's input is assumed to be unquantized. The model is traced with + torch.fx; PQuantML layers and Quantizer modules are treated as leaves. + For each node the output's quantization state is propagated forward: + - PQuantML layers (PQWeightBiasBase, PQAvgPoolBase, PQBatchNorm1d, + PQBatchNorm2d, PQActivation): output is quantized iff quantize_output + is True. If quantize_input is False and the incoming data is not + already quantized, that is reported as missing quantization. + - Quantizer modules always produce quantized output. + - Non-PQuantML quant-sensitive ops — arithmetic/combining + (add/sub/mul/div/matmul/bmm/mm/einsum/cat/stack/pow) and off-grid + nonlinearities (sigmoid/tanh/softmax/exp/log/sqrt/rsqrt/reciprocal/ + gelu/silu/elu/selu/softplus/mish/hardsigmoid/hardswish/leaky_relu/ + sin/cos/tan), whether used as torch functions, torch.nn.functional + functions, operator.* functions, or tensor methods — require each + input to already be quantized; their own output is marked unquantized + so the next consumer's per-edge input check flags it if needed. + - The model's output node(s) are required to be quantized: any + unquantized return value is flagged. + - Other ops (shape-only views, indexing, non-arithmetic torch calls) + propagate the quantization state of their inputs. + + When `add_missing_quantizers` is False (default), returns True if every + location in the graph has the required quantization, otherwise a list of + strings describing each missing quantization. + + When `add_missing_quantizers` is True, `config` must be provided. For + every producer whose output feeds an edge with missing quantization: + - if the producer is a PQuantML layer, its `quantize_output` flag is + set to True (and, for built `PQWeightBiasBase` subclasses that lacked + an output_quantizer, one is constructed in-place from the layer's + own k_output/i_output/f_output and mode settings); + - otherwise, a fresh `Quantizer` module is instantiated with the + config's default data k/i/f and data round/overflow modes, attached + to the transformed module as `_auto_missing_quantizer_`, and a + `call_module` to that new quantizer is inserted on each affected + edge. Using one Quantizer per insertion site (rather than sharing + one) allows heterogeneous quantization to specialize per site. + Returns the transformed torch.fx GraphModule. + """ + traced, node_issues, edges_to_quantize = _analyze_quantization(model) + + if add_missing_quantizers: + if config is None: + raise ValueError("check_quantization(add_missing_quantizers=True) requires config") + if edges_to_quantize: + _insert_missing_quantizers(traced, edges_to_quantize, config) + return traced + + if not node_issues: + return True + messages = [] + for node, msgs in node_issues.items(): + for msg in msgs: + messages.append(f"'{node.name}' ({node.op}): {msg}") + return messages + + +def print_quantization_check(model, use_color=True): + """Print the traced model's graph with missing-quantization nodes flagged. + + Runs `check_quantization(model)` and prints the torch.fx graph line by + line. Nodes with missing quantization are highlighted in red (git-diff + style) and followed by one annotated line per issue. + + Set use_color=False to disable ANSI escape codes (e.g. when redirecting + to a file or a terminal that does not support colors). + + Returns the same value as `check_quantization(model)`. + """ + traced, node_issues, _ = _analyze_quantization(model) + + RED = "\033[31m" if use_color else "" + BOLD = "\033[1m" if use_color else "" + RESET = "\033[0m" if use_color else "" + + for node in traced.graph.nodes: + line = node.format_node() or f"{node.name} = {node.op} {node.target}" + if node in node_issues: + logging.info(f"{RED}{BOLD}- {line}{RESET}") + for msg in node_issues[node]: + logging.info(f"{RED} ! {msg}{RESET}") + else: + logging.info(f" {line}") + + if not node_issues: + return True + messages = [] + for node, msgs in node_issues.items(): + for msg in msgs: + messages.append(f"'{node.name}' ({node.op}): {msg}") + return messages diff --git a/src/pquant/core/torch/utils.py b/src/pquant/core/torch/utils.py new file mode 100644 index 0000000..6f6f466 --- /dev/null +++ b/src/pquant/core/torch/utils.py @@ -0,0 +1,28 @@ +from pquant.core.torch.pruning_methods.activation_pruning import ActivationPruning +from pquant.core.torch.pruning_methods.autosparse import AutoSparse +from pquant.core.torch.pruning_methods.cs import ContinuousSparsification +from pquant.core.torch.pruning_methods.dst import DST +from pquant.core.torch.pruning_methods.fitcompress import FITCompress +from pquant.core.torch.pruning_methods.mdmm import MDMM +from pquant.core.torch.pruning_methods.pdp import PDP +from pquant.core.torch.pruning_methods.wanda import Wanda + + +def get_pruning_layer(config, layer_type): + pruning_method = config.pruning_parameters.pruning_method + if pruning_method == "dst": + return DST(config, layer_type) + elif pruning_method == "autosparse": + return AutoSparse(config, layer_type) + elif pruning_method == "cs": + return ContinuousSparsification(config, layer_type) + elif pruning_method == "fitcompress": + return FITCompress(config) + elif pruning_method == "pdp": + return PDP(config, layer_type) + elif pruning_method == "activation_pruning": + return ActivationPruning(config, layer_type) + elif pruning_method == "wanda": + return Wanda(config, layer_type) + elif pruning_method == "mdmm": + return MDMM(config, layer_type) diff --git a/src/pquant/core/utils.py b/src/pquant/core/utils.py index fe9e575..072e7bd 100644 --- a/src/pquant/core/utils.py +++ b/src/pquant/core/utils.py @@ -2,35 +2,6 @@ import yaml -from pquant.pruning_methods.activation_pruning import ActivationPruning -from pquant.pruning_methods.autosparse import AutoSparse -from pquant.pruning_methods.cs import ContinuousSparsification -from pquant.pruning_methods.dst import DST -from pquant.pruning_methods.fitcompress import FITCompress -from pquant.pruning_methods.mdmm import MDMM -from pquant.pruning_methods.pdp import PDP -from pquant.pruning_methods.wanda import Wanda - - -def get_pruning_layer(config, layer_type): - pruning_method = config.pruning_parameters.pruning_method - if pruning_method == "dst": - return DST(config, layer_type) - elif pruning_method == "autosparse": - return AutoSparse(config, layer_type) - elif pruning_method == "cs": - return ContinuousSparsification(config, layer_type) - elif pruning_method == "pdp": - return PDP(config, layer_type) - elif pruning_method == "activation_pruning": - return ActivationPruning(config, layer_type) - elif pruning_method == "wanda": - return Wanda(config, layer_type) - elif pruning_method == "mdmm": - return MDMM(config, layer_type) - elif pruning_method == "fitcompress": - return FITCompress(config) - def get_default_config(pruning_method: str): assert pruning_method in [ @@ -38,7 +9,6 @@ def get_default_config(pruning_method: str): "ap", "cs", "dst", - "fitcompress", "pdp", "wanda", "mdmm", diff --git a/src/pquant/data_models/fitcompress_model.py b/src/pquant/data_models/fitcompress_model.py index 8ea61d6..f0b0124 100644 --- a/src/pquant/data_models/fitcompress_model.py +++ b/src/pquant/data_models/fitcompress_model.py @@ -1,5 +1,3 @@ -from typing import List - from pydantic import BaseModel, Field @@ -12,7 +10,7 @@ class PruningSchedule(BaseModel): class BaseFitCompressModel(BaseModel): enable_fitcompress: bool = Field(default=False) optimize_quantization: bool = Field(default=True) - quantization_schedule: List[float] = Field(default_factory=lambda: [7.0, 4.0, 3.0, 2.0]) + quantization_schedule: list[float] = Field(default_factory=lambda: [7.0, 4.0, 3.0, 2.0]) pruning_schedule: PruningSchedule = Field(default_factory=PruningSchedule) compression_goal: float = Field(default=0.10) optimize_pruning: bool = Field(default=False) diff --git a/src/pquant/data_models/hyperparameter_optimization_model.py b/src/pquant/data_models/hyperparameter_optimization_model.py index 1c53ca7..258ecd9 100644 --- a/src/pquant/data_models/hyperparameter_optimization_model.py +++ b/src/pquant/data_models/hyperparameter_optimization_model.py @@ -1,16 +1,16 @@ -from typing import Any, Dict, List, Optional, Union +from typing import Any from pydantic import BaseModel, Field class HyperparameterSearch(BaseModel): - numerical: Dict[str, List[Union[int, float]]] = Field(default_factory=dict) - categorical: Optional[Dict[str, List[str]]] = Field(default_factory=dict) + numerical: dict[str, list[int | float]] = Field(default_factory=dict) + categorical: dict[str, list[str]] | None = Field(default_factory=dict) class Sampler(BaseModel): type: str = Field(default="TPESampler") - params: Dict[str, Any] = Field(default_factory=dict) + params: dict[str, Any] = Field(default_factory=dict) class BaseHyperparameterOptimizationModel(BaseModel): diff --git a/src/pquant/data_models/pruning_model.py b/src/pquant/data_models/pruning_model.py index 21e6607..11b797a 100644 --- a/src/pquant/data_models/pruning_model.py +++ b/src/pquant/data_models/pruning_model.py @@ -1,11 +1,11 @@ from enum import Enum -from typing import List, Literal, Optional +from typing import Literal from pydantic import BaseModel, Field class BasePruningModel(BaseModel): - disable_pruning_for_layers: List[str] = Field(default_factory=list) + disable_pruning_for_layers: list[str] = Field(default_factory=list) enable_pruning: bool = Field(default=True) threshold_decay: float = Field(default=0.0) @@ -39,8 +39,8 @@ class PDPPruningModel(BasePruningModel): class WandaPruningModel(BasePruningModel): pruning_method: Literal["wanda"] = "wanda" - M: Optional[int] = (Field(default=None),) - N: Optional[int] = (Field(default=None),) + M: int | None = (Field(default=None),) + N: int | None = (Field(default=None),) sparsity: float = Field(default=0.9) t_delta: int = Field(default=100) t_start_collecting_batch: int = Field(default=100) diff --git a/src/pquant/data_models/quantization_model.py b/src/pquant/data_models/quantization_model.py index 31cefd8..bcd3976 100644 --- a/src/pquant/data_models/quantization_model.py +++ b/src/pquant/data_models/quantization_model.py @@ -1,12 +1,6 @@ -from typing import List -from enum import Enum from pydantic import BaseModel, Field - -class QuantizationGranularity(str, Enum): - PER_TENSOR = "per_tensor" - PER_CHANNEL = "per_channel" - PER_WEIGHT = "per_weight" +from pquant.core.constants import QuantizationGranularity class BaseQuantizationModel(BaseModel): @@ -19,6 +13,7 @@ class BaseQuantizationModel(BaseModel): quantize_input: bool = Field(default=True) quantize_output: bool = Field(default=False) granularity: QuantizationGranularity = Field(default=QuantizationGranularity.PER_TENSOR) + dynamic_data_quantization: bool = Field(default=False) enable_quantization: bool = Field(default=True) hgq_gamma: float = Field(default=0.0003) hgq_beta: float = Field(default=1e-5) diff --git a/src/pquant/data_models/training_model.py b/src/pquant/data_models/training_model.py index f841d70..abd70a9 100644 --- a/src/pquant/data_models/training_model.py +++ b/src/pquant/data_models/training_model.py @@ -1,10 +1,8 @@ -from typing import Literal - from pydantic import BaseModel, ConfigDict, Field class BaseTrainingModel(BaseModel): - model_config = ConfigDict(extra='allow') + model_config = ConfigDict(extra="allow") epochs: int = Field(default=200) fine_tuning_epochs: int = Field(default=0) pretraining_epochs: int = Field(default=50) diff --git a/src/pquant/pruning_methods/fitcompress.py b/src/pquant/pruning_methods/fitcompress.py deleted file mode 100644 index f2b3b5c..0000000 --- a/src/pquant/pruning_methods/fitcompress.py +++ /dev/null @@ -1,52 +0,0 @@ -import keras - - -@keras.saving.register_keras_serializable(package="Layers") -class FITCompress(keras.layers.Layer): - def __init__(self, config, *args, **kwargs): - super().__init__(*args, **kwargs) - if isinstance(config, dict): - from pquant.core.hyperparameter_optimization import PQConfig - - config = PQConfig.load_from_config(config) - self.config = config - self.is_pretraining = True - self.is_finetuning = False - - def build(self, input_shape): - self.mask = self.add_weight(shape=input_shape, initializer="ones", trainable=False) - super().build(input_shape) - - def call(self, weight): - return self.mask * weight - - def get_hard_mask(self, weight=None): - return self.mask - - def pre_epoch_function(self, epoch, total_epochs): - pass - - def calculate_additional_loss(*args, **kwargs): - return 0 - - def pre_finetune_function(self): - self.is_finetuning = True - - def post_round_function(self): - pass - - def post_pre_train_function(self): - self.is_pretraining = False - - def post_epoch_function(self, epoch, total_epochs): - pass - - def get_config(self): - config = super().get_config() - - config.update( - { - "config": self.config.get_dict(), - } - ) - return config diff --git a/tests/conftest.py b/tests/conftest.py index 479ab00..c26d7ec 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,13 +8,10 @@ @pytest.fixture(scope="function", autouse=True) def set_image_data_format(): - if "DATA_FORMAT" in os.environ: - keras.backend.set_image_data_format(os.environ["DATA_FORMAT"]) - else: - keras.backend.set_image_data_format("channels_first") + keras.backend.set_image_data_format(os.environ.get("DATA_FORMAT", "channels_last")) -@pytest.fixture(scope='session', autouse=True, params=[42]) +@pytest.fixture(scope="session", autouse=True, params=[42]) def set_random_seed(request): """Set random seeds for reproducibility""" @@ -23,31 +20,36 @@ def set_random_seed(request): random.seed(seed) backend = keras.backend.backend() match backend: - case 'tensorflow': + case "tensorflow": import tensorflow as tf tf.random.set_seed(seed) - case 'torch': + case "torch": import torch torch.manual_seed(seed) case _: - raise ValueError(f'Unknown backend: {backend}') + raise ValueError(f"Unknown backend: {backend}") -@pytest.fixture(scope='session', autouse=True) +@pytest.fixture(scope="session", autouse=True) def configure_backend(): backend = keras.backend.backend() match backend: - case 'tensorflow': - pass - case 'torch': + case "tensorflow": + import tensorflow as tf + + # Use full float32 matmul precision to match numpy/onnxruntime. Without this, TF uses + # TF32 on Ampere+ GPUs and the ~1e-3 relative error exceeds tight test tolerances + # (e.g. the keras→ONNX parity tests). Mirrors the torch 'highest' setting below. + tf.config.experimental.enable_tensor_float_32_execution(False) + case "torch": import torch - torch.set_float32_matmul_precision('highest') + torch.set_float32_matmul_precision("highest") device = "cuda" if torch.cuda.is_available() else "cpu" torch.set_default_device(device) torch.set_default_dtype(torch.float32) case _: - raise ValueError(f'Unknown backend: {backend}') + raise ValueError(f"Unknown backend: {backend}") diff --git a/tests/run_tests.sh b/tests/run_tests.sh index df2f5fc..34fedbc 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -6,6 +6,13 @@ pytest test_pdp.py KERAS_BACKEND="torch" pytest test_pdp.py pytest test_wanda.py KERAS_BACKEND="torch" pytest test_wanda.py +pytest test_torch_pruning_layers.py +pytest test_quantizer_parity.py pytest test_keras_compression_layers.py -DATA_FORMAT=channels_last pytest test_keras_compression_layers.py KERAS_BACKEND="torch" pytest test_torch_compression_layers.py +KERAS_BACKEND=torch pytest test_torch_onnx_converter.py +pytest test_keras_onnx_converter.py +KERAS_BACKEND=torch pytest test_torch_alkaid_conversion.py +pytest test_keras_alkaid_conversion.py +KERAS_BACKEND=torch pytest test_hgq_torch.py +pytest test_hgq_keras.py diff --git a/tests/test_ap.py b/tests/test_ap.py index d126898..afd4d9c 100644 --- a/tests/test_ap.py +++ b/tests/test_ap.py @@ -2,7 +2,7 @@ from keras import ops from keras.random import shuffle -from pquant.pruning_methods.activation_pruning import ActivationPruning +from pquant.core.keras.pruning_methods.activation_pruning import ActivationPruning @pytest.fixture diff --git a/tests/test_hgq_keras.py b/tests/test_hgq_keras.py new file mode 100644 index 0000000..780891e --- /dev/null +++ b/tests/test_hgq_keras.py @@ -0,0 +1,241 @@ +"""Tests for HGQ (high granularity quantization) per-tensor / per-channel / per-weight granularity. + +The granularity controls the shape of the trainable bitwidth tensors (`i` and `f`) inside the +HGQ quantizer: + - per_tensor: a single shared value for the whole tensor -> shape collapses to all-ones + - per_channel: one value per output channel -> only the output-channel axis is kept + - per_weight: one value per element -> shape matches the tensor itself + (for data the batch axis is always shared, so it collapses to 1) + +In Keras, kernels are stored output-channel-last and we assume channels_last data, so the +output-channel axis is -1 for both weights and activations. +""" + +import keras +import numpy as np +import pytest +from keras import ops +from pquant.layers import ( + PQAvgPool1d, + PQBatchNormalization, + PQConv1d, + PQConv2d, + PQDense, + PQMultiheadAttention, +) + +from pquant import pdp_config +from pquant.core.keras.quantizer import Quantizer + +BATCH_SIZE = 4 +IN_FEATURES = 16 +OUT_FEATURES = 32 +KERNEL_SIZE = 3 +STEPS = 8 + +# HGQ supports only per_tensor and per_weight; per_channel is rejected. +GRANULARITIES = ["per_tensor", "per_weight"] + + +@pytest.fixture(autouse=True) +def run_around_tests(): + keras.backend.clear_session() + + +def hgq_config(granularity): + """A PQConfig with high granularity quantization enabled for the given granularity.""" + config = pdp_config() + config.quantization_parameters.use_high_granularity_quantization = True + config.quantization_parameters.enable_quantization = True + config.quantization_parameters.granularity = granularity + return config + + +def shape_of(tensor): + return tuple(int(d) for d in ops.shape(tensor)) + + +def is_single_value(shape): + """per_tensor: every axis collapsed to 1.""" + return int(np.prod(shape)) == 1 + + +# ----------------------------------------------------------------------------- weights + + +def assert_weight_granularity(layer, granularity, kernel_shape): + """The weight quantizer's i/f bitwidth tensors must match the expected shape for `granularity`.""" + _, i, f = layer.get_weight_quantization_bits() + for name, t in (("i", i), ("f", f)): + shape = shape_of(t) + if granularity == "per_tensor": + assert is_single_value(shape), f"weight {name}: expected single value, got {shape}" + else: # per_weight + assert shape == kernel_shape, f"weight {name}: expected {kernel_shape}, got {shape}" + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +def test_dense_weight_granularity(granularity): + layer = PQDense(hgq_config(granularity), units=OUT_FEATURES) + layer.build((BATCH_SIZE, IN_FEATURES)) + assert_weight_granularity(layer, granularity, kernel_shape=(IN_FEATURES, OUT_FEATURES)) + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +def test_conv1d_weight_granularity(granularity): + layer = PQConv1d(hgq_config(granularity), filters=OUT_FEATURES, kernel_size=KERNEL_SIZE, data_format="channels_last") + layer.build((BATCH_SIZE, STEPS, IN_FEATURES)) + assert_weight_granularity(layer, granularity, kernel_shape=(KERNEL_SIZE, IN_FEATURES, OUT_FEATURES)) + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +def test_conv2d_weight_granularity(granularity): + layer = PQConv2d(hgq_config(granularity), filters=OUT_FEATURES, kernel_size=KERNEL_SIZE, data_format="channels_last") + layer.build((BATCH_SIZE, STEPS, STEPS, IN_FEATURES)) + assert_weight_granularity(layer, granularity, kernel_shape=(KERNEL_SIZE, KERNEL_SIZE, IN_FEATURES, OUT_FEATURES)) + + +def test_per_channel_rejected_for_hgq(): + # per_channel is not a valid HGQ granularity; constructing the HGQ weight quantizer must raise. + with pytest.raises(ValueError, match="per_channel"): + PQDense(hgq_config("per_channel"), units=OUT_FEATURES) + + +# ------------------------------------------------------------------------------- data +# A data quantizer built on a tensor shaped like a layer's output. The batch axis is always +# shared, so per_weight keeps every non-batch axis and collapses only the batch axis to 1. + + +def assert_data_granularity(output_shape, granularity): + quantizer = Quantizer(k=0.0, i=0.0, f=7.0, is_heterogeneous=True, is_data=True, granularity=granularity) + quantizer.build(output_shape) + _, i, f = quantizer.get_quantization_bits() + for name, t in (("i", i), ("f", f)): + shape = shape_of(t) + if granularity == "per_tensor": + assert is_single_value(shape), f"data {name}: expected single value, got {shape}" + else: # per_weight: batch axis shared, rest per-element + assert shape == (1,) + tuple(output_shape[1:]), f"data {name}: expected batch-collapsed, got {shape}" + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +def test_dense_data_granularity(granularity): + assert_data_granularity((BATCH_SIZE, OUT_FEATURES), granularity) + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +def test_conv1d_data_granularity(granularity): + # channels_last, "valid" padding output length + assert_data_granularity((BATCH_SIZE, STEPS - KERNEL_SIZE + 1, OUT_FEATURES), granularity) + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +def test_conv2d_data_granularity(granularity): + out_len = STEPS - KERNEL_SIZE + 1 + assert_data_granularity((BATCH_SIZE, out_len, out_len, OUT_FEATURES), granularity) + + +# --------------------------------------------------- per-quantizer granularity override +# Each quantizer follows the config granularity unless its per-quantizer override is set. + + +def test_per_quantizer_granularity_override(): + # input per_tensor, everything else per_weight. + config = hgq_config("per_weight") + layer = PQDense(config, units=OUT_FEATURES, quantize_output=True, in_quant_granularity="per_tensor") + layer.build((BATCH_SIZE, IN_FEATURES)) + + # weight follows config (per_weight) -> full kernel shape + _, wi, _ = layer.get_weight_quantization_bits() + assert shape_of(wi) == (IN_FEATURES, OUT_FEATURES) + # input overridden to per_tensor -> single value + _, ii, _ = layer.get_input_quantization_bits() + assert is_single_value(shape_of(ii)) + # output keeps config per_weight -> batch-collapsed data shape + _, oi, _ = layer.get_output_quantization_bits() + assert shape_of(oi) == (1, OUT_FEATURES) + + +def test_per_quantizer_granularity_defaults_to_config(): + config = hgq_config("per_tensor") + layer = PQDense(config, units=OUT_FEATURES, weight_quant_granularity="per_weight") + layer.build((BATCH_SIZE, IN_FEATURES)) + # weight overridden to per_weight; input uses config per_tensor + _, wi, _ = layer.get_weight_quantization_bits() + assert shape_of(wi) == (IN_FEATURES, OUT_FEATURES) + _, ii, _ = layer.get_input_quantization_bits() + assert is_single_value(shape_of(ii)) + + +# --------------------------------------------------- granularity override on boundary layers +# A model can start/end with a BatchNorm / AvgPool whose input/output quantizer is effectively +# the model's I/O quantizer, so that must be overridable independently of the config granularity. + + +def test_batchnorm_input_granularity_override(): + config = hgq_config("per_weight") + layer = PQBatchNormalization(config, in_quant_granularity="per_tensor") + layer.build((BATCH_SIZE, IN_FEATURES)) + # input overridden to per_tensor -> single value + _, ii, _ = layer.get_input_quantization_bits() + assert is_single_value(shape_of(ii)) + # weight follows config per_weight -> one value per feature + _, wi, _ = layer.get_weight_quantization_bits() + assert shape_of(wi) == (IN_FEATURES,) + + +def test_avgpool_io_granularity_override(): + config = hgq_config("per_weight") + layer = PQAvgPool1d( + config, + pool_size=2, + quantize_output=True, + in_quant_granularity="per_tensor", + out_quant_granularity="per_tensor", + ) + layer.build((BATCH_SIZE, STEPS, OUT_FEATURES)) + _, ii, _ = layer.get_input_quantization_bits() + assert is_single_value(shape_of(ii)) # input overridden to per_tensor + _, oi, _ = layer.get_output_quantization_bits() + assert is_single_value(shape_of(oi)) # output overridden to per_tensor + + +def test_mha_io_param_granularity_override(): + config = hgq_config("per_weight") + # in/out are the model-boundary granularities; param (weight+bias) is uniform across projections. + layer = PQMultiheadAttention( + config, + embed_dim=IN_FEATURES, + num_heads=4, + quantize_output=True, + in_quant_granularity="per_tensor", + out_quant_granularity="per_tensor", + ) + # Build the projections directly (a full forward would hit an unrelated dtype issue in + # PQDense.ebops on 3-D input); granularity shapes are fixed at build time. + for proj in (layer.q_proj, layer.k_proj, layer.v_proj, layer.out_proj): + proj.build((2, 5, IN_FEATURES)) + # Q/K/V projection inputs (boundary) overridden to per_tensor + for proj in (layer.q_proj, layer.k_proj, layer.v_proj): + _, ii, _ = proj.get_input_quantization_bits() + assert is_single_value(shape_of(ii)) + # out_proj output (boundary) overridden to per_tensor + _, oi, _ = layer.out_proj.get_output_quantization_bits() + assert is_single_value(shape_of(oi)) + # weights follow param granularity = config per_weight + _, wi, _ = layer.q_proj.get_weight_quantization_bits() + assert shape_of(wi) == (IN_FEATURES, IN_FEATURES) + # out_proj input is internal -> stays config per_weight (batch-collapsed, not single) + _, oii, _ = layer.out_proj.get_input_quantization_bits() + assert not is_single_value(shape_of(oii)) + + +def test_mha_qkv_always_output_quantized(): + config = hgq_config("per_tensor") + # MHA-level quantize_output=False (default): Q/K/V outputs are matmul operands and stay + # output-quantized; only out_proj follows the MHA-level flag. + layer = PQMultiheadAttention(config, embed_dim=IN_FEATURES, num_heads=4) + assert layer.q_proj.quantize_output is True + assert layer.k_proj.quantize_output is True + assert layer.v_proj.quantize_output is True + assert layer.out_proj.quantize_output is False diff --git a/tests/test_hgq_torch.py b/tests/test_hgq_torch.py new file mode 100644 index 0000000..456d6aa --- /dev/null +++ b/tests/test_hgq_torch.py @@ -0,0 +1,490 @@ +""" +Parity tests: PyTorch `HGQQuantizer` vs Keras `hgq.Quantizer` (hgq2 library). + +Both implementations should produce matching forward outputs, matching +gradients on the fractional-bit parameter `f`, and should follow a similar +training trajectory when fed the same data with the same initial state. +""" + +import pytest # noqa: E402 +import torch # noqa: E402 + +hgq = pytest.importorskip("hgq") +from hgq.quantizer import Quantizer as KerasHGQQuantizer # noqa: E402 +from hgq.quantizer import QuantizerConfig # noqa: E402 + +from pquant.core.torch.hgq_quantizer import HGQQuantizer # noqa: E402 + +RTOL = 1e-4 +ATOL = 1e-5 + +# HGQ supports only per_tensor and per_weight; per_channel is rejected. +GRANULARITIES = ["per_tensor", "per_weight"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_keras_q(k, i, f, overflow, round_mode, is_data, place="datalane"): + """Construct the Keras hgq2 quantizer mirroring create_hgq_*_quantizer().""" + homogeneous_axis = (0,) if is_data else () + cfg = QuantizerConfig( + q_type="kif", + place="datalane" if is_data else place, + k0=k, + i0=i, + f0=f, + overflow_mode=overflow, + round_mode=round_mode, + homogeneous_axis=homogeneous_axis, + ) + return KerasHGQQuantizer(config=cfg) + + +def _as_torch(x): + """Unwrap a keras tensor produced with the torch backend to a torch.Tensor.""" + if isinstance(x, torch.Tensor): + return x + import keras.ops as ops + + return ops.convert_to_tensor(x) + + +# --------------------------------------------------------------------------- +# Forward parity +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "overflow,round_mode,is_data", + [ + ("SAT", "RND", False), + ("SAT", "RND_CONV", False), + ("SAT_SYM", "RND", False), + ("WRAP", "RND", False), + ("WRAP", "RND", True), + ("SAT", "RND", True), + ], +) +def test_forward_matches_keras(overflow, round_mode, is_data): + torch.manual_seed(0) + k, i, f = 1, 2, 4 + x = torch.randn(8, 5) * 0.8 + + torch_q = HGQQuantizer( + k0=k, + i0=i, + f0=f, + overflow_mode=overflow, + round_mode=round_mode, + is_data=is_data, + ) + keras_q = _make_keras_q(k, i, f, overflow, round_mode, is_data) + + out_torch = torch_q(x, training=False).detach() + # Keras quantizer: call with training=False + out_keras = _as_torch(keras_q(x, training=False)).detach() + + assert out_torch.shape == out_keras.shape, f"shape mismatch: {out_torch.shape} vs {out_keras.shape}" + assert torch.allclose(out_torch, out_keras, rtol=RTOL, atol=ATOL), ( + f"[{overflow}/{round_mode}/is_data={is_data}] max diff = {(out_torch - out_keras).abs().max().item():.6g}" + ) + + +def test_forward_matches_keras_training_sat(): + torch.manual_seed(1) + x = torch.randn(4, 6) * 1.2 + torch_q = HGQQuantizer(k0=1, i0=1, f0=3, overflow_mode="SAT", round_mode="RND", is_data=False) + keras_q = _make_keras_q(1, 1, 3, "SAT", "RND", is_data=False) + + out_torch = torch_q(x, training=True).detach() + out_keras = _as_torch(keras_q(x, training=True)).detach() + + assert torch.allclose(out_torch, out_keras, rtol=RTOL, atol=ATOL), ( + f"training forward diverged, max diff = {(out_torch - out_keras).abs().max().item():.6g}" + ) + + +# --------------------------------------------------------------------------- +# Backward parity — gradient through f +# --------------------------------------------------------------------------- + + +def test_backward_f_gradient_sat(): + torch.manual_seed(2) + x_data = torch.randn(4, 5) * 0.6 + + torch_q = HGQQuantizer(k0=1, i0=2, f0=4, overflow_mode="SAT", round_mode="RND", is_data=False) + keras_q = _make_keras_q(1, 2, 4, "SAT", "RND", is_data=False) + + # Forward to trigger build for both + x_t = x_data.clone().requires_grad_(True) + out_t = torch_q(x_t, training=True) + loss_t = out_t.pow(2).sum() + loss_t.backward() + grad_f_torch = torch_q._f.grad.detach().clone() + + x_k = x_data.clone().requires_grad_(True) + out_k = _as_torch(keras_q(x_k, training=True)) + loss_k = out_k.pow(2).sum() + loss_k.backward() + + # hgq2 stores the f parameter as keras_q.quantizer._f — access via state + keras_f = _find_f_param(keras_q) + grad_f_keras = keras_f.grad.detach().clone() if keras_f.grad is not None else None + + assert grad_f_keras is not None, "Keras hgq quantizer produced no gradient on _f" + assert grad_f_torch.shape == grad_f_keras.shape, f"grad f shape mismatch: {grad_f_torch.shape} vs {grad_f_keras.shape}" + # Gradient direction/magnitude should match up to STE discretisation noise. + assert torch.allclose(grad_f_torch, grad_f_keras, rtol=1e-3, atol=1e-5), ( + f"grad f mismatch, max diff = {(grad_f_torch - grad_f_keras).abs().max().item():.6g}" + ) + + +def test_backward_input_gradient_ste(): + """STE: gradient w.r.t. input should be approximately identity within sat range.""" + torch.manual_seed(3) + x = (torch.randn(3, 4) * 0.3).requires_grad_(True) # small → within sat bounds + + torch_q = HGQQuantizer(k0=1, i0=2, f0=4, overflow_mode="SAT", round_mode="RND", is_data=False) + out = torch_q(x, training=True) + out.sum().backward() + + assert x.grad is not None + assert torch.allclose(x.grad, torch.ones_like(x), atol=1e-5), ( + f"STE grad should be ~1 inside sat range, got max deviation {(x.grad - 1).abs().max().item():.6g}" + ) + + +# --------------------------------------------------------------------------- +# Training trajectory parity +# --------------------------------------------------------------------------- + + +def test_training_trajectory_matches(): + """Run several SGD steps on both implementations with same data/init. The + learned f parameter should stay close (exact equality not guaranteed due to + STE non-determinism in rounding ties, but trajectories should track).""" + torch.manual_seed(4) + k, i, f = 1, 2, 4 + lr = 0.05 + steps = 20 + + torch_q = HGQQuantizer(k0=k, i0=i, f0=f, overflow_mode="SAT", round_mode="RND", is_data=False) + keras_q = _make_keras_q(k, i, f, "SAT", "RND", is_data=False) + + # Prime build with a dummy pass + dummy = torch.zeros(4, 3) + torch_q(dummy, training=False) + keras_q(dummy, training=False) + + # Optimizer: plain SGD on f (and i for SAT) for both + torch_opt = torch.optim.SGD([p for p in torch_q.parameters() if p.requires_grad], lr=lr) + keras_f = _find_f_param(keras_q) + keras_i = _find_i_param(keras_q) + keras_params = [p for p in (keras_f, keras_i) if p is not None and p.requires_grad] + keras_opt = torch.optim.SGD(keras_params, lr=lr) + + for step in range(steps): + torch.manual_seed(100 + step) + x = torch.randn(4, 3) * 0.5 + + torch_opt.zero_grad() + loss_t = torch_q(x, training=True).pow(2).sum() + loss_t.backward() + torch_opt.step() + + keras_opt.zero_grad() + loss_k = _as_torch(keras_q(x, training=True)).pow(2).sum() + loss_k.backward() + keras_opt.step() + + # After training, the learned f should match within reasonable tolerance. + diff_f = (torch_q._f.detach() - keras_f.detach()).abs().max().item() + assert diff_f < 0.25, f"f parameters diverged after {steps} steps: max diff = {diff_f:.4f}" + + # Forward outputs on a held-out batch should also be close. + torch.manual_seed(999) + x_test = torch.randn(4, 3) * 0.5 + out_t = torch_q(x_test, training=False).detach() + out_k = _as_torch(keras_q(x_test, training=False)).detach() + diff_out = (out_t - out_k).abs().max().item() + assert diff_out < 0.1, f"trained outputs diverged: max diff = {diff_out:.4f}" + + +# --------------------------------------------------------------------------- +# Granularity → bit-width tensor shape +# +# per_tensor : one shared bit-width -> shape collapses to all-ones +# per_weight : one per element -> shape matches the tensor +# (data always shares the batch axis, so axis 0 collapses to 1) +# per_channel is not supported via granularity (see axis-override tests). +# --------------------------------------------------------------------------- + + +def _is_single_value(shape): + return int(torch.tensor(shape).prod().item()) == 1 + + +def _build_hgq(shape, is_data, granularity): + q = HGQQuantizer(k0=1, i0=2, f0=4, overflow_mode="SAT", round_mode="RND", is_data=is_data, granularity=granularity) + q.build(shape) + return tuple(q.i.shape), tuple(q.f.shape) + + +# weight layout (output-channel first): Linear (out, in), Conv1d (out, in, k), Conv2d (out, in, kH, kW) +WEIGHT_SHAPES = [(32, 16), (32, 16, 3), (32, 16, 3, 3)] +# data layout (channels-first): (batch, channels, *spatial) +DATA_SHAPES = [(4, 32), (4, 32, 6), (4, 32, 6, 6)] + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +@pytest.mark.parametrize("shape", WEIGHT_SHAPES) +def test_weight_granularity_shape(granularity, shape): + for bw_shape in _build_hgq(shape, is_data=False, granularity=granularity): + if granularity == "per_tensor": + assert _is_single_value(bw_shape), f"expected single value, got {bw_shape}" + else: # per_weight + assert bw_shape == shape, f"expected {shape}, got {bw_shape}" + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +@pytest.mark.parametrize("shape", DATA_SHAPES) +def test_data_granularity_shape(granularity, shape): + for bw_shape in _build_hgq(shape, is_data=True, granularity=granularity): + if granularity == "per_tensor": + assert _is_single_value(bw_shape), f"expected single value, got {bw_shape}" + else: # per_weight: batch axis shared, rest per-element + assert bw_shape == (1,) + shape[1:], f"expected batch-collapsed, got {bw_shape}" + + +@pytest.mark.parametrize("is_data", [False, True]) +def test_per_channel_rejected_for_hgq(is_data): + # per_channel is not a valid HGQ granularity; building must raise. + q = HGQQuantizer(k0=1, i0=2, f0=4, overflow_mode="SAT", round_mode="RND", is_data=is_data, granularity="per_channel") + with pytest.raises(ValueError, match="per_channel"): + q.build((32, 16, 3, 3)) + + +# --------------------------------------------------------------------------- +# Per-quantizer granularity override +# --------------------------------------------------------------------------- + + +def test_per_quantizer_granularity_override(): + from pquant.core.torch.layers import PQDense + + # input per_tensor, everything else per_weight + layer = PQDense( + _hgq_config("per_weight"), + in_features=16, + out_features=32, + quantize_output=True, + in_quant_granularity="per_tensor", + ) + layer(torch.randn(4, 16)) + _, wi, _ = layer.get_weight_quantization_bits() + assert tuple(wi.shape) == (32, 16) # weight follows config per_weight + _, ii, _ = layer.get_input_quantization_bits() + assert _is_single_value(tuple(ii.shape)) # input overridden to per_tensor + _, oi, _ = layer.get_output_quantization_bits() + assert tuple(oi.shape) == (1, 32) # output keeps config per_weight (batch-collapsed) + + +def test_per_quantizer_granularity_defaults_to_config(): + from pquant.core.torch.layers import PQDense + + layer = PQDense(_hgq_config("per_tensor"), in_features=16, out_features=32, weight_quant_granularity="per_weight") + layer(torch.randn(4, 16)) + _, wi, _ = layer.get_weight_quantization_bits() + assert tuple(wi.shape) == (32, 16) # weight overridden to per_weight + _, ii, _ = layer.get_input_quantization_bits() + assert _is_single_value(tuple(ii.shape)) # input uses config per_tensor + + +# --------------------------------------------------------------------------- +# Per-quantizer granularity override on auxiliary (boundary) layers. +# A model can start/end with a BatchNorm / LayerNorm / AvgPool whose input/output +# quantizer is effectively the model's I/O quantizer, so that must be overridable. +# --------------------------------------------------------------------------- + + +def test_batchnorm_input_granularity_override(): + from pquant.core.torch.layers import PQBatchNorm1d + + # config per_weight, but this BN sits on the model boundary -> input per_tensor + layer = PQBatchNorm1d(_hgq_config("per_weight"), num_features=16, in_quant_granularity="per_tensor") + layer(torch.randn(4, 16)) + _, ii, _ = layer.get_input_quantization_bits() + assert _is_single_value(tuple(ii.shape)) # input overridden to per_tensor + _, wi, _ = layer.get_weight_quantization_bits() + assert tuple(wi.shape) == (16,) # weight follows config per_weight + + +def test_layernorm_io_granularity_override(): + from pquant.core.torch.layers import PQLayerNorm + + layer = PQLayerNorm( + _hgq_config("per_weight"), + normalized_shape=16, + quantize_output=True, + in_quant_granularity="per_tensor", + out_quant_granularity="per_tensor", + ) + layer(torch.randn(4, 16)) + _, ii, _ = layer.get_input_quantization_bits() + assert _is_single_value(tuple(ii.shape)) # input overridden to per_tensor + _, oi, _ = layer.get_output_quantization_bits() + assert _is_single_value(tuple(oi.shape)) # output overridden to per_tensor + _, wi, _ = layer.get_weight_quantization_bits() + assert tuple(wi.shape) == (16,) # weight follows config per_weight + + +def test_avgpool_io_granularity_override(): + from pquant.core.torch.layers import PQAvgPool1d + + layer = PQAvgPool1d( + _hgq_config("per_weight"), + kernel_size=2, + quantize_output=True, + in_quant_granularity="per_tensor", + out_quant_granularity="per_tensor", + ) + layer(torch.randn(4, 16, 8)) + _, ii, _ = layer.get_input_quantization_bits() + assert _is_single_value(tuple(ii.shape)) # input overridden to per_tensor + _, oi, _ = layer.get_output_quantization_bits() + assert _is_single_value(tuple(oi.shape)) # output overridden to per_tensor + + +def test_mha_io_param_granularity_override(): + from pquant.core.torch.layers import PQMultiheadAttention + + # in/out are the model-boundary granularities; param (weight+bias) is uniform across projections. + layer = PQMultiheadAttention( + _hgq_config("per_weight"), + embed_dim=16, + num_heads=4, + batch_first=True, + quantize_output=True, + in_quant_granularity="per_tensor", + out_quant_granularity="per_tensor", + ) + x = torch.randn(2, 5, 16) + layer(x, x, x) + # Q/K/V projection inputs (boundary) overridden to per_tensor + for proj in (layer.q_proj, layer.k_proj, layer.v_proj): + _, ii, _ = proj.get_input_quantization_bits() + assert _is_single_value(tuple(ii.shape)) + # out_proj output (boundary) overridden to per_tensor + _, oi, _ = layer.out_proj.get_output_quantization_bits() + assert _is_single_value(tuple(oi.shape)) + # weights follow param granularity = config per_weight + _, wi, _ = layer.q_proj.get_weight_quantization_bits() + assert tuple(wi.shape) == (16, 16) + # out_proj input is internal -> stays config per_weight (batch-collapsed, not single) + _, oii, _ = layer.out_proj.get_input_quantization_bits() + assert not _is_single_value(tuple(oii.shape)) + + +def test_mha_qkv_always_output_quantized(): + from pquant.core.torch.layers import PQMultiheadAttention + + # MHA-level quantize_output=False (default): Q/K/V outputs are matmul operands and stay + # output-quantized; only out_proj follows the MHA-level flag. + layer = PQMultiheadAttention(_hgq_config("per_tensor"), embed_dim=16, num_heads=4, batch_first=True) + assert layer.q_proj.quantize_output is True + assert layer.k_proj.quantize_output is True + assert layer.v_proj.quantize_output is True + assert layer.out_proj.quantize_output is False + + +def test_mha_param_granularity_override(): + from pquant.core.torch.layers import PQMultiheadAttention + + # param overrides weight+bias across all projections; inputs/outputs follow config per_tensor. + layer = PQMultiheadAttention( + _hgq_config("per_tensor"), + embed_dim=16, + num_heads=4, + batch_first=True, + param_quant_granularity="per_weight", + ) + x = torch.randn(2, 5, 16) + layer(x, x, x) + _, wi, _ = layer.out_proj.get_weight_quantization_bits() + assert tuple(wi.shape) == (16, 16) # weight overridden to per_weight + _, ii, _ = layer.q_proj.get_input_quantization_bits() + assert _is_single_value(tuple(ii.shape)) # input uses config per_tensor + + +# --------------------------------------------------------------------------- +# Granularity end-to-end through the PQ layers +# --------------------------------------------------------------------------- + + +def _hgq_config(granularity): + from pquant import pdp_config + + config = pdp_config() + config.quantization_parameters.use_high_granularity_quantization = True + config.quantization_parameters.enable_quantization = True + config.quantization_parameters.granularity = granularity + config.pruning_parameters.enable_pruning = False + return config + + +@pytest.mark.parametrize("granularity", GRANULARITIES) +def test_layer_weight_granularity(granularity): + from pquant.core.torch.layers import PQConv1d, PQConv2d, PQDense + + cases = [ + (PQDense(_hgq_config(granularity), in_features=16, out_features=32), torch.randn(4, 16), (32, 16)), + (PQConv1d(_hgq_config(granularity), 16, 32, kernel_size=3), torch.randn(4, 16, 8), (32, 16, 3)), + (PQConv2d(_hgq_config(granularity), 16, 32, kernel_size=3), torch.randn(4, 16, 8, 8), (32, 16, 3, 3)), + ] + for layer, x, weight_shape in cases: + layer(x) # triggers lazy build of the weight quantizer + _, i, f = layer.get_weight_quantization_bits() + for bw_shape in (tuple(i.shape), tuple(f.shape)): + if granularity == "per_tensor": + assert _is_single_value(bw_shape), f"{type(layer).__name__}: expected single value, got {bw_shape}" + else: # per_weight + assert bw_shape == weight_shape, f"{type(layer).__name__}: expected {weight_shape}, got {bw_shape}" + + +# --------------------------------------------------------------------------- +# Utility: locate f / i parameters inside the Keras hgq2 quantizer +# --------------------------------------------------------------------------- + + +def _find_f_param(keras_q): + """Locate the fractional-bit parameter on the Keras hgq2 quantizer. + + hgq2 stores it as `keras_q.quantizer._f` (Keras Variable, which the torch + backend exposes as a torch.nn.Parameter). + """ + inner = getattr(keras_q, "quantizer", keras_q) + for name in ("_f", "f"): + p = getattr(inner, name, None) + if p is None: + continue + # Keras Variables expose a torch `.value` parameter under the torch backend + val = getattr(p, "value", p) + if isinstance(val, torch.nn.Parameter) or isinstance(val, torch.Tensor): + return val + raise RuntimeError("Could not locate f parameter on Keras hgq quantizer") + + +def _find_i_param(keras_q): + inner = getattr(keras_q, "quantizer", keras_q) + for name in ("_i", "i"): + p = getattr(inner, name, None) + if p is None: + continue + val = getattr(p, "value", p) + if isinstance(val, torch.nn.Parameter) or isinstance(val, torch.Tensor): + return val + return None diff --git a/tests/test_keras_alkaid_conversion.py b/tests/test_keras_alkaid_conversion.py new file mode 100644 index 0000000..8285169 --- /dev/null +++ b/tests/test_keras_alkaid_conversion.py @@ -0,0 +1,328 @@ +"""Convert a pruned + quantized PQuant Keras model with Alkaid""" + +import keras +import numpy as np +import pytest +from alkaid.codegen import RTLModel # noqa: E402 +from alkaid.converter import trace_model +from alkaid.trace import trace # noqa: E402 +from pquant.activations import PQActivation +from pquant.layers import ( + PQAvgPool1d, + PQAvgPool2d, + PQBatchNormalization, + PQConv1d, + PQConv2d, + PQDense, + PQDepthwiseConv2d, + PQMultiheadAttention, + PQSeparableConv2d, + PQSoftmax, + apply_final_compression, +) + +from pquant import pdp_config +from pquant._alkaid_plugin import _alkaid_keras_plugin # noqa: E402 +from pquant.core.keras.quantizer import Quantizer + +_alkaid_keras_plugin.register() + +IN_FEATURES = 3 +OUT_FEATURES = 4 +KERNEL_SIZE = 3 +H = W = 6 +SEQ_LEN = H * W + +PRUNE_FRACTION = 0.9 +INPUT_KIF = (1, 4, 4) + +IMG_SHAPE = (H, W, IN_FEATURES) +SEQ_SHAPE = (SEQ_LEN, IN_FEATURES) + + +@pytest.fixture(autouse=True) +def channels_last(): + # We assume Keras models are channels_last + keras.backend.set_image_data_format("channels_last") + + +def build_model(config): + img_in = keras.Input(shape=IMG_SHAPE, name="img") + a = PQConv2d(config, OUT_FEATURES, KERNEL_SIZE, padding="same")(img_in) + a = PQActivation(config, activation="relu", quantize_input=True, quantize_output=True)(a) + a = keras.layers.Flatten()(a) + + seq_in = keras.Input(shape=SEQ_SHAPE, name="seq") + b = PQConv1d(config, OUT_FEATURES, KERNEL_SIZE, padding="same")(seq_in) + b = PQActivation(config, activation="relu", quantize_input=True, quantize_output=True)(b) + b = keras.layers.Flatten()(b) + + x = keras.layers.Add()([a, b]) + x = PQDense(config, units=OUT_FEATURES)(x) + x = PQActivation(config, activation="relu", quantize_input=True, quantize_output=True)(x) + return keras.Model([img_in, seq_in], x) + + +def rtl_predict(comb, path, data): + """Write the RTL project, compile the simulation emulator, and run bit-accurate inference.""" + rtl_model = RTLModel(comb, str(path), "model", flavor="verilog", latency_cutoff=5, clock_period=5.0, print_latency=False) + rtl_model.write() + rtl_model.compile() + if isinstance(data, list): + data = [a.astype(np.float64) for a in data] + else: + data = data.astype(np.float64) + return rtl_model.predict(data) + + +def random_prune(layer, fraction, rng): + """Zero exactly ``fraction`` of the layer's weights via its pruning mask.""" + mask = layer.pruning_layer.mask + numel = int(np.prod(mask.shape)) + n_zero = int(round(fraction * numel)) + flat = np.ones(numel, dtype="float32") + flat[rng.permutation(numel)[:n_zero]] = 0.0 + mask.assign(flat.reshape(mask.shape)) + return n_zero / numel + + +def build_pruned_compressed_model(config): + model = build_model(config) + img = np.zeros((1,) + IMG_SHAPE, dtype="float32") + seq = np.zeros((1,) + SEQ_SHAPE, dtype="float32") + model([img, seq]) + apply_final_compression(model) + return model + + +def test_alkaid_conversion_pruned_quantized_model(): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + model = build_pruned_compressed_model(config) + inp, out = trace_model(model, inputs_kif=INPUT_KIF) + assert out.shape == (OUT_FEATURES,) + assert inp.shape == (int(np.prod(IMG_SHAPE)) + int(np.prod(SEQ_SHAPE)),) + + +def test_alkaid_rtl_matches_model(tmp_path): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + + model = build_pruned_compressed_model(config) + + inp_fv, out_fv = trace_model(model, inputs_kif=INPUT_KIF) + comb = trace(inp_fv, out_fv, optimize=True) + + n_samples = 16 + rng = np.random.default_rng(0) + img = rng.integers(0, 16, size=(n_samples,) + IMG_SHAPE).astype("float32") / 16.0 + seq = rng.integers(0, 16, size=(n_samples,) + SEQ_SHAPE).astype("float32") / 16.0 + + reference = np.asarray(model([img, seq]), dtype=np.float64) # (n_samples, OUT_FEATURES) + emulated = rtl_predict(comb, tmp_path, [img, seq]) + assert (tmp_path / "src" / "model.v").exists() + + assert np.any(reference != 0) # the comparison is non-trivial + np.testing.assert_allclose(emulated, reference, rtol=0, atol=1e-9) + + +ALL_C = 4 +ALL_H = ALL_W = 8 +ALL_LIN = (ALL_H // 2) * (ALL_W // 2) * 2 + +ALL_IMG_SHAPE = (ALL_H, ALL_W, IN_FEATURES) +ALL_SEQ_SHAPE = (ALL_LIN, IN_FEATURES) + +ALL_KERAS_LAYER_TYPES = { + "PQConv2d", + "PQBatchNormalization", + "PQDepthwiseConv2d", + "PQSeparableConv2d", + "PQAvgPool2d", + "PQConv1d", + "PQAvgPool1d", + "PQDense", + "PQActivation", +} + + +def build_all_layers_model(config): + """Model exercising every PQ layer type the keras Alkaid plugin handles.""" + img_in = keras.Input(shape=ALL_IMG_SHAPE, name="img") + a = PQConv2d(config, ALL_C, KERNEL_SIZE, padding="same")(img_in) + a = PQBatchNormalization(config, axis=-1)(a) + a = PQActivation(config, activation="relu", quantize_input=True, quantize_output=True)(a) + a = PQDepthwiseConv2d(config, KERNEL_SIZE, padding="same")(a) + a = PQSeparableConv2d(config, ALL_C, KERNEL_SIZE, padding="same")(a) + a = PQAvgPool2d(config, pool_size=2, strides=2, padding="valid")(a) + a = keras.layers.Flatten()(a) + + seq_in = keras.Input(shape=ALL_SEQ_SHAPE, name="seq") + b = PQConv1d(config, ALL_C, KERNEL_SIZE, padding="same")(seq_in) + b = PQActivation(config, activation="relu", quantize_input=True, quantize_output=True)(b) + b = PQAvgPool1d(config, pool_size=2, strides=2, padding="valid")(b) + b = keras.layers.Flatten()(b) + + x = keras.layers.Add()([a, b]) + x = PQDense(config, units=OUT_FEATURES)(x) + x = PQActivation(config, activation="relu", quantize_input=True, quantize_output=True)(x) + return keras.Model([img_in, seq_in], x) + + +def test_alkaid_conversion_all_layer_types(tmp_path): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + + model = build_all_layers_model(config) + rng = np.random.default_rng(0) + + model( + [ + rng.standard_normal((4,) + ALL_IMG_SHAPE).astype("float32"), + rng.standard_normal((4,) + ALL_SEQ_SHAPE).astype("float32"), + ] + ) + + for layer in model._flatten_layers(): + if not hasattr(layer, "pruning_layer"): + continue + layer._kernel.assign(rng.standard_normal(layer._kernel.shape).astype("float32")) + random_prune(layer, PRUNE_FRACTION, rng) + + apply_final_compression(model) + + inp_fv, out_fv = trace_model(model, inputs_kif=INPUT_KIF) + comb = trace(inp_fv, out_fv, optimize=True) + + n_samples = 16 + img = rng.integers(0, 16, size=(n_samples,) + ALL_IMG_SHAPE).astype("float32") / 16.0 + seq = rng.integers(0, 16, size=(n_samples,) + ALL_SEQ_SHAPE).astype("float32") / 16.0 + reference = np.asarray(model([img, seq]), dtype=np.float64) + emulated = rtl_predict(comb, tmp_path, [img, seq]) + assert (tmp_path / "src" / "model.v").exists() + + assert np.any(reference != 0) + np.testing.assert_allclose(emulated, reference, rtol=0, atol=1e-9) + + +def make_data_quantizer(config): + qp = config.quantization_parameters + return Quantizer( + k=qp.default_data_keep_negatives, + i=qp.default_data_integer_bits, + f=qp.default_data_fractional_bits, + overflow=qp.overflow_mode_data, + round_mode=qp.round_mode, + is_heterogeneous=qp.use_high_granularity_quantization, + is_data=True, + hgq_gamma=qp.hgq_gamma, + place="datalane", + dynamic_data=qp.dynamic_data_quantization, + ) + + +def create_single_layer_model(input_shape, layer, out_quantizer=False): + """A keras model that is one PQ layer, optionally followed by a data Quantizer.""" + inp = keras.Input(shape=input_shape) + x = layer(inp) + if out_quantizer: + x = make_data_quantizer(layer.config)(x) + return keras.Model(inp, x) + + +SINGLE_LAYER_CASES = { + "conv2d": lambda c: ( + (4, 4, 2), + create_single_layer_model((4, 4, 2), PQConv2d(c, 3, KERNEL_SIZE, padding="same", quantize_output=True)), + ), + "conv1d": lambda c: ( + (8, 2), + create_single_layer_model((8, 2), PQConv1d(c, 3, KERNEL_SIZE, padding="same", quantize_output=True)), + ), + "dense": lambda c: ((6,), create_single_layer_model((6,), PQDense(c, units=OUT_FEATURES, quantize_output=True))), + "depthwise2d": lambda c: ( + (4, 4, 3), + create_single_layer_model((4, 4, 3), PQDepthwiseConv2d(c, KERNEL_SIZE, padding="same", quantize_output=True)), + ), + "separable2d": lambda c: ( + (4, 4, 2), + create_single_layer_model((4, 4, 2), PQSeparableConv2d(c, 3, KERNEL_SIZE, padding="same", quantize_output=True)), + ), + "batchnorm": lambda c: ((6,), create_single_layer_model((6,), PQBatchNormalization(c, axis=-1), out_quantizer=True)), + "avgpool2d": lambda c: ( + (4, 4, 3), + create_single_layer_model((4, 4, 3), PQAvgPool2d(c, pool_size=2, strides=2, quantize_output=True)), + ), + "avgpool1d": lambda c: ( + (8, 3), + create_single_layer_model((8, 3), PQAvgPool1d(c, pool_size=2, strides=2, quantize_output=True)), + ), + "activation": lambda c: ( + (6,), + create_single_layer_model((6,), PQActivation(c, activation="relu", quantize_input=True, quantize_output=True)), + ), + "quantizer": lambda c: ((6,), create_single_layer_model((6,), make_data_quantizer(c))), + "softmax": lambda c: ((6,), create_single_layer_model((6,), PQSoftmax(c, axis=-1))), +} + + +@pytest.mark.parametrize("case_id", list(SINGLE_LAYER_CASES)) +def test_alkaid_single_layer(case_id, tmp_path): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + input_shape, model = SINGLE_LAYER_CASES[case_id](config) + rng = np.random.default_rng(0) + + model(rng.standard_normal((4,) + input_shape).astype("float32")) + apply_final_compression(model) + + inp_fv, out_fv = trace_model(model, inputs_kif=INPUT_KIF) + comb = trace(inp_fv, out_fv, optimize=True) + + n_samples = 16 + x = rng.integers(0, 16, size=(n_samples,) + input_shape).astype("float32") / 16.0 + reference = np.asarray(model(x), dtype=np.float64).reshape(n_samples, -1) + emulated = rtl_predict(comb, tmp_path, x) + + assert np.any(reference != 0) + np.testing.assert_allclose(emulated, reference, rtol=0, atol=1e-9) + + +MHA_SEQ_LEN = 4 +MHA_EMBED_DIM = 4 +MHA_NUM_HEADS = 2 + + +def build_mha_model(config, rng): + inp = keras.Input(shape=(MHA_SEQ_LEN, MHA_EMBED_DIM)) + out, _ = PQMultiheadAttention( + config, + embed_dim=MHA_EMBED_DIM, + num_heads=MHA_NUM_HEADS, + quantize_output=True, + )(inp) + model = keras.Model(inp, out) + model(rng.standard_normal((4, MHA_SEQ_LEN, MHA_EMBED_DIM)).astype("float32")) + apply_final_compression(model) + return model + + +def test_alkaid_multihead_attention(tmp_path): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + + rng = np.random.default_rng(0) + model = build_mha_model(config, rng) + + inp_fv, out_fv = trace_model(model, inputs_kif=INPUT_KIF) + comb = trace(inp_fv, out_fv, optimize=True) + assert out_fv.shape == (MHA_SEQ_LEN * MHA_EMBED_DIM,) + + n_samples = 16 + x = rng.integers(0, 16, size=(n_samples, MHA_SEQ_LEN, MHA_EMBED_DIM)).astype("float32") / 16.0 + reference = np.asarray(model(x), dtype=np.float64).reshape(n_samples, -1) + emulated = rtl_predict(comb, tmp_path, x) + assert (tmp_path / "src" / "model.v").exists() + + assert np.any(reference != 0) + np.testing.assert_allclose(emulated, reference, rtol=0, atol=1e-9) diff --git a/tests/test_keras_compression_layers.py b/tests/test_keras_compression_layers.py index 2cee1b0..6df0073 100644 --- a/tests/test_keras_compression_layers.py +++ b/tests/test_keras_compression_layers.py @@ -15,18 +15,7 @@ ReLU, SeparableConv2D, ) - -from pquant import ( - ap_config, - autosparse_config, - cs_config, - dst_config, - mdmm_config, - pdp_config, - wanda_config, -) from pquant.activations import PQActivation -from pquant.core.hyperparameter_optimization import PQConfig from pquant.layers import ( PQAvgPool1d, PQAvgPool2d, @@ -35,6 +24,7 @@ PQConv2d, PQDense, PQDepthwiseConv2d, + PQMultiheadAttention, PQSeparableConv2d, add_compression_layers, apply_final_compression, @@ -43,6 +33,17 @@ pre_finetune_functions, ) +from pquant import ( + ap_config, + autosparse_config, + cs_config, + dst_config, + mdmm_config, + pdp_config, + wanda_config, +) +from pquant.core.hyperparameter_optimization import PQConfig + BATCH_SIZE = 4 OUT_FEATURES = 32 IN_FEATURES = 16 @@ -301,8 +302,7 @@ def test_separable_conv2d_call(config_pdp, conv2d_input): layer_to_replace.pointwise_constraint, layer_to_replace.bias_constraint, ) - layer.depthwise_conv.build(conv2d_input.shape) - layer.pointwise_conv.build(conv2d_input.shape) + layer.build(conv2d_input.shape) layer.depthwise_conv._kernel.assign(layer_to_replace.depthwise_kernel) layer.pointwise_conv._kernel.assign(layer_to_replace.pointwise_kernel) @@ -316,9 +316,15 @@ def test_separable_conv2d_add_remove_layers(config_pdp, conv2d_input): inputs = keras.Input(shape=conv2d_input.shape[1:]) out = SeparableConv2D(OUT_FEATURES, KERNEL_SIZE, use_bias=False, padding="same")(inputs) model = keras.Model(inputs=inputs, outputs=out, name="test_conv2d") + depthwise_kernel = ops.copy(model.layers[1].depthwise_kernel) + pointwise_kernel = ops.copy(model.layers[1].pointwise_kernel) model = add_compression_layers(model, config_pdp, conv2d_input.shape) model(conv2d_input) + # The original layer's weights must survive the replacement. + assert ops.all(ops.equal(model.layers[1].depthwise_conv._kernel, depthwise_kernel)) + assert ops.all(ops.equal(model.layers[1].pointwise_conv._kernel, pointwise_kernel)) + post_pretrain_functions(model, config_pdp) pre_finetune_functions(model) @@ -1387,6 +1393,8 @@ def test_trigger_post_pretraining(config_pdp, conv2d_input): def test_hgq_weight_shape(config_pdp, dense_input): config_pdp.quantization_parameters.enable_quantization = True config_pdp.quantization_parameters.use_high_granularity_quantization = True + # Per-weight granularity → one bit-width per kernel element. + config_pdp.quantization_parameters.granularity = "per_weight" inputs = keras.Input(shape=dense_input.shape[1:]) out = Dense(OUT_FEATURES, use_bias=False)(inputs) act1 = Activation("tanh")(out) @@ -1466,13 +1474,13 @@ def test_set_activation_custom_bits_hgq(config_pdp, conv2d_input): assert ops.all(f_input == 7.0) config_pdp.quantization_parameters.layer_specific = { - 'conv2d': { - 'weight': {'integer_bits': 1.0, 'fractional_bits': 3.0}, - 'bias': {'integer_bits': 2.0, 'fractional_bits': 4.0}, + "conv2d": { + "weight": {"integer_bits": 1.0, "fractional_bits": 3.0}, + "bias": {"integer_bits": 2.0, "fractional_bits": 4.0}, }, - 're_lu': {"input": {'integer_bits': 1.0, 'fractional_bits': 3.0}}, - 'average_pooling2d': {"input": {'integer_bits': 1.0, 'fractional_bits': 3.0}}, - 'activation': {"input": {'integer_bits': 0.0, 'fractional_bits': 3.0}}, + "re_lu": {"input": {"integer_bits": 1.0, "fractional_bits": 3.0}}, + "average_pooling2d": {"input": {"integer_bits": 1.0, "fractional_bits": 3.0}}, + "activation": {"input": {"integer_bits": 0.0, "fractional_bits": 3.0}}, } keras.backend.clear_session() inputs = keras.Input(shape=conv2d_input.shape[1:]) @@ -1533,13 +1541,13 @@ def test_set_activation_custom_bits_quantizer(config_pdp, conv2d_input): assert m.f_input == 7.0 config_pdp.quantization_parameters.layer_specific = { - 'conv2d': { - 'weight': {'integer_bits': 1.0, 'fractional_bits': 3.0}, - 'bias': {'integer_bits': 2.0, 'fractional_bits': 4.0}, + "conv2d": { + "weight": {"integer_bits": 1.0, "fractional_bits": 3.0}, + "bias": {"integer_bits": 2.0, "fractional_bits": 4.0}, }, - 're_lu': {"input": {'integer_bits': 1.0, 'fractional_bits': 3.0}}, - 'average_pooling2d': {"input": {'integer_bits': 1.0, 'fractional_bits': 3.0}}, - 'activation': {"input": {'integer_bits': 0.0, 'fractional_bits': 3.0}}, + "re_lu": {"input": {"integer_bits": 1.0, "fractional_bits": 3.0}}, + "average_pooling2d": {"input": {"integer_bits": 1.0, "fractional_bits": 3.0}}, + "activation": {"input": {"integer_bits": 0.0, "fractional_bits": 3.0}}, } keras.backend.clear_session() inputs = keras.Input(shape=conv2d_input.shape[1:]) @@ -1707,7 +1715,6 @@ def test_avg_pool1d(config_pdp, conv1d_input): class DummyLayer(keras.layers.Layer): - def __init__(self, *args, **kwargs): super().__init__() self.built = True @@ -1732,7 +1739,7 @@ def extra_repr(self): def test_avgpool_quant_called(config_pdp, conv1d_input): config_pdp.quantization_parameters.enable_quantization = True - with patch('pquant.layers.Quantizer', DummyLayer): + with patch("pquant.layers.Quantizer", DummyLayer): layer = PQAvgPool1d(config_pdp, KERNEL_SIZE, quantize_input=True) layer(conv1d_input) assert layer.input_quantizer.layer_called == 1 @@ -1754,7 +1761,7 @@ def test_avgpool_quant_called(config_pdp, conv1d_input): def test_batchnorm_quant_called(config_pdp, conv2d_input): config_pdp.quantization_parameters.enable_quantization = True axis = -1 if keras.backend.image_data_format() == "channels_last" else 1 - with patch('pquant.layers.Quantizer', DummyLayer): + with patch("pquant.layers.Quantizer", DummyLayer): layer = PQBatchNormalization(config_pdp, axis=axis, quantize_input=True) layer(conv2d_input) assert layer.input_quantizer.layer_called == 1 @@ -1778,7 +1785,7 @@ def test_batchnorm_quant_called(config_pdp, conv2d_input): def test_pqconv2d_quant_called(config_pdp, conv2d_input): config_pdp.quantization_parameters.enable_quantization = True - with patch('pquant.layers.Quantizer', DummyLayer): + with patch("pquant.layers.Quantizer", DummyLayer): layer = PQConv2d(config_pdp, OUT_FEATURES, KERNEL_SIZE, quantize_input=True, use_bias=True) layer.post_pre_train_function() layer(conv2d_input) @@ -1809,7 +1816,7 @@ def test_pqconv2d_quant_called(config_pdp, conv2d_input): def test_pqdepthwiseconv2d_quant_called(config_pdp, conv2d_input): config_pdp.quantization_parameters.enable_quantization = True - with patch('pquant.layers.Quantizer', DummyLayer): + with patch("pquant.layers.Quantizer", DummyLayer): layer = PQDepthwiseConv2d(config_pdp, KERNEL_SIZE, quantize_input=True, use_bias=True) layer.post_pre_train_function() layer(conv2d_input) @@ -1839,7 +1846,7 @@ def test_pqdepthwiseconv2d_quant_called(config_pdp, conv2d_input): def test_pqconv1d_quant_called(config_pdp, conv1d_input): config_pdp.quantization_parameters.enable_quantization = True - with patch('pquant.layers.Quantizer', DummyLayer): + with patch("pquant.layers.Quantizer", DummyLayer): layer = PQConv1d(config_pdp, OUT_FEATURES, KERNEL_SIZE, quantize_input=True, use_bias=True) layer.post_pre_train_function() layer(conv1d_input) @@ -1869,7 +1876,7 @@ def test_pqconv1d_quant_called(config_pdp, conv1d_input): def test_dense_quant_called(config_pdp, dense_input): config_pdp.quantization_parameters.enable_quantization = True - with patch('pquant.layers.Quantizer', DummyLayer): + with patch("pquant.layers.Quantizer", DummyLayer): layer = PQDense(config_pdp, OUT_FEATURES, quantize_input=True, use_bias=True) layer.post_pre_train_function() layer(dense_input) @@ -1902,7 +1909,7 @@ def test_layer_replacement_quant_called(config_pdp, conv2d_input): config_pdp.quantization_parameters.quantize_input = True config_pdp.quantization_parameters.quantize_output = True config_pdp.quantization_parameters.use_high_granularity_quantization = True - with patch('pquant.layers.Quantizer', DummyLayer): + with patch("pquant.layers.Quantizer", DummyLayer): inp = keras.Input(shape=conv2d_input.shape[1:]) x = Conv2D(OUT_FEATURES, KERNEL_SIZE)(inp) @@ -2134,3 +2141,149 @@ def test_model_fit(config_fn): model.compile(optimizer="adam", loss="mse", jit_compile=False) callback = PQuantCallback(config, log_ebops=False, log_keep_ratio=False) model.fit(dummy_x, dummy_y, epochs=callback.total_epochs, callbacks=[callback], verbose=0) + + +def build_single_layer_model(layer_type, config): + channels_first = keras.backend.image_data_format() == "channels_first" + conv2d_shape = (IN_FEATURES, 8, 8) if channels_first else (8, 8, IN_FEATURES) + conv1d_shape = (IN_FEATURES, STEPS) if channels_first else (STEPS, IN_FEATURES) + bn_axis = 1 if channels_first else -1 + + if layer_type == "dense": + inputs = keras.Input(shape=(IN_FEATURES,)) + outputs = PQDense(config, units=OUT_FEATURES)(inputs) + elif layer_type == "conv1d": + inputs = keras.Input(shape=conv1d_shape) + outputs = PQConv1d(config, OUT_FEATURES, KERNEL_SIZE, padding="same")(inputs) + elif layer_type == "conv2d": + inputs = keras.Input(shape=conv2d_shape) + outputs = PQConv2d(config, OUT_FEATURES, KERNEL_SIZE, padding="same")(inputs) + elif layer_type == "depthwise_conv2d": + inputs = keras.Input(shape=conv2d_shape) + outputs = PQDepthwiseConv2d(config, KERNEL_SIZE, padding="same")(inputs) + elif layer_type == "separable_conv2d": + inputs = keras.Input(shape=conv2d_shape) + outputs = PQSeparableConv2d(config, OUT_FEATURES, KERNEL_SIZE, padding="same")(inputs) + elif layer_type == "batchnorm": + inputs = keras.Input(shape=conv2d_shape) + outputs = PQBatchNormalization(config, axis=bn_axis)(inputs) + elif layer_type == "avgpool1d": + inputs = keras.Input(shape=conv1d_shape) + outputs = PQAvgPool1d(config, pool_size=2, strides=2, padding="same")(inputs) + elif layer_type == "avgpool2d": + inputs = keras.Input(shape=conv2d_shape) + outputs = PQAvgPool2d(config, pool_size=2, strides=2, padding="same")(inputs) + elif layer_type in ("activation_relu", "activation_tanh", "activation_hard_tanh"): + activation = layer_type.replace("activation_", "") + inputs = keras.Input(shape=(IN_FEATURES,)) + outputs = PQActivation(config, activation=activation, quantize_input=True, quantize_output=True)(inputs) + elif layer_type == "mha": + inputs = keras.Input(shape=(STEPS, IN_FEATURES)) + outputs = PQMultiheadAttention(config, embed_dim=IN_FEATURES, num_heads=2)(inputs)[0] + else: + raise ValueError(f"unknown layer type {layer_type}") + + model = keras.Model(inputs, outputs) + dummy = np.zeros((1,) + tuple(inputs.shape[1:]), dtype=np.float32) + return model, dummy + + +def randomize_weights(model, seed): + rng = np.random.default_rng(seed) + for w in model.weights: + # Lifecycle flags must keep their 0/1 values. + if w.name.endswith(("is_pretraining", "is_finetuning")): + continue + w.assign(rng.standard_normal(w.shape).astype(w.dtype)) + + +def assert_weights_equal(model, reloaded, label): + assert len(model.weights) == len(reloaded.weights), f"[{label}] weight count mismatch" + for orig_w, loaded_w in zip(model.weights, reloaded.weights): + np.testing.assert_array_equal( + np.array(orig_w), + np.array(loaded_w), + err_msg=f"[{label}] Weight mismatch: {orig_w.name}", + ) + + +LAYER_TYPES = [ + "dense", + "conv1d", + "conv2d", + "depthwise_conv2d", + "separable_conv2d", + "batchnorm", + "avgpool1d", + "avgpool2d", + "activation_relu", + "activation_tanh", + "activation_hard_tanh", + "mha", +] + + +@pytest.mark.parametrize("use_hgq", [False, True], ids=["kif", "hgq"]) +@pytest.mark.parametrize("layer_type", LAYER_TYPES) +def test_layer_type_serialization(tmp_path, layer_type, use_hgq): + """Every PQ layer type must survive a full .keras save/load round-trip. + + Covers all quantizer variables (k/i/f/b, HGQ bitwidth variables) and pruning + state carried as model weights, at each lifecycle stage. + """ + # dst: no global pruning setup, so single-layer models without a prunable + # weight layer (batchnorm/pool/activation/mha) pass the lifecycle functions. + config = dst_config() + config.quantization_parameters.enable_quantization = True + config.quantization_parameters.use_high_granularity_quantization = use_hgq + + model, dummy = build_single_layer_model(layer_type, config) + model(dummy) + randomize_weights(model, seed=42) + + def roundtrip(label): + save_path = tmp_path / f"{label}.keras" + model.save(save_path) + reloaded = keras.models.load_model(save_path) + assert_weights_equal(model, reloaded, f"{layer_type}/{label}") + # The reloaded model must be usable, not just structurally equal. + np.testing.assert_allclose( + np.array(model(dummy)), + np.array(reloaded(dummy)), + rtol=0, + atol=0, + err_msg=f"[{layer_type}/{label}] forward pass mismatch after reload", + ) + + roundtrip("initial") + post_pretrain_functions(model, config) + roundtrip("post_pretrain") + pre_finetune_functions(model) + roundtrip("pre_finetune") + apply_final_compression(model) + roundtrip("final_compression") + + +@pytest.mark.parametrize("use_hgq", [False, True], ids=["kif", "hgq"]) +@pytest.mark.parametrize("layer_type", LAYER_TYPES) +def test_layer_type_checkpoint_save_load(tmp_path, layer_type, use_hgq): + """Every PQ layer type must survive a save_weights/load_weights round-trip.""" + # dst: no global pruning setup, so single-layer models without a prunable + # weight layer (batchnorm/pool/activation/mha) pass the lifecycle functions. + config = dst_config() + config.quantization_parameters.enable_quantization = True + config.quantization_parameters.use_high_granularity_quantization = use_hgq + + model, dummy = build_single_layer_model(layer_type, config) + model(dummy) + randomize_weights(model, seed=0) + original_weights = [np.array(w) for w in model.weights] + + path = str(tmp_path / "ckpt.weights.h5") + model.save_weights(path) + for w in model.weights: + w.assign(np.zeros(w.shape, dtype=w.dtype)) + model.load_weights(path) + + for orig, w in zip(original_weights, model.weights): + np.testing.assert_array_equal(orig, np.array(w), err_msg=f"[{layer_type}] Checkpoint weight mismatch: {w.name}") diff --git a/tests/test_keras_onnx_converter.py b/tests/test_keras_onnx_converter.py new file mode 100644 index 0000000..91dc46f --- /dev/null +++ b/tests/test_keras_onnx_converter.py @@ -0,0 +1,197 @@ +"""Tests for the Keras → ONNX converter (convert_to_onnx). + +Each test builds a small functional Keras model, runs a forward pass to +initialise all sublayer state, calls apply_final_compression, exports to ONNX +via convert_to_onnx(), and verifies that onnxruntime produces the same output +as the Keras model. + +bias=True/False is tested via parametrize where applicable. +""" + +import keras +import numpy as np +import pytest + +import pquant +from pquant.core.keras.convert_to_onnx import convert_to_onnx +from pquant.core.keras.layers import ( + PQBatchNormalization, + PQConv1d, + PQConv2d, + PQDense, + PQDepthwiseConv2d, + apply_final_compression, +) + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime not installed") + +ATOL = 1e-4 + + +# --------------------------------------------------------------------------- +# fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def cfg(): + c = pquant.cs_config() + c.quantization_parameters.enable_quantization = False + return c + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _channels_first(): + return keras.backend.image_data_format() == "channels_first" + + +def _keras_out(model, x: np.ndarray) -> np.ndarray: + from keras import ops + + return ops.convert_to_numpy(model(x, training=False)) + + +def _onnx_run(model, x: np.ndarray, input_shape: tuple, tmp_path) -> np.ndarray: + path = str(tmp_path / "model.onnx") + convert_to_onnx(model, input_shape=input_shape, output_path=path) + sess = ort.InferenceSession(path) + in_name = sess.get_inputs()[0].name + return sess.run(None, {in_name: x})[0] + + +# --------------------------------------------------------------------------- +# PQDense +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bias", [True, False]) +def test_dense_onnx(cfg, bias, tmp_path): + IN, OUT = 16, 8 + inputs = keras.Input(shape=(IN,)) + x = PQDense(cfg, units=OUT, use_bias=bias)(inputs) + model = keras.Model(inputs, x) + + dummy = np.zeros((1, IN), dtype=np.float32) + model(dummy) + apply_final_compression(model) + + x_np = np.random.randn(4, IN).astype(np.float32) + keras_out = _keras_out(model, x_np) + onnx_out = _onnx_run(model, x_np, input_shape=(IN,), tmp_path=tmp_path) + np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg=f"PQDense bias={bias}: keras vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQConv2d +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bias", [True, False]) +def test_conv2d_onnx(cfg, bias, tmp_path): + IN_C, OUT_C, H, W = 3, 8, 8, 8 + if _channels_first(): + input_shape = (IN_C, H, W) + x_np = np.random.randn(2, IN_C, H, W).astype(np.float32) + else: + input_shape = (H, W, IN_C) + x_np = np.random.randn(2, H, W, IN_C).astype(np.float32) + + inputs = keras.Input(shape=input_shape) + x = PQConv2d(cfg, OUT_C, kernel_size=3, padding="same", use_bias=bias)(inputs) + model = keras.Model(inputs, x) + + dummy = np.zeros((1, *input_shape), dtype=np.float32) + model(dummy) + apply_final_compression(model) + + keras_out = _keras_out(model, x_np) + onnx_out = _onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg=f"PQConv2d bias={bias}: keras vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQConv1d +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bias", [True, False]) +def test_conv1d_onnx(cfg, bias, tmp_path): + IN_C, OUT_C, L = 4, 8, 16 + if _channels_first(): + input_shape = (IN_C, L) + x_np = np.random.randn(2, IN_C, L).astype(np.float32) + else: + input_shape = (L, IN_C) + x_np = np.random.randn(2, L, IN_C).astype(np.float32) + + inputs = keras.Input(shape=input_shape) + x = PQConv1d(cfg, OUT_C, kernel_size=3, padding="same", use_bias=bias)(inputs) + model = keras.Model(inputs, x) + + dummy = np.zeros((1, *input_shape), dtype=np.float32) + model(dummy) + apply_final_compression(model) + + keras_out = _keras_out(model, x_np) + onnx_out = _onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg=f"PQConv1d bias={bias}: keras vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQBatchNormalization +# --------------------------------------------------------------------------- + + +def test_batchnorm_onnx(cfg, tmp_path): + IN_C, H, W = 8, 4, 4 + if _channels_first(): + input_shape = (IN_C, H, W) + x_np = np.random.randn(4, IN_C, H, W).astype(np.float32) + bn_axis = 1 + else: + input_shape = (H, W, IN_C) + x_np = np.random.randn(4, H, W, IN_C).astype(np.float32) + bn_axis = -1 + + inputs = keras.Input(shape=input_shape) + x = PQBatchNormalization(cfg, axis=bn_axis)(inputs) + model = keras.Model(inputs, x) + + dummy = np.zeros((1, *input_shape), dtype=np.float32) + model(dummy, training=True) # warm up running stats + apply_final_compression(model) + + keras_out = _keras_out(model, x_np) + onnx_out = _onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg="PQBatchNormalization: keras vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQDepthwiseConv2d +# --------------------------------------------------------------------------- + + +def test_depthwise_conv2d_onnx(cfg, tmp_path): + IN_C, H, W = 4, 8, 8 + if _channels_first(): + input_shape = (IN_C, H, W) + x_np = np.random.randn(2, IN_C, H, W).astype(np.float32) + else: + input_shape = (H, W, IN_C) + x_np = np.random.randn(2, H, W, IN_C).astype(np.float32) + + inputs = keras.Input(shape=input_shape) + x = PQDepthwiseConv2d(cfg, kernel_size=3, padding="same")(inputs) + model = keras.Model(inputs, x) + + dummy = np.zeros((1, *input_shape), dtype=np.float32) + model(dummy) + apply_final_compression(model) + + keras_out = _keras_out(model, x_np) + onnx_out = _onnx_run(model, x_np, input_shape=input_shape, tmp_path=tmp_path) + np.testing.assert_allclose(keras_out, onnx_out, atol=ATOL, err_msg="PQDepthwiseConv2d: keras vs ONNX mismatch") diff --git a/tests/test_pdp.py b/tests/test_pdp.py index 4a6a081..3a1ec1c 100644 --- a/tests/test_pdp.py +++ b/tests/test_pdp.py @@ -2,7 +2,7 @@ from keras import ops from keras.random import shuffle -from pquant.pruning_methods.pdp import PDP +from pquant.core.keras.pruning_methods.pdp import PDP @pytest.fixture diff --git a/tests/test_quantizer_parity.py b/tests/test_quantizer_parity.py new file mode 100644 index 0000000..baf879f --- /dev/null +++ b/tests/test_quantizer_parity.py @@ -0,0 +1,357 @@ +""" +Parity tests: PyTorch `Quantizer` vs Keras `Quantizer` in non-HGQ (fixed-point) mode. + +Both implementations must produce matching forward outputs, matching gradients +w.r.t. the quantized tensor, and must freeze their dynamically-computed +bitwidths identically once training ends (``apply_final_compression``). + +``per_channel`` granularity needs special handling: keras stores weights +output-channel-last while torch stores them output-channel-first, so the raw +``Quantizer`` reduces over different axes in each backend. A real layer +reconciles this with a transpose before quantizing (keras ``_handle_transpose``); +here we do the same transpose manually (see ``to_channel_last``/``to_channel_first``) +so both backends reduce over the same channel groups before comparing. +""" + +import os + +os.environ.setdefault("KERAS_BACKEND", "tensorflow") + +import numpy as np # noqa: E402 +import pytest # noqa: E402 +import tensorflow as tf # noqa: E402 +import torch # noqa: E402 +from keras import ops # noqa: E402 + +from pquant.core.keras.quantizer import Quantizer as KQuantizer # noqa: E402 +from pquant.core.torch.quantizer import Quantizer as TQuantizer # noqa: E402 + +ABSOLUTE_TOLERANCE = 1e-5 +RELATIVE_TOLERANCE = 1e-4 + + +def to_numpy(x): + if isinstance(x, torch.Tensor): + return x.detach().cpu().numpy() + if isinstance(x, np.ndarray): + return x + return np.asarray(ops.convert_to_numpy(x)) + + +def to_channel_last(x): + """Move torch's channel-first axis (0) to keras's channel-last axis (-1).""" + return np.moveaxis(x, 0, -1) + + +def to_channel_first(x): + """Move keras's channel-last axis (-1) back to torch's channel-first axis (0).""" + return np.moveaxis(to_numpy(x), -1, 0) + + +def assert_close(a, b, atol=ABSOLUTE_TOLERANCE, rtol=RELATIVE_TOLERANCE, msg=""): + a_np = to_numpy(a) + b_np = to_numpy(b) + assert a_np.shape == b_np.shape, f"{msg}: shape mismatch: {a_np.shape} vs {b_np.shape}" + np.testing.assert_allclose(a_np, b_np, atol=atol, rtol=rtol, err_msg=msg) + + +def keras_tensor(arr): + return ops.convert_to_tensor(np.asarray(arr).astype(np.float32)) + + +def torch_tensor(arr, requires_grad=False): + return torch.as_tensor(np.asarray(arr).astype(np.float32)).requires_grad_(requires_grad) + + +def reset_seed(seed=0): + np.random.seed(seed) + torch.manual_seed(seed) + + +def keras_grad(fn, x, training=True): + """Gradient of scalar-reducing ``fn(x).sum()`` w.r.t. keras tensor ``x``.""" + xt = tf.convert_to_tensor(x) + with tf.GradientTape() as tape: + tape.watch(xt) + loss = ops.sum(fn(xt, training=training)) + return tape.gradient(loss, xt) + + +def make_quantizers(shape, k=1.0, i=2.0, f=2.0, overflow="SAT", round_mode="RND", is_data=False, granularity="per_tensor"): + k_layer = KQuantizer( + k=k, + i=i, + f=f, + overflow=overflow, + round_mode=round_mode, + is_heterogeneous=False, + is_data=is_data, + granularity=granularity, + ) + k_layer.build(shape) + + t_layer = TQuantizer( + k=k, + i=i, + f=f, + overflow=overflow, + round_mode=round_mode, + is_heterogeneous=False, + is_data=is_data, + granularity=granularity, + ) + return k_layer, t_layer + + +@pytest.mark.parametrize( + "shape,granularity,is_data", + [ + ((8, 4), "per_tensor", False), + ((8, 4), "per_weight", False), + ((4, 8, 3, 3), "per_weight", False), + ((16,), "per_tensor", True), + ], +) +def test_quantizer_matches_keras(shape, granularity, is_data): + reset_seed() + x_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + + k_layer, t_layer = make_quantizers(shape, is_data=is_data, granularity=granularity) + + k_x = keras_tensor(x_np) + t_x = torch_tensor(x_np, requires_grad=True) + k_out = k_layer(k_x, training=True) + t_out = t_layer(t_x) + assert_close(k_out, t_out, msg=f"Quantizer forward ({granularity}, is_data={is_data})") + assert_close(k_layer.i, t_layer.i, msg=f"Quantizer i ({granularity}, is_data={is_data})") + assert_close(k_layer.f, t_layer.f, msg=f"Quantizer f ({granularity}, is_data={is_data})") + + k_grad = keras_grad(k_layer, x_np, training=True) + t_out.sum().backward() + assert_close(k_grad, t_x.grad, msg=f"Quantizer backward ({granularity}, is_data={is_data})") + + +@pytest.mark.parametrize("shape", [(8, 4), (4, 8, 3, 3)]) +def test_quantizer_per_channel_matches_keras(shape): + reset_seed() + x_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + x_np_channel_last = to_channel_last(x_np) + + k_layer, t_layer = make_quantizers(x_np_channel_last.shape, is_data=False, granularity="per_channel") + + k_x = keras_tensor(x_np_channel_last) + t_x = torch_tensor(x_np, requires_grad=True) + k_out = k_layer(k_x, training=True) + t_out = t_layer(t_x) + assert_close(to_channel_first(k_out), t_out, msg=f"Quantizer per_channel forward {shape}") + assert_close(to_channel_first(k_layer.i), t_layer.i, msg=f"Quantizer per_channel i {shape}") + assert_close(to_channel_first(k_layer.f), t_layer.f, msg=f"Quantizer per_channel f {shape}") + + k_grad = keras_grad(k_layer, x_np_channel_last, training=True) + t_out.sum().backward() + assert_close(to_channel_first(k_grad), t_x.grad, msg=f"Quantizer per_channel backward {shape}") + + +@pytest.mark.parametrize("shape", [(8, 4), (4, 8, 3, 3)]) +def test_quantizer_per_channel_freezes_bits_after_final_compression(shape): + reset_seed() + train_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + eval_np = (np.random.randn(*shape) * 500.0).astype(np.float32) + train_np_channel_last = to_channel_last(train_np) + eval_np_channel_last = to_channel_last(eval_np) + + k_layer, t_layer = make_quantizers(train_np_channel_last.shape, is_data=False, granularity="per_channel") + + k_layer(keras_tensor(train_np_channel_last), training=True) + t_layer.train() + t_layer(torch_tensor(train_np)) + + k_layer.apply_final_compression() + t_layer.apply_final_compression() + assert_close(to_channel_first(k_layer.i), t_layer.i, msg=f"per_channel i after apply_final_compression {shape}") + assert_close(to_channel_first(k_layer.f), t_layer.f, msg=f"per_channel f after apply_final_compression {shape}") + + frozen_k_i, frozen_k_f = to_numpy(k_layer.i).copy(), to_numpy(k_layer.f).copy() + + k_layer(keras_tensor(eval_np_channel_last), training=False) + t_layer.eval() + t_layer(torch_tensor(eval_np)) + + assert_close(k_layer.i, frozen_k_i, msg=f"keras per_channel i drifted after eval forward {shape}") + assert_close(k_layer.f, frozen_k_f, msg=f"keras per_channel f drifted after eval forward {shape}") + assert_close(t_layer.i, to_channel_first(frozen_k_i), msg=f"torch per_channel i drifted after eval forward {shape}") + assert_close(t_layer.f, to_channel_first(frozen_k_f), msg=f"torch per_channel f drifted after eval forward {shape}") + + +@pytest.mark.parametrize("granularity", ["per_tensor", "per_channel", "per_weight"]) +def test_is_data_ignores_granularity(granularity): + reset_seed() + shape = (8, 4) + x_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + + k_layer, t_layer = make_quantizers(shape, is_data=True, granularity=granularity) + + k_out = k_layer(keras_tensor(x_np), training=True) + t_out = t_layer(torch_tensor(x_np)) + + assert to_numpy(k_layer.i).size == 1, f"keras is_data i should stay scalar for granularity={granularity}" + assert to_numpy(t_layer.i).size == 1, f"torch is_data i should stay scalar for granularity={granularity}" + assert_close(k_out, t_out, msg=f"Quantizer forward (is_data=True, granularity={granularity})") + + +@pytest.mark.parametrize("overflow", ["SAT", "SAT_SYM", "WRAP", "WRAP_SM"]) +def test_quantizer_matches_keras_overflow_modes(overflow): + reset_seed() + shape = (8, 4) + x_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + + k_layer, t_layer = make_quantizers(shape, overflow=overflow, is_data=False, granularity="per_tensor") + + k_x = keras_tensor(x_np) + t_x = torch_tensor(x_np, requires_grad=True) + k_out = k_layer(k_x, training=True) + t_out = t_layer(t_x) + assert_close(k_out, t_out, msg=f"Quantizer forward (overflow={overflow})") + + k_grad = keras_grad(k_layer, x_np, training=True) + t_out.sum().backward() + assert_close(k_grad, t_x.grad, msg=f"Quantizer backward (overflow={overflow})") + + +def test_quantizer_matches_keras_wrap_eval_saturates(): + """WRAP overflow only saturates when training=False -- exercise that branch specifically.""" + reset_seed() + shape = (8, 4) + x_np = (np.random.randn(*shape) * 50.0).astype(np.float32) # large enough to trigger saturation + + k_layer, t_layer = make_quantizers(shape, overflow="WRAP", is_data=False, granularity="per_tensor") + + k_out = k_layer(keras_tensor(x_np), training=False) + t_layer.eval() + t_out = t_layer(torch_tensor(x_np)) + assert_close(k_out, t_out, msg="Quantizer WRAP eval-mode saturate branch") + + +@pytest.mark.parametrize( + "round_mode", + ["RND", "TRN", "RND_CONV", "TRN_ZERO", "RND_ZERO", "RND_MIN_INF", "RND_INF"], +) +def test_quantizer_matches_keras_round_modes(round_mode): + reset_seed() + shape = (8, 4) + x_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + + k_layer, t_layer = make_quantizers(shape, round_mode=round_mode, is_data=False, granularity="per_tensor") + + k_x = keras_tensor(x_np) + t_x = torch_tensor(x_np, requires_grad=True) + k_out = k_layer(k_x, training=True) + t_out = t_layer(t_x) + assert_close(k_out, t_out, msg=f"Quantizer forward (round_mode={round_mode})") + + k_grad = keras_grad(k_layer, x_np, training=True) + t_out.sum().backward() + assert_close(k_grad, t_x.grad, msg=f"Quantizer backward (round_mode={round_mode})") + + +@pytest.mark.parametrize("k", [0.0, 1.0]) +def test_quantizer_matches_keras_k_values(k): + reset_seed() + shape = (8, 4) + x_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + if k == 0.0: + x_np = np.abs(x_np) + + k_layer, t_layer = make_quantizers(shape, k=k, is_data=False, granularity="per_tensor") + + k_x = keras_tensor(x_np) + t_x = torch_tensor(x_np, requires_grad=True) + k_out = k_layer(k_x, training=True) + t_out = t_layer(t_x) + assert_close(k_out, t_out, msg=f"Quantizer forward (k={k})") + + k_grad = keras_grad(k_layer, x_np, training=True) + t_out.sum().backward() + assert_close(k_grad, t_x.grad, msg=f"Quantizer backward (k={k})") + + +def test_get_total_bits_matches_keras_non_hgq(): + reset_seed() + shape = (8, 4) + x_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + + k_layer, t_layer = make_quantizers(shape, is_data=False, granularity="per_tensor") + k_layer(keras_tensor(x_np), training=True) + t_layer(torch_tensor(x_np)) + + assert_close(k_layer.get_total_bits(shape), t_layer.get_total_bits(shape), msg="get_total_bits (non-HGQ)") + + +def test_get_total_bits_matches_keras_hgq(): + reset_seed() + shape = (8, 4) + x_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + + k_layer = KQuantizer( + k=1.0, + i=2.0, + f=2.0, + overflow="SAT", + round_mode="RND", + is_heterogeneous=True, + is_data=False, + granularity="per_tensor", + place="weight", + ) + k_layer.build(shape) + t_layer = TQuantizer( + k=1.0, + i=2.0, + f=2.0, + overflow="SAT", + round_mode="RND", + is_heterogeneous=True, + is_data=False, + granularity="per_tensor", + place="weight", + ) + + k_layer(keras_tensor(x_np), training=True) + t_layer(torch_tensor(x_np)) + + assert_close(k_layer.get_total_bits(shape), t_layer.get_total_bits(shape), msg="get_total_bits (HGQ)") + + +@pytest.mark.parametrize( + "shape,granularity", + [ + ((8, 4), "per_weight"), + ((4, 8, 3, 3), "per_weight"), + ], +) +def test_quantizer_freezes_bits_after_final_compression(shape, granularity): + reset_seed() + train_np = (np.random.randn(*shape) * 2.0).astype(np.float32) + eval_np = (np.random.randn(*shape) * 500.0).astype(np.float32) + + k_layer, t_layer = make_quantizers(shape, is_data=False, granularity=granularity) + + k_layer(keras_tensor(train_np), training=True) + t_layer.train() + t_layer(torch_tensor(train_np)) + + k_layer.apply_final_compression() + t_layer.apply_final_compression() + assert_close(k_layer.i, t_layer.i, msg=f"i after apply_final_compression ({granularity})") + assert_close(k_layer.f, t_layer.f, msg=f"f after apply_final_compression ({granularity})") + + frozen_k_i, frozen_k_f = to_numpy(k_layer.i).copy(), to_numpy(k_layer.f).copy() + + k_layer(keras_tensor(eval_np), training=False) + t_layer.eval() + t_layer(torch_tensor(eval_np)) + + assert_close(k_layer.i, frozen_k_i, msg=f"keras i drifted after eval forward ({granularity})") + assert_close(k_layer.f, frozen_k_f, msg=f"keras f drifted after eval forward ({granularity})") + assert_close(t_layer.i, frozen_k_i, msg=f"torch i drifted after eval forward ({granularity})") + assert_close(t_layer.f, frozen_k_f, msg=f"torch f drifted after eval forward ({granularity})") diff --git a/tests/test_torch_alkaid_conversion.py b/tests/test_torch_alkaid_conversion.py new file mode 100644 index 0000000..c9097a2 --- /dev/null +++ b/tests/test_torch_alkaid_conversion.py @@ -0,0 +1,343 @@ +"""Convert a pruned + quantized PQuant Torch model with Alkaid""" + +import numpy as np +import pytest +import torch +import torch.nn as nn +from alkaid.codegen import RTLModel # noqa: E402 +from alkaid.converter import trace_model +from alkaid.trace import ( # noqa: E402 + FVArray, + HWConfig, + trace, # noqa: E402 +) + +from pquant import pdp_config +from pquant._alkaid_plugin import _alkaid_torch_plugin # noqa: E402 +from pquant.core.torch.activations import PQActivation +from pquant.core.torch.layers import ( + PQAvgPool1d, + PQAvgPool2d, + PQBatchNorm1d, + PQBatchNorm2d, + PQConv1d, + PQConv2d, + PQDense, + PQMultiheadAttention, + PQSoftmax, + apply_final_compression, +) +from pquant.core.torch.quantizer import Quantizer + +_alkaid_torch_plugin.register() + +IN_FEATURES = 3 +OUT_FEATURES = 4 +KERNEL_SIZE = 3 +H = W = 6 +SEQ_LEN = H * W + +PRUNE_FRACTION = 0.9 +HWCONF = HWConfig(1, 1, -1) +INPUT_KIF = (1, 4, 4) + + +class TwoBranchNet(nn.Module): + def __init__(self, config): + super().__init__() + self.conv2d = PQConv2d(config, IN_FEATURES, OUT_FEATURES, KERNEL_SIZE, padding="same") + self.act2d = PQActivation(config, "relu", quantize_input=True, quantize_output=True) + self.flat2d = nn.Flatten() + + self.conv1d = PQConv1d(config, IN_FEATURES, OUT_FEATURES, KERNEL_SIZE, padding="same") + self.act1d = PQActivation(config, "relu", quantize_input=True, quantize_output=True) + self.flat1d = nn.Flatten() + + self.dense = PQDense(config, OUT_FEATURES * SEQ_LEN, OUT_FEATURES) + self.act = PQActivation(config, "relu", quantize_input=True, quantize_output=True) + + def forward(self, img, seq): + a = self.flat2d(self.act2d(self.conv2d(img))) + b = self.flat1d(self.act1d(self.conv1d(seq))) + x = a + b + return self.act(self.dense(x)) + + +def random_prune(layer, fraction, rng): + mask = layer.pruning_layer.mask + numel = int(np.prod(tuple(mask.shape))) + n_zero = int(round(fraction * numel)) + flat = np.ones(numel, dtype="float32") + flat[rng.permutation(numel)[:n_zero]] = 0.0 + mask.copy_(torch.tensor(flat.reshape(tuple(mask.shape)), dtype=mask.dtype, device=mask.device)) + return n_zero / numel + + +def fixed_point_input(shape, kif=INPUT_KIF): + """Bounded fixed-point symbolic input so the SAT input quantizer can be replayed.""" + k, i, f = (np.full(shape, v, dtype=np.int8) for v in kif) + return FVArray.from_kif(k, i, f, HWCONF, 0, None) + + +def rtl_predict(comb, path, data): + rtl_model = RTLModel(comb, str(path), "model", flavor="verilog", latency_cutoff=5, clock_period=5.0, print_latency=False) + rtl_model.write() + rtl_model.compile() + if isinstance(data, list): + data = [a.astype(np.float64) for a in data] + else: + data = data.astype(np.float64) + return rtl_model.predict(data) + + +def build_pruned_compressed_model(config): + model = TwoBranchNet(config) + device = next(model.parameters()).device + img = torch.zeros(1, IN_FEATURES, H, W, device=device) + seq = torch.zeros(1, IN_FEATURES, SEQ_LEN, device=device) + with torch.no_grad(): + model(img, seq) + apply_final_compression(model) + model.eval() + return model, device + + +def test_alkaid_conversion_pruned_quantized_model(): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + + model, _ = build_pruned_compressed_model(config) + + inputs = (fixed_point_input((1, IN_FEATURES, H, W)), fixed_point_input((1, IN_FEATURES, SEQ_LEN))) + inp, out = trace_model(model, hwconf=HWCONF, inputs=inputs, framework="torch") + + assert out.shape == (OUT_FEATURES,) + expected_inputs = IN_FEATURES * H * W + IN_FEATURES * SEQ_LEN + assert inp.shape == (expected_inputs,) + + +def test_alkaid_rtl_matches_model(tmp_path): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + + rng = np.random.default_rng(0) + model, device = build_pruned_compressed_model(config) + + inputs = (fixed_point_input((1, IN_FEATURES, H, W)), fixed_point_input((1, IN_FEATURES, SEQ_LEN))) + inp_fv, out_fv = trace_model(model, hwconf=HWCONF, inputs=inputs, framework="torch") + comb = trace(inp_fv, out_fv, optimize=True) + n_samples = 16 + img = rng.integers(0, 16, size=(n_samples, IN_FEATURES, H, W)).astype("float32") / 16.0 + seq = rng.integers(0, 16, size=(n_samples, IN_FEATURES, SEQ_LEN)).astype("float32") / 16.0 + + with torch.no_grad(): + reference = ( + model(torch.tensor(img, device=device), torch.tensor(seq, device=device)).cpu().numpy().astype(np.float64) + ) # (n_samples, OUT_FEATURES) + + emulated = rtl_predict(comb, tmp_path, [img, seq]) + assert (tmp_path / "src" / "model.v").exists() + + assert np.any(reference != 0) # the comparison is non-trivial + np.testing.assert_allclose(emulated, reference, rtol=0, atol=1e-9) + + +ALL_C = 4 +ALL_H = ALL_W = 6 +ALL_LIN = (ALL_H // 2) * (ALL_W // 2) * 2 + +ALL_TORCH_LAYER_TYPES = { + "PQConv2d", + "PQBatchNorm2d", + "PQAvgPool2d", + "PQConv1d", + "PQBatchNorm1d", + "PQAvgPool1d", + "PQDense", + "PQActivation", +} + + +class AllLayersNet(nn.Module): + def __init__(self, config): + super().__init__() + self.conv2d = PQConv2d(config, IN_FEATURES, ALL_C, KERNEL_SIZE, padding="same") + self.bn2d = PQBatchNorm2d(config, ALL_C) + self.act2d = PQActivation(config, "relu", quantize_input=True, quantize_output=True) + self.pool2d = PQAvgPool2d(config, kernel_size=2, stride=2) + self.flat2d = nn.Flatten() + + self.conv1d = PQConv1d(config, IN_FEATURES, ALL_C, KERNEL_SIZE, padding="same") + self.bn1d = PQBatchNorm1d(config, ALL_C) + self.act1d = PQActivation(config, "relu", quantize_input=True, quantize_output=True) + self.pool1d = PQAvgPool1d(config, kernel_size=2, stride=2) + self.flat1d = nn.Flatten() + + self.dense = PQDense(config, ALL_C * (ALL_H // 2) * (ALL_W // 2), OUT_FEATURES) + self.act = PQActivation(config, "relu", quantize_input=True, quantize_output=True) + + def forward(self, img, seq): + a = self.flat2d(self.pool2d(self.act2d(self.bn2d(self.conv2d(img))))) + b = self.flat1d(self.pool1d(self.act1d(self.bn1d(self.conv1d(seq))))) + return self.act(self.dense(a + b)) + + +def test_alkaid_conversion_all_layer_types(tmp_path): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + + rng = np.random.default_rng(0) + model = AllLayersNet(config) + device = next(model.parameters()).device + + model.train() + with torch.no_grad(): + model( + torch.tensor(rng.standard_normal((4, IN_FEATURES, ALL_H, ALL_W)), dtype=torch.float32, device=device), + torch.tensor(rng.standard_normal((4, IN_FEATURES, ALL_LIN)), dtype=torch.float32, device=device), + ) + + with torch.no_grad(): + for module in model.modules(): + if not hasattr(module, "pruning_layer"): + continue + random_prune(module, PRUNE_FRACTION, rng) + apply_final_compression(model) + model.eval() + + inputs = (fixed_point_input((1, IN_FEATURES, ALL_H, ALL_W)), fixed_point_input((1, IN_FEATURES, ALL_LIN))) + inp_fv, out_fv = trace_model(model, hwconf=HWCONF, inputs=inputs, framework="torch") + comb = trace(inp_fv, out_fv, optimize=True) + assert out_fv.shape == (OUT_FEATURES,) + + n_samples = 16 + img = rng.integers(0, 16, size=(n_samples, IN_FEATURES, ALL_H, ALL_W)).astype("float32") / 16.0 + seq = rng.integers(0, 16, size=(n_samples, IN_FEATURES, ALL_LIN)).astype("float32") / 16.0 + with torch.no_grad(): + reference = ( + model(torch.tensor(img, device=device), torch.tensor(seq, device=device)).cpu().numpy().astype(np.float64) + ) + emulated = rtl_predict(comb, tmp_path, [img, seq]) + assert np.any(reference != 0) + np.testing.assert_allclose(emulated, reference, rtol=0, atol=1e-9) + + +def make_data_quantizer(config): + qp = config.quantization_parameters + return Quantizer( + k=qp.default_data_keep_negatives, + i=qp.default_data_integer_bits, + f=qp.default_data_fractional_bits, + overflow=qp.overflow_mode_data, + round_mode=qp.round_mode, + is_heterogeneous=qp.use_high_granularity_quantization, + is_data=True, + hgq_gamma=qp.hgq_gamma, + place="datalane", + dynamic_data=qp.dynamic_data_quantization, + ) + + +class SingleLayer(nn.Module): + def __init__(self, layer, tail=None): + super().__init__() + self.layer = layer + self.tail = tail + + def forward(self, x): + x = self.layer(x) + return x if self.tail is None else self.tail(x) + + +SINGLE_LAYER_CASES = { + "conv2d": lambda c: ((1, 2, 4, 4), SingleLayer(PQConv2d(c, 2, 3, KERNEL_SIZE, padding="same", quantize_output=True))), + "conv1d": lambda c: ((1, 2, 8), SingleLayer(PQConv1d(c, 2, 3, KERNEL_SIZE, padding="same", quantize_output=True))), + "dense": lambda c: ((1, 6), SingleLayer(PQDense(c, 6, OUT_FEATURES, quantize_output=True))), + "batchnorm2d": lambda c: ((1, 3, 4, 4), SingleLayer(PQBatchNorm2d(c, 3), make_data_quantizer(c))), + "batchnorm1d": lambda c: ((1, 3, 8), SingleLayer(PQBatchNorm1d(c, 3), make_data_quantizer(c))), + "avgpool2d": lambda c: ((1, 3, 4, 4), SingleLayer(PQAvgPool2d(c, kernel_size=2, stride=2, quantize_output=True))), + "avgpool1d": lambda c: ((1, 3, 8), SingleLayer(PQAvgPool1d(c, kernel_size=2, stride=2, quantize_output=True))), + "activation": lambda c: ((1, 6), SingleLayer(PQActivation(c, "relu", quantize_input=True, quantize_output=True))), + "quantizer": lambda c: ((1, 6), SingleLayer(make_data_quantizer(c))), + "softmax": lambda c: ((1, 6), SingleLayer(PQSoftmax(c, axis=-1))), +} + + +@pytest.mark.parametrize("case_id", list(SINGLE_LAYER_CASES)) +def test_alkaid_single_layer(case_id, tmp_path): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + shape, model = SINGLE_LAYER_CASES[case_id](config) + rng = np.random.default_rng(0) + + model.train() + with torch.no_grad(): + model(torch.tensor(rng.standard_normal((4,) + shape[1:]), dtype=torch.float32)) + apply_final_compression(model) + model.eval() + + inp_fv, out_fv = trace_model(model, hwconf=HWCONF, inputs=(fixed_point_input(shape),), framework="torch") + comb = trace(inp_fv, out_fv, optimize=True) + + n_samples = 16 + x = rng.integers(0, 16, size=(n_samples,) + shape[1:]).astype("float32") / 16.0 + with torch.no_grad(): + reference = model(torch.tensor(x)).cpu().numpy().reshape(n_samples, -1).astype(np.float64) + emulated = rtl_predict(comb, tmp_path, x) + + assert np.any(reference != 0) + np.testing.assert_allclose(emulated, reference, rtol=0, atol=1e-9) + + +MHA_SEQ_LEN = 4 +MHA_EMBED_DIM = 4 +MHA_NUM_HEADS = 2 + + +class MHANet(nn.Module): + """Self-attention PQMultiheadAttention with every data quantizer enabled. + + The MHA lives inside a wrapper module so torch.fx inlines its forward with + concrete (None) mask arguments; tracing the MHA as the fx root would turn the + masks into proxies and hit data-dependent control flow. + """ + + def __init__(self, config): + super().__init__() + self.mha = PQMultiheadAttention( + config, + embed_dim=MHA_EMBED_DIM, + num_heads=MHA_NUM_HEADS, + batch_first=True, + quantize_output=True, + ) + + def forward(self, x): + out, _ = self.mha(x, x, x) + return out + + +def test_alkaid_multihead_attention(tmp_path): + config = pdp_config() + config.quantization_parameters.enable_quantization = True + + rng = np.random.default_rng(0) + model = MHANet(config) + with torch.no_grad(): + model(torch.tensor(rng.standard_normal((4, MHA_SEQ_LEN, MHA_EMBED_DIM)), dtype=torch.float32)) # build + apply_final_compression(model) + model.eval() + + shape = (1, MHA_SEQ_LEN, MHA_EMBED_DIM) + inp_fv, out_fv = trace_model(model, hwconf=HWCONF, inputs=(fixed_point_input(shape),), framework="torch") + comb = trace(inp_fv, out_fv, optimize=True) + assert out_fv.shape == (MHA_SEQ_LEN * MHA_EMBED_DIM,) + + n_samples = 16 + x = rng.integers(0, 16, size=(n_samples, MHA_SEQ_LEN, MHA_EMBED_DIM)).astype("float32") / 16.0 + with torch.no_grad(): + reference = model(torch.tensor(x)).cpu().numpy().reshape(n_samples, -1).astype(np.float64) + emulated = rtl_predict(comb, tmp_path, x) + + assert np.any(reference != 0) + np.testing.assert_allclose(emulated, reference, rtol=0, atol=1e-9) diff --git a/tests/test_torch_checkpoint.py b/tests/test_torch_checkpoint.py new file mode 100644 index 0000000..29624a7 --- /dev/null +++ b/tests/test_torch_checkpoint.py @@ -0,0 +1,171 @@ +import os + +import numpy as np +import pytest +import torch +from torch import nn + +os.environ["KERAS_BACKEND"] = "torch" + +from pquant.activations import PQActivation # noqa: E402 +from pquant.layers import ( # noqa: E402 + PQAvgPool1d, + PQAvgPool2d, + PQBatchNorm1d, + PQBatchNorm2d, + PQConv1d, + PQConv2d, + PQDense, + PQLayerNorm, + PQMultiheadAttention, + apply_final_compression, + post_pretrain_functions, + pre_finetune_functions, +) + +from pquant import dst_config # noqa: E402 + +BATCH_SIZE = 2 +OUT_FEATURES = 8 +IN_FEATURES = 4 +KERNEL_SIZE = 3 +STEPS = 6 + + +class SingleLayerModel(nn.Module): + def __init__(self, layer, is_mha=False): + super().__init__() + self.layer = layer + self.is_mha = is_mha + + def forward(self, x): + if self.is_mha: + x, _ = self.layer(x, x, x) + else: + x = self.layer(x) + return x + + +def build_model_and_input(layer_type, config): + if layer_type == "dense": + layer = PQDense(config, IN_FEATURES, OUT_FEATURES) + x = torch.randn(BATCH_SIZE, IN_FEATURES) + elif layer_type == "conv1d": + layer = PQConv1d(config, IN_FEATURES, OUT_FEATURES, KERNEL_SIZE, padding=1) + x = torch.randn(BATCH_SIZE, IN_FEATURES, STEPS) + elif layer_type == "conv2d": + layer = PQConv2d(config, IN_FEATURES, OUT_FEATURES, KERNEL_SIZE, padding=1) + x = torch.randn(BATCH_SIZE, IN_FEATURES, STEPS, STEPS) + elif layer_type == "batchnorm1d": + layer = PQBatchNorm1d(config, IN_FEATURES) + x = torch.randn(BATCH_SIZE, IN_FEATURES, STEPS) + elif layer_type == "batchnorm2d": + layer = PQBatchNorm2d(config, IN_FEATURES) + x = torch.randn(BATCH_SIZE, IN_FEATURES, STEPS, STEPS) + elif layer_type == "layernorm": + layer = PQLayerNorm(config, IN_FEATURES) + x = torch.randn(BATCH_SIZE, STEPS, IN_FEATURES) + elif layer_type == "avgpool1d": + layer = PQAvgPool1d(config, kernel_size=2) + x = torch.randn(BATCH_SIZE, IN_FEATURES, STEPS) + elif layer_type == "avgpool2d": + layer = PQAvgPool2d(config, kernel_size=2) + x = torch.randn(BATCH_SIZE, IN_FEATURES, STEPS, STEPS) + elif layer_type.startswith("activation_"): + layer = PQActivation(config, activation=layer_type.replace("activation_", ""), quantize_output=True) + x = torch.randn(BATCH_SIZE, IN_FEATURES) + elif layer_type == "mha": + layer = PQMultiheadAttention(config, embed_dim=IN_FEATURES, num_heads=2) + x = torch.randn(STEPS, BATCH_SIZE, IN_FEATURES) + return SingleLayerModel(layer, is_mha=True), x + else: + raise ValueError(f"unknown layer kind {layer_type}") + + return SingleLayerModel(layer), x + + +STAGE_FLAGS = ("is_pretraining", "is_finetuning", "final_compression_done") + + +def randomize_state(model, seed): + gen = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for name, t in model.state_dict().items(): + if any(k in name for k in STAGE_FLAGS): + continue + if not torch.is_floating_point(t): + continue + t.copy_(torch.rand(t.shape, generator=gen, device="cpu") + 0.5) + + +def advance_to_stage(model, config, stage): + if stage == "initial": + return + post_pretrain_functions(model, config) + if stage == "post_pretrain": + return + pre_finetune_functions(model) + if stage == "pre_finetune": + return + apply_final_compression(model) + assert stage == "final_compression" + + +def make_config(use_hgq): + config = dst_config() + config.quantization_parameters.enable_quantization = True + config.quantization_parameters.use_high_granularity_quantization = use_hgq + return config + + +LAYER_TYPES = [ + "dense", + "conv1d", + "conv2d", + "batchnorm1d", + "batchnorm2d", + "layernorm", + "avgpool1d", + "avgpool2d", + "activation_relu", + "activation_tanh", + "activation_hard_tanh", + "mha", +] +STAGES = ["initial", "post_pretrain", "pre_finetune", "final_compression"] + + +@pytest.mark.parametrize("use_hgq", [False, True], ids=["kif", "hgq"]) +@pytest.mark.parametrize("stage", STAGES) +@pytest.mark.parametrize("layer_type", LAYER_TYPES) +def test_state_dict_roundtrip(tmp_path, layer_type, stage, use_hgq): + torch.manual_seed(0) + config = make_config(use_hgq) + + model, x = build_model_and_input(layer_type, config) + model(x) # HGQ quantizers build lazily on first forward + advance_to_stage(model, config, stage) + randomize_state(model, seed=42) + + path = tmp_path / "ckpt.pt" + torch.save(model.state_dict(), path) + torch.manual_seed(1) + fresh_config = make_config(use_hgq) + fresh, _ = build_model_and_input(layer_type, fresh_config) + fresh(x) + advance_to_stage(fresh, fresh_config, stage) + missing, unexpected = fresh.load_state_dict(torch.load(path, weights_only=True), strict=True) + assert not missing and not unexpected + + saved = model.state_dict() + reloaded = fresh.state_dict() + assert saved.keys() == reloaded.keys() + for name in saved: + np.testing.assert_array_equal( + reloaded[name].cpu(), saved[name].cpu(), err_msg=f"state mismatch: {name}", strict=True + ) + + model.eval() + fresh.eval() + with torch.no_grad(): + np.testing.assert_array_equal(fresh(x).cpu(), model(x).cpu(), strict=True) diff --git a/tests/test_torch_compression_layers.py b/tests/test_torch_compression_layers.py index cb4d1ff..778ef73 100644 --- a/tests/test_torch_compression_layers.py +++ b/tests/test_torch_compression_layers.py @@ -1,8 +1,8 @@ -import keras +import os + import numpy as np import pytest import torch -from keras import ops from torch import nn from torch.nn import ( AvgPool2d, @@ -14,10 +14,12 @@ Tanh, ) -from pquant import post_training_prune -from pquant.activations import PQActivation -from pquant.core.hyperparameter_optimization import PQConfig -from pquant.layers import ( +os.environ["KERAS_BACKEND"] = "torch" + +import keras # noqa: E402 +from keras import ops # noqa: E402 +from pquant.activations import PQActivation # noqa: E402 +from pquant.layers import ( # noqa: E402 PQAvgPool1d, PQAvgPool2d, PQBatchNorm2d, @@ -33,6 +35,9 @@ pre_finetune_functions, ) +from pquant import post_training_prune # noqa: E402 +from pquant.core.hyperparameter_optimization import PQConfig # noqa: E402 + BATCH_SIZE = 4 OUT_FEATURES = 32 IN_FEATURES = 16 @@ -599,6 +604,8 @@ def _to_bool(val): def test_hgq_weight_shape(config_pdp, dense_input): config_pdp.quantization_parameters.enable_quantization = True config_pdp.quantization_parameters.use_high_granularity_quantization = True + # Per-weight granularity → one bit-width per weight element. + config_pdp.quantization_parameters.granularity = "per_weight" layer = Linear(IN_FEATURES, OUT_FEATURES, bias=False) layer2 = Linear(OUT_FEATURES, OUT_FEATURES, bias=False) model = TestModel2(layer, layer2, "relu", "tanh") @@ -606,20 +613,22 @@ def test_hgq_weight_shape(config_pdp, dense_input): model = add_compression_layers(model, config_pdp, dense_input.shape) post_pretrain_functions(model, config_pdp) - assert model.submodule.weight_quantizer.quantizer.quantizer._i.shape == model.submodule.weight.shape - assert model.activation.input_quantizer.quantizer.quantizer._i.shape == (1, OUT_FEATURES) + assert model.submodule.weight_quantizer.quantizer._i.shape == model.submodule.weight.shape + assert model.activation.input_quantizer.quantizer._i.shape == (1, OUT_FEATURES) def test_qbn_build(config_pdp, conv2d_input): config_pdp.quantization_parameters.enable_quantization = True config_pdp.quantization_parameters.use_high_granularity_quantization = True + # Per-weight granularity → one bit-width per conv-kernel element. + config_pdp.quantization_parameters.granularity = "per_weight" layer = Conv2d(IN_FEATURES, OUT_FEATURES, KERNEL_SIZE, bias=False) layer2 = BatchNorm2d(OUT_FEATURES) model = TestModel2(layer, layer2, None, "tanh") model = add_compression_layers(model, config_pdp, conv2d_input.shape) post_pretrain_functions(model, config_pdp) - assert model.submodule.weight_quantizer.quantizer.quantizer._i.shape == model.submodule.weight.shape + assert model.submodule.weight_quantizer.quantizer._i.shape == model.submodule.weight.shape def test_set_activation_custom_bits_hgq(config_pdp, conv2d_input): @@ -634,13 +643,13 @@ def test_set_activation_custom_bits_hgq(config_pdp, conv2d_input): if isinstance(m, (PQWeightBiasBase)): assert m.i_weight == 0.0 assert m.i_bias == 0.0 - assert torch.all(m.weight_quantizer.quantizer.quantizer.i == 0.0) - assert torch.all(m.weight_quantizer.quantizer.quantizer.i == 0.0) + assert torch.all(m.weight_quantizer.quantizer.i == 0.0) + assert torch.all(m.weight_quantizer.quantizer.i == 0.0) assert m.f_weight == 7.0 assert m.f_bias == 7.0 - assert torch.all(m.weight_quantizer.quantizer.quantizer.f == 7.0) - assert torch.all(m.weight_quantizer.quantizer.quantizer.f == 7.0) + assert torch.all(m.weight_quantizer.quantizer.f == 7.0) + assert torch.all(m.weight_quantizer.quantizer.f == 7.0) elif isinstance(m, PQActivation) and m.activation_name == "tanh": k_input, i_input, f_input = m.get_input_quantization_bits() @@ -655,17 +664,17 @@ def test_set_activation_custom_bits_hgq(config_pdp, conv2d_input): elif isinstance(m, PQAvgPool2d): assert m.i_input == 0.0 assert m.f_input == 7.0 - assert torch.all(m.input_quantizer.quantizer.quantizer.i == 0.0) - assert torch.all(m.input_quantizer.quantizer.quantizer.f == 7.0) + assert torch.all(m.input_quantizer.quantizer.i == 0.0) + assert torch.all(m.input_quantizer.quantizer.f == 7.0) config_pdp.quantization_parameters.layer_specific = { - 'submodule': { - 'weight': {'integer_bits': 1, 'fractional_bits': 3}, - 'bias': {'integer_bits': 2, 'fractional_bits': 4}, + "submodule": { + "weight": {"integer_bits": 1, "fractional_bits": 3}, + "bias": {"integer_bits": 2, "fractional_bits": 4}, }, - 'submodule2': {"input": {'integer_bits': 1, 'fractional_bits': 3}}, - 'activation': {"input": {'integer_bits': 1, 'fractional_bits': 4}}, - 'activation2': {"input": {'integer_bits': 0, 'fractional_bits': 3}}, + "submodule2": {"input": {"integer_bits": 1, "fractional_bits": 3}}, + "activation": {"input": {"integer_bits": 1, "fractional_bits": 4}}, + "activation2": {"input": {"integer_bits": 0, "fractional_bits": 3}}, } model = TestModel2(layer, layer2, "relu", "tanh") @@ -675,13 +684,13 @@ def test_set_activation_custom_bits_hgq(config_pdp, conv2d_input): if isinstance(m, (PQWeightBiasBase)): assert m.i_weight == 1.0 assert m.i_bias == 2.0 - assert torch.all(m.weight_quantizer.quantizer.quantizer.i == 1.0) - assert torch.all(m.bias_quantizer.quantizer.quantizer.i == 2.0) + assert torch.all(m.weight_quantizer.quantizer.i == 1.0) + assert torch.all(m.bias_quantizer.quantizer.i == 2.0) assert m.f_weight == 3.0 assert m.f_bias == 4.0 - assert torch.all(m.weight_quantizer.quantizer.quantizer.f == 3.0) - assert torch.all(m.bias_quantizer.quantizer.quantizer.f == 4.0) + assert torch.all(m.weight_quantizer.quantizer.f == 3.0) + assert torch.all(m.bias_quantizer.quantizer.f == 4.0) elif isinstance(m, PQActivation) and m.activation_name == "tanh": k_input, i_input, f_input = m.get_input_quantization_bits() @@ -695,8 +704,8 @@ def test_set_activation_custom_bits_hgq(config_pdp, conv2d_input): elif isinstance(m, PQAvgPool2d): assert m.i_input == 1.0 assert m.f_input == 3.0 - assert torch.all(m.input_quantizer.quantizer.quantizer.i == 1.0) - assert torch.all(m.input_quantizer.quantizer.quantizer.f == 3.0) + assert torch.all(m.input_quantizer.quantizer.i == 1.0) + assert torch.all(m.input_quantizer.quantizer.f == 3.0) def test_disable_pruning_from_single_layer(config_pdp, conv2d_input): @@ -739,13 +748,13 @@ def test_set_activation_custom_bits_quantizer(config_pdp, conv2d_input): assert m.f_input == 8.0 config_pdp.quantization_parameters.layer_specific = { - 'submodule': { - 'weight': {'integer_bits': 1.0, 'fractional_bits': 3.0}, - 'bias': {'integer_bits': 1.0, 'fractional_bits': 3.0}, + "submodule": { + "weight": {"integer_bits": 1.0, "fractional_bits": 3.0}, + "bias": {"integer_bits": 1.0, "fractional_bits": 3.0}, }, - 'submodule2': {"input": {'integer_bits': 1.0, 'fractional_bits': 3.0}}, - 'activation': {"input": {'integer_bits': 0.0, 'fractional_bits': 4.0}}, - 'activation2': {"input": {'integer_bits': 0.0, 'fractional_bits': 3.0}}, + "submodule2": {"input": {"integer_bits": 1.0, "fractional_bits": 3.0}}, + "activation": {"input": {"integer_bits": 0.0, "fractional_bits": 4.0}}, + "activation2": {"input": {"integer_bits": 0.0, "fractional_bits": 3.0}}, } model = TestModel2(layer, layer2, "relu", "tanh") @@ -1166,7 +1175,6 @@ def test_batchnorm2d_direct_hgq(config_pdp, conv2d_input): class DummyLayer(nn.Module): - def __init__(self, is_pretraining=False): super().__init__() self.built = True @@ -1736,7 +1744,6 @@ def dummy_hgq_loss(): class ModelWithAllLayers(nn.Module): - def __init__(self, use_bias=True): super().__init__() self.conv = Conv2d(IN_FEATURES, OUT_FEATURES, KERNEL_SIZE, bias=use_bias) diff --git a/tests/test_torch_missing_quantizer_tracing.py b/tests/test_torch_missing_quantizer_tracing.py new file mode 100644 index 0000000..635a9fb --- /dev/null +++ b/tests/test_torch_missing_quantizer_tracing.py @@ -0,0 +1,445 @@ +import os + +import pytest +import torch +from torch import nn + +os.environ["KERAS_BACKEND"] = "torch" + +from pquant.activations import PQActivation # noqa: E402 +from pquant.layers import PQDense # noqa: E402 + +from pquant.core.hyperparameter_optimization import PQConfig # noqa: E402 +from pquant.core.torch.quantizer import Quantizer # noqa: E402 +from pquant.core.torch.tracing import check_quantization # noqa: E402 + +BATCH_SIZE = 4 +OUT_FEATURES = 32 +IN_FEATURES = 16 + + +@pytest.fixture +def config_pdp(): + cfg = { + "pruning_parameters": { + "disable_pruning_for_layers": [], + "enable_pruning": True, + "epsilon": 1.0, + "pruning_method": "pdp", + "sparsity": 0.75, + "temperature": 1e-5, + "threshold_decay": 0.0, + "structured_pruning": False, + }, + "quantization_parameters": { + "default_weight_integer_bits": 0.0, + "default_weight_fractional_bits": 7.0, + "default_data_integer_bits": 0.0, + "default_data_fractional_bits": 7.0, + "default_data_keep_negatives": 0.0, + "default_weight_keep_negatives": 1.0, + "quantize_input": True, + "quantize_output": False, + "enable_quantization": True, + "hgq_gamma": 0.0003, + "hgq_beta": 1e-5, + "hgq_heterogeneous": True, + "layer_specific": {}, + "use_high_granularity_quantization": False, + "use_real_tanh": False, + "use_relu_multiplier": True, + "use_symmetric_quantization": False, + "round_mode": "RND", + "overflow_mode_parameters": "SAT", + "overflow_mode_data": "SAT", + "granularity": "per_tensor", + }, + "training_parameters": {"pruning_first": False}, + "fitcompress_parameters": {"enable_fitcompress": False}, + } + return PQConfig.load_from_config(cfg) + + +def test_check_quantization_passes(config_pdp): + class GoodModel(nn.Module): + def __init__(self): + super().__init__() + self.d1 = PQDense(config_pdp, IN_FEATURES, OUT_FEATURES, quantize_input=True, quantize_output=True) + self.act = PQActivation(config_pdp, "relu", quantize_input=False, quantize_output=True) + self.d2 = PQDense(config_pdp, OUT_FEATURES, OUT_FEATURES, quantize_input=False, quantize_output=True) + + def forward(self, x): + return self.d2(self.act(self.d1(x))) + + assert check_quantization(GoodModel()) is True + + +def test_check_quantization_fails(config_pdp): + # d1.quantize_output=False, so the `a + a` add receives unquantized inputs. + class BadModel(nn.Module): + def __init__(self): + super().__init__() + self.d1 = PQDense(config_pdp, IN_FEATURES, OUT_FEATURES, quantize_input=True, quantize_output=False) + self.d2 = PQDense(config_pdp, OUT_FEATURES, OUT_FEATURES, quantize_input=False, quantize_output=True) + + def forward(self, x): + a = self.d1(x) + b = a + a + return self.d2(b) + + result = check_quantization(BadModel()) + assert isinstance(result, list) + assert len(result) >= 1 + joined = "\n".join(result) + assert "add" in joined + assert "d2" in joined + assert "not quantized" in joined + + +def _build_chain_model(config, a_qin, a_qout, b_qin, b_qout, op="add"): + class ChainModel(nn.Module): + def __init__(self): + super().__init__() + self.a = PQDense(config, IN_FEATURES, OUT_FEATURES, quantize_input=a_qin, quantize_output=a_qout) + self.b = PQDense(config, OUT_FEATURES, OUT_FEATURES, quantize_input=b_qin, quantize_output=b_qout) + + def forward(self, x): + y = self.a(x) + if op == "add": + y = y + y + elif op == "sigmoid": + y = torch.sigmoid(y) + else: + raise ValueError(op) + return self.b(y) + + return ChainModel() + + +def test_check_quantization_chain_all_true_passes(config_pdp): + model = _build_chain_model(config_pdp, True, True, True, True) + assert check_quantization(model) is True + + +def test_check_quantization_chain_a_qin_false(config_pdp): + # 'a' has quantize_input=False but receives the raw (unquantized) model input. + model = _build_chain_model(config_pdp, False, True, True, True) + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'a'" in joined + assert "quantize_input=False" in joined + assert "'b'" not in joined + assert "add" not in joined + + +def test_check_quantization_chain_b_qin_false(config_pdp): + # 'b' has quantize_input=False but its input comes from the unquantized `add` op output. + model = _build_chain_model(config_pdp, True, True, True, True) + model.b.quantize_input = False + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'b'" in joined + assert "quantize_input=False" in joined + assert "'a'" not in joined + + +def test_check_quantization_chain_a_qout_false(config_pdp): + # 'a'.quantize_output=False, so the `add` op consuming 'a' receives unquantized input. + model = _build_chain_model(config_pdp, True, False, True, True) + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'add'" in joined + assert "'a'" in joined + assert "not quantized" in joined + assert "'b'" not in joined + + +def test_check_quantization_chain_b_qout_false(config_pdp): + # 'b'.quantize_output=False, so the model output is left unquantized. + model = _build_chain_model(config_pdp, True, True, True, False) + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'output'" in joined + assert "'b'" in joined + assert "not quantized" in joined + + +def test_check_quantization_chain_both_qin_false(config_pdp): + # Both 'a' and 'b' have quantize_input=False while their inputs are unquantized. + model = _build_chain_model(config_pdp, False, True, False, True) + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'a'" in joined + assert "'b'" in joined + assert joined.count("quantize_input=False") == 2 + + +def test_check_quantization_chain_both_qout_false(config_pdp): + # First layer output and model output unquantized + model = _build_chain_model(config_pdp, True, False, True, False) + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'add'" in joined + assert "'output'" in joined + assert joined.count("not quantized") == 2 + + +def test_check_quantization_chain_all_false(config_pdp): + # All quantize_input/output flags are False. + model = _build_chain_model(config_pdp, False, False, False, False) + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'a'" in joined + assert "'add'" in joined + assert "'b'" in joined + assert "'output'" in joined + assert joined.count("quantize_input=False") == 2 + + +def test_check_quantization_unary_op_between(config_pdp): + # 'a'.quantize_output=False. + model = _build_chain_model(config_pdp, True, False, True, True, op="sigmoid") + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "sigmoid" in joined + assert "'a'" in joined + + +def test_check_quantization_fix_pqml_producer_flips_flag_unbuilt(config_pdp): + # 'a'.quantize_output=False feeds the `add` op, change quantize_output to True. + model = _build_chain_model(config_pdp, True, False, True, True) + assert model.a.quantize_output is False + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + assert model.a.quantize_output is True + added = [name for name, _ in traced.named_modules() if name.startswith("_auto_missing_quantizer_")] + assert added == [] + assert check_quantization(traced) is True + + +def test_check_quantization_fix_pqml_producer_flips_flag_built(config_pdp): + # Same as above but 'a' is already built: the fix flips the flag and constructs its output_quantizer in-place. + model = _build_chain_model(config_pdp, True, False, True, True) + x = torch.randn(BATCH_SIZE, IN_FEATURES) + model(x) + assert model.a.built is True + assert not hasattr(model.a, "output_quantizer") + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + assert model.a.quantize_output is True + assert hasattr(model.a, "output_quantizer") + added = [name for name, _ in traced.named_modules() if name.startswith("_auto_missing_quantizer_")] + assert added == [] + assert check_quantization(traced) is True + + +def test_check_quantization_fix_output_flips_pqml_flag(config_pdp): + # 'b'.quantize_output=False leaves the model output unquantized, so the fix flips 'b'.quantize_output. + model = _build_chain_model(config_pdp, True, True, True, False) + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + assert model.b.quantize_output is True + added = [name for name, _ in traced.named_modules() if name.startswith("_auto_missing_quantizer_")] + assert added == [] + assert check_quantization(traced) is True + + +def test_check_quantization_fix_placeholder_producer(config_pdp): + # `x + x` consumes the raw placeholder input (unquantized), so the fix inserts a standalone quantizer. + class PHModel(nn.Module): + def __init__(self): + super().__init__() + self.b = PQDense(config_pdp, IN_FEATURES, OUT_FEATURES, quantize_input=True, quantize_output=True) + + def forward(self, x): + return self.b(x + x) + + model = PHModel() + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + added = [name for name, _ in traced.named_modules() if name.startswith("_auto_missing_quantizer_")] + assert len(added) == 1 + assert check_quantization(traced) is True + + +def test_check_quantization_fix_functional_producer_inserts_quantizer(config_pdp): + # `add` feeds `sigmoid` and `sigmoid` feeds 'b' (both function outputs are unquantized). + class ChainedOps(nn.Module): + def __init__(self): + super().__init__() + self.a = PQDense(config_pdp, IN_FEATURES, OUT_FEATURES, quantize_input=True, quantize_output=True) + self.b = PQDense(config_pdp, OUT_FEATURES, OUT_FEATURES, quantize_input=False, quantize_output=True) + + def forward(self, x): + y = self.a(x) + y = torch.sigmoid(y + y) + return self.b(y) + + model = ChainedOps() + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + added = [name for name, _ in traced.named_modules() if name.startswith("_auto_missing_quantizer_")] + assert len(added) == 2 + assert check_quantization(traced) is True + + +def test_check_quantization_fix_roundtrip_all_false(config_pdp): + # All flags False, so the fix must flip both layers' quantize_output to fully quantize the graph. + model = _build_chain_model(config_pdp, False, False, False, False) + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + assert model.a.quantize_output is True + assert model.b.quantize_output is True + assert check_quantization(traced) is True + + +def test_check_quantization_fix_requires_config(config_pdp): + # Fixing missing quantizer requires a config (here omitted, so it raises an error). + model = _build_chain_model(config_pdp, True, False, True, True) + with pytest.raises(ValueError): + check_quantization(model, add_missing_quantizers=True) + + +def test_check_quantization_pqactivation_relu_preserves_quantization(config_pdp): + # A relu PQActivation with quantize_input=True quantizes its input, and relu is a + # grid-preserving clip, so its output is already quantized: the `add` it feeds needs + # no extra quantizer. The model passes as-is and the activation's quantize_output stays False. + class ActChain(nn.Module): + def __init__(self): + super().__init__() + self.act = PQActivation(config_pdp, "relu", quantize_input=True, quantize_output=False) + self.b = PQDense(config_pdp, IN_FEATURES, OUT_FEATURES, quantize_input=True, quantize_output=True) + + def forward(self, x): + y = self.act(x) + y = y + y + return self.b(y) + + model = ActChain() + assert check_quantization(model) is True + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + assert model.act.quantize_output is False # relu needs no output quantizer + added = [name for name, _ in traced.named_modules() if name.startswith("_auto_missing_quantizer_")] + assert added == [] + assert check_quantization(traced) is True + + +def test_check_quantization_pqactivation_nonlinear_producer_fix(config_pdp): + # A tanh PQActivation is NOT grid-preserving, so its quantize_output=False output feeds + # the `add` unquantized and the fix flips the activation's quantize_output to True. + class ActChain(nn.Module): + def __init__(self): + super().__init__() + self.act = PQActivation(config_pdp, "tanh", quantize_input=True, quantize_output=False) + self.b = PQDense(config_pdp, IN_FEATURES, OUT_FEATURES, quantize_input=True, quantize_output=True) + + def forward(self, x): + y = self.act(x) + y = y + y + return self.b(y) + + model = ActChain() + result = check_quantization(model) + assert isinstance(result, list) + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + assert model.act.quantize_output is True + assert check_quantization(traced) is True + + +def test_check_quantization_multiple_consumers_of_pqml_output(config_pdp): + # 'a'.quantize_output=False feeds both the `add` op and the model output, quantize a output. + class MultiConsumer(nn.Module): + def __init__(self): + super().__init__() + self.a = PQDense(config_pdp, IN_FEATURES, OUT_FEATURES, quantize_input=True, quantize_output=False) + self.b = PQDense(config_pdp, OUT_FEATURES, OUT_FEATURES, quantize_input=True, quantize_output=True) + + def forward(self, x): + y = self.a(x) + z = y + y + return self.b(z), y + + model = MultiConsumer() + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'add'" in joined + assert "'output'" in joined + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + assert model.a.quantize_output is True + added = [name for name, _ in traced.named_modules() if name.startswith("_auto_missing_quantizer_")] + assert added == [] + assert check_quantization(traced) is True + + +def _make_quantizer(config): + qp = config.quantization_parameters + return Quantizer( + k=qp.default_data_keep_negatives, + i=qp.default_data_integer_bits, + f=qp.default_data_fractional_bits, + overflow=qp.overflow_mode_data, + round_mode=qp.round_mode, + is_heterogeneous=False, + is_data=True, + granularity="per_tensor", + hgq_gamma=qp.hgq_gamma, + ) + + +def _build_dense_matmul_skip_model(config, dense_qout): + class DenseMatmulSkip(nn.Module): + def __init__(self): + super().__init__() + self.d = PQDense(config, IN_FEATURES, IN_FEATURES, quantize_input=True, quantize_output=dense_qout) + self.register_buffer("w", torch.randn(IN_FEATURES, IN_FEATURES)) + # Route the constant matrix through a Quantizer so it counts as + # "assumed quantized" (the tracer treats a raw get_attr as unquantized). + self.wq = _make_quantizer(config) + + def forward(self, x): + y = self.d(x) + y = torch.matmul(y, self.wq(self.w)) # matmul with constant matrix (assumed quantized) + y = y + x # skip connection from the (unquantized) model input + return y + + return DenseMatmulSkip() + + +def test_check_quantization_dense_matmul_skip_from_input(config_pdp): + # matmul output is unquantized and the skip connection feeds the raw model input 'x' into the + # `add`, so the add inputs and the model output are flagged; the assumed-quantized constant is not. + model = _build_dense_matmul_skip_model(config_pdp, dense_qout=True) + result = check_quantization(model) + assert isinstance(result, list) + joined = "\n".join(result) + assert "'add'" in joined + assert "input 'x' is not quantized" in joined + assert "input 'matmul' is not quantized" in joined + assert "'output'" in joined + # The constant matmul operand is quantized, so matmul itself reports no missing input. + assert not any(r.startswith("'matmul'") for r in result) + + +def test_check_quantization_dense_matmul_skip_from_input_fix(config_pdp): + # Fixing inserts standalone quantizers on the unquantized skip input, the matmul output, and the + # model output (the assumed-quantized constant needs none). + model = _build_dense_matmul_skip_model(config_pdp, dense_qout=True) + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + added = [name for name, _ in traced.named_modules() if name.startswith("_auto_missing_quantizer_")] + assert len(added) == 3 + assert check_quantization(traced) is True + + +def test_check_quantization_dense_matmul_skip_unquantized_dense(config_pdp): + # With dense.quantize_output=False the matmul also sees an unquantized data input from 'd'; the fix + # flips the PQDense flag (a PQuantML producer) and inserts quantizers for the rest. + model = _build_dense_matmul_skip_model(config_pdp, dense_qout=False) + result = check_quantization(model) + assert isinstance(result, list) + assert any(r.startswith("'matmul'") and "input 'd' is not quantized" in r for r in result) + traced = check_quantization(model, add_missing_quantizers=True, config=config_pdp) + assert model.d.quantize_output is True + assert check_quantization(traced) is True diff --git a/tests/test_torch_onnx_converter.py b/tests/test_torch_onnx_converter.py new file mode 100644 index 0000000..c3570bd --- /dev/null +++ b/tests/test_torch_onnx_converter.py @@ -0,0 +1,385 @@ +"""Tests for convert_to_onnx / convert_to_onnx_fx. + +Each test builds a small model (one PQ layer + ReLU where applicable), runs a +forward pass to initialise any running statistics, calls apply_final_compression +on every PQ module, exports to ONNX with convert_to_onnx(), and then verifies +that onnxruntime produces the same output as the PyTorch model. + +The same check is repeated with bias=True and bias=False via parametrize. +""" + +import os + +import numpy as np +import pytest +import torch +import torch.nn as nn + +os.environ["KERAS_BACKEND"] = "torch" + +from pquant.layers import ( # noqa: E402 + PQAvgPool1d, + PQAvgPool2d, + PQBatchNorm1d, + PQBatchNorm2d, + PQConv1d, + PQConv2d, + PQDense, + PQMultiheadAttention, +) + +import pquant # noqa: E402 +from pquant.core.torch.convert_to_onnx import ( # noqa: E402 + convert_to_onnx, + convert_to_onnx_fx, + export_qdq_layernorm, +) + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime not installed") + +ATOL = 1e-4 # float32 Gemm/Conv can differ by ~1 ULP; keep some slack + + +# --------------------------------------------------------------------------- +# fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def cfg(): + c = pquant.cs_config() + c.quantization_parameters.enable_quantization = False + return c + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _apply_compression(model: nn.Module): + for m in model.modules(): + if hasattr(m, "apply_final_compression"): + m.apply_final_compression() + + +def _onnx_run(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path) -> np.ndarray: + """Export model → ONNX file in tmp_path, run with onnxruntime, return output.""" + path = str(tmp_path / "model.onnx") + convert_to_onnx(model, input_shape=input_shape, output_path=path) + sess = ort.InferenceSession(path) + in_name = sess.get_inputs()[0].name + return sess.run(None, {in_name: x.cpu().numpy()})[0] + + +def _onnx_run_fx(model: nn.Module, x: torch.Tensor, input_shape: tuple, tmp_path) -> np.ndarray: + """FX-based export → ONNX, run with onnxruntime.""" + path = str(tmp_path / "model_fx.onnx") + convert_to_onnx_fx(model, input_shape=input_shape, output_path=path) + sess = ort.InferenceSession(path) + in_name = sess.get_inputs()[0].name + return sess.run(None, {in_name: x.cpu().numpy()})[0] + + +def _torch_out(model: nn.Module, x: torch.Tensor) -> np.ndarray: + model.eval() + with torch.no_grad(): + return model(x).cpu().numpy() + + +# --------------------------------------------------------------------------- +# PQDense +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bias", [True, False]) +def test_dense_onnx(cfg, bias, tmp_path): + IN, OUT = 16, 8 + model = nn.Sequential( + PQDense(cfg, in_features=IN, out_features=OUT, bias=bias), + nn.ReLU(), + ) + x = torch.randn(4, IN) + with torch.no_grad(): + model(x) # warm-up (needed for any running stats) + _apply_compression(model) + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(IN,), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQDense bias={bias}: torch vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQConv2d +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bias", [True, False]) +def test_conv2d_onnx(cfg, bias, tmp_path): + IN_C, OUT_C, H, W = 3, 8, 8, 8 + model = nn.Sequential( + PQConv2d(cfg, in_channels=IN_C, out_channels=OUT_C, kernel_size=3, padding=1, bias=bias), + nn.ReLU(), + ) + x = torch.randn(2, IN_C, H, W) + with torch.no_grad(): + model(x) + _apply_compression(model) + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(IN_C, H, W), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQConv2d bias={bias}: torch vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQConv1d +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bias", [True, False]) +def test_conv1d_onnx(cfg, bias, tmp_path): + IN_C, OUT_C, L = 4, 8, 16 + model = nn.Sequential( + PQConv1d(cfg, in_channels=IN_C, out_channels=OUT_C, kernel_size=3, padding=1, bias=bias), + nn.ReLU(), + ) + x = torch.randn(2, IN_C, L) + with torch.no_grad(): + model(x) + _apply_compression(model) + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(IN_C, L), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg=f"PQConv1d bias={bias}: torch vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQBatchNorm2d +# --------------------------------------------------------------------------- + + +def test_batchnorm2d_onnx(cfg, tmp_path): + C, H, W = 8, 4, 4 + model = nn.Sequential( + PQBatchNorm2d(cfg, num_features=C), + nn.ReLU(), + ) + x = torch.randn(4, C, H, W) + with torch.no_grad(): + model(x) + _apply_compression(model) + model.eval() # switch BN to use running stats + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(C, H, W), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQBatchNorm2d: torch vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQBatchNorm1d +# --------------------------------------------------------------------------- + + +def test_batchnorm1d_onnx(cfg, tmp_path): + C, L = 8, 16 + model = nn.Sequential( + PQBatchNorm1d(cfg, num_features=C), + nn.ReLU(), + ) + x = torch.randn(4, C, L) + with torch.no_grad(): + model(x) + _apply_compression(model) + model.eval() + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(C, L), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQBatchNorm1d: torch vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQAvgPool2d +# --------------------------------------------------------------------------- + + +def test_avgpool2d_onnx(cfg, tmp_path): + C, H, W = 8, 8, 8 + model = nn.Sequential( + PQAvgPool2d(cfg, kernel_size=2, stride=2), + ) + x = torch.randn(2, C, H, W) + with torch.no_grad(): + model(x) + _apply_compression(model) + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(C, H, W), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQAvgPool2d: torch vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQAvgPool1d +# --------------------------------------------------------------------------- + + +def test_avgpool1d_onnx(cfg, tmp_path): + C, L = 8, 16 + model = nn.Sequential( + PQAvgPool1d(cfg, kernel_size=2, stride=2), + ) + x = torch.randn(2, C, L) + with torch.no_grad(): + model(x) + _apply_compression(model) + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run(model, x, input_shape=(C, L), tmp_path=tmp_path) + np.testing.assert_allclose(torch_out, onnx_out, atol=ATOL, err_msg="PQAvgPool1d: torch vs ONNX mismatch") + + +# --------------------------------------------------------------------------- +# PQMultiheadAttention (uses FX converter; self-attention, batch_first=True) +# --------------------------------------------------------------------------- + + +class _SelfAttnModel(nn.Module): + """Thin wrapper so FX tracing sees a single-input model.""" + + def __init__(self, mha: PQMultiheadAttention): + super().__init__() + self.mha = mha + + def forward(self, x): + out, _ = self.mha(x, x, x) + return out + + +@pytest.mark.parametrize("bias", [True, False]) +def test_mha_onnx(cfg, bias, tmp_path): + E, H, T = 16, 4, 8 + mha = PQMultiheadAttention(cfg, embed_dim=E, num_heads=H, bias=bias, batch_first=True) + model = _SelfAttnModel(mha) + + x = torch.randn(2, T, E) + with torch.no_grad(): + model(x) + _apply_compression(model) + + torch_out = _torch_out(model, x) + onnx_out = _onnx_run_fx(model, x, input_shape=(T, E), tmp_path=tmp_path) + np.testing.assert_allclose( + torch_out, onnx_out, atol=ATOL, err_msg=f"PQMultiheadAttention bias={bias}: torch vs ONNX mismatch" + ) + + +# --------------------------------------------------------------------------- +# Static-QDQ LayerNormalization graph +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("input_shape", [(4, 64), (1, 4, 64)]) +def test_qdq_layernorm_export(input_shape, tmp_path): + import onnx + + D = input_shape[-1] + rng = np.random.default_rng(0) + # Q7 representable: gamma = k / 128, k integer, |k| < 32768 + gamma_q = rng.integers(low=64, high=192, size=(D,), dtype=np.int32) # ~0.5 .. 1.5 + gamma = (gamma_q.astype(np.float32)) / (1 << 7) + # Q15 representable: beta = k / 32768, k integer, |k| < 32768 (so |beta| < 1) + beta_q = rng.integers(low=-1024, high=1024, size=(D,), dtype=np.int32) + beta = (beta_q.astype(np.float32)) / (1 << 15) + + input_scale_log2 = -7 # input_scale = 2**-7 + output_scale_log2 = -6 # output_scale = 2**-6 + eps_q0 = 1 + + path = str(tmp_path / "qdq_layernorm.onnx") + model_proto = export_qdq_layernorm( + output_path=path, + input_shape=input_shape, + gamma=gamma, + beta=beta, + input_scale_log2=input_scale_log2, + output_scale_log2=output_scale_log2, + eps_q0=eps_q0, + ) + + # ----- structural checks ----- + op_types = [n.op_type for n in model_proto.graph.node] + assert op_types == ["DequantizeLinear", "LayerNormalization", "QuantizeLinear", "DequantizeLinear"] + + ln_node = model_proto.graph.node[1] + axis = next(a.i for a in ln_node.attribute if a.name == "axis") + eps_attr = next(a.f for a in ln_node.attribute if a.name == "epsilon") + assert axis == -1 + expected_eps = eps_q0 * (2.0**input_scale_log2) ** 2 + assert abs(eps_attr - expected_eps) < 1e-12 + + # input must be int8, output float + assert len(model_proto.graph.input) == 1 + assert model_proto.graph.input[0].type.tensor_type.elem_type == onnx.TensorProto.INT8 + assert model_proto.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT + in_dims = [d.dim_value for d in model_proto.graph.input[0].type.tensor_type.shape.dim] + assert tuple(in_dims) == input_shape + + # zero-points must be int8 zero + inits = {t.name: t for t in model_proto.graph.initializer} + for zp_name in ("input_zero_point", "output_zero_point"): + zp = onnx.numpy_helper.to_array(inits[zp_name]) + assert zp.dtype == np.int8 + assert int(zp) == 0 + + # scales must be exact powers of two + in_scale = float(onnx.numpy_helper.to_array(inits["input_scale"])) + out_scale = float(onnx.numpy_helper.to_array(inits["output_scale"])) + assert in_scale == 2.0**input_scale_log2 + assert out_scale == 2.0**output_scale_log2 + + # ----- numerical check via onnxruntime ----- + sess = ort.InferenceSession(path) + in_name = sess.get_inputs()[0].name + x_q = rng.integers(low=-64, high=64, size=input_shape, dtype=np.int8) + onnx_out = sess.run(None, {in_name: x_q})[0] + + # Reference: dequantize -> layernorm(axis=-1) -> quantize -> dequantize + x_f = x_q.astype(np.float32) * in_scale + mean = x_f.mean(axis=-1, keepdims=True) + var = x_f.var(axis=-1, keepdims=True) + x_norm = (x_f - mean) / np.sqrt(var + expected_eps) + y_f = x_norm * gamma + beta + y_q = np.clip(np.round(y_f / out_scale), -128, 127).astype(np.int8) + y_ref = y_q.astype(np.float32) * out_scale + + np.testing.assert_allclose(onnx_out, y_ref, atol=out_scale * 0.5) + + +def test_qdq_layernorm_validation(tmp_path): + path = str(tmp_path / "bad.onnx") + D = 64 + gamma = np.ones(D, dtype=np.float32) + beta = np.zeros(D, dtype=np.float32) + + # rank-1 input: rejected + with pytest.raises(ValueError, match="rank"): + export_qdq_layernorm(path, (D,), gamma, beta, -7, -6) + + # last dim not multiple of 32 + with pytest.raises(ValueError, match="multiple of 32"): + export_qdq_layernorm(path, (4, 16), np.ones(16, np.float32), np.zeros(16, np.float32), -7, -6) + + # last dim not power of two (96 = 32*3) + with pytest.raises(ValueError, match="power of two"): + export_qdq_layernorm(path, (4, 96), np.ones(96, np.float32), np.zeros(96, np.float32), -7, -6) + + # gamma not Q7-representable (1/3 is not k/128 exactly) + with pytest.raises(ValueError, match="gamma"): + export_qdq_layernorm(path, (4, D), np.full(D, 1.0 / 3.0, np.float32), beta, -7, -6) + + # beta not Q15-representable (1/3 is not k/32768 exactly) + with pytest.raises(ValueError, match="beta"): + export_qdq_layernorm(path, (4, D), gamma, np.full(D, 1.0 / 3.0, np.float32), -7, -6) + + # eps_q0 < 1 + with pytest.raises(ValueError, match="eps_q0"): + export_qdq_layernorm(path, (4, D), gamma, beta, -7, -6, eps_q0=0) diff --git a/tests/test_torch_pruning_layers.py b/tests/test_torch_pruning_layers.py new file mode 100644 index 0000000..88a8eb3 --- /dev/null +++ b/tests/test_torch_pruning_layers.py @@ -0,0 +1,638 @@ +""" +Parity tests: PyTorch pruning methods vs Keras reference. + +For each pruning layer (Activation, PDP, CS, DST, Wanda, AutoSparse, MDMM) +the torch implementation must produce numerically matching outputs and +masks when given identical state and inputs as the keras version. Metric +functions (Structured/Unstructured sparsity) are also compared directly. +""" + +import os + +# The keras layers are expected to run on the tensorflow backend (the long-term +# target for the keras implementation). Run this test file with +# ``KERAS_BACKEND=tensorflow`` so the keras forward paths (including those +# routed through ``ops.custom_gradient``) exercise their intended backend. +os.environ.setdefault("KERAS_BACKEND", "tensorflow") + +import numpy as np # noqa: E402 +import pytest # noqa: E402 +import torch # noqa: E402 +from keras import ops # noqa: E402 + +from pquant.core.keras.pruning_methods.activation_pruning import ( # noqa: E402 + ActivationPruning as KerasActivationPruning, +) +from pquant.core.keras.pruning_methods.autosparse import ( # noqa: E402 + AutoSparse as KerasAutoSparse, +) +from pquant.core.keras.pruning_methods.cs import ( # noqa: E402 + ContinuousSparsification as KerasContinuousSparsification, +) +from pquant.core.keras.pruning_methods.dst import DST as KerasDST # noqa: E402 +from pquant.core.keras.pruning_methods.mdmm import MDMM as KerasMDMM # noqa: E402 +from pquant.core.keras.pruning_methods.metric_functions import ( # noqa: E402 + StructuredSparsityMetric as KerasStructuredSparsityMetric, +) +from pquant.core.keras.pruning_methods.metric_functions import ( # noqa: E402 + UnstructuredSparsityMetric as KerasUnstructuredSparsityMetric, +) +from pquant.core.keras.pruning_methods.pdp import PDP as KerasPDP # noqa: E402 +from pquant.core.keras.pruning_methods.wanda import Wanda as KerasWanda # noqa: E402 +from pquant.core.torch.pruning_methods.activation_pruning import ( # noqa: E402 + ActivationPruning as TorchActivationPruning, +) +from pquant.core.torch.pruning_methods.autosparse import ( # noqa: E402 + AutoSparse as TorchAutoSparse, +) +from pquant.core.torch.pruning_methods.cs import ( # noqa: E402 + ContinuousSparsification as TorchContinuousSparsification, +) +from pquant.core.torch.pruning_methods.dst import DST as TorchDST # noqa: E402 +from pquant.core.torch.pruning_methods.mdmm import MDMM as TorchMDMM # noqa: E402 +from pquant.core.torch.pruning_methods.metric_functions import ( # noqa: E402 + StructuredSparsityMetric as TorchStructuredSparsityMetric, +) +from pquant.core.torch.pruning_methods.metric_functions import ( # noqa: E402 + UnstructuredSparsityMetric as TorchUnstructuredSparsityMetric, +) +from pquant.core.torch.pruning_methods.pdp import PDP as TorchPDP # noqa: E402 +from pquant.core.torch.pruning_methods.wanda import Wanda as TorchWanda # noqa: E402 + +ABSOLUTE_TOLERANCE = 1e-5 +RELATIVE_TOLERANCE = 1e-4 + + +def to_numpy(x): + if isinstance(x, torch.Tensor): + return x.detach().cpu().numpy() + return np.asarray(ops.convert_to_numpy(x)) + + +def assert_close(a, b, atol=ABSOLUTE_TOLERANCE, rtol=RELATIVE_TOLERANCE, msg=""): + a_np = to_numpy(a) + b_np = to_numpy(b) + assert a_np.shape == b_np.shape, f"{msg}: shape mismatch: {a_np.shape} vs {b_np.shape}" + np.testing.assert_allclose(a_np, b_np, atol=atol, rtol=rtol, err_msg=msg) + + +def keras_tensor(arr): + return ops.convert_to_tensor(np.asarray(arr).astype(np.float32)) + + +def torch_tensor(arr, requires_grad=False): + return torch.as_tensor(np.asarray(arr).astype(np.float32)).requires_grad_(requires_grad) + + +def reset_seed(seed=0): + np.random.seed(seed) + torch.manual_seed(seed) + + +def keras_grad(fn, x): + """Gradient of scalar-reducing ``fn(x).sum()`` w.r.t. keras tensor ``x``.""" + import keras + + if keras.backend.backend() == "tensorflow": + import tensorflow as tf + + xt = tf.convert_to_tensor(x) + with tf.GradientTape() as tape: + tape.watch(xt) + loss = ops.sum(fn(xt)) + return tape.gradient(loss, xt) + raise RuntimeError("keras_grad only supports the tensorflow backend") + + +def ap_config(): + return { + "pruning_parameters": { + "pruning_method": "activation_pruning", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "threshold": 0.3, + "t_start_collecting_batch": 0, + "threshold_decay": 0.0, + "t_delta": 2, + }, + } + + +@pytest.mark.parametrize( + "layer_type,shape", + [ + ("linear", (16, 8)), + ("conv", (16, 8, 3, 3)), + ], +) +def test_activation_pruning_matches_keras(layer_type, shape): + cfg = ap_config() + out_channels = shape[0] + batch = 32 + + reset_seed() + weight_np = np.random.randn(*shape).astype(np.float32) + # Construct outputs with distinct per-channel activity levels so the + # resulting mask is non-trivial (some channels pct_active > threshold, + # some below). linspace from -0.5 to 1 puts roughly 1/3 of the channels + # below zero, guaranteeing those get pruned while the rest survive. + per_channel = np.linspace(-0.5, 1.0, num=out_channels, dtype=np.float32) + np.random.shuffle(per_channel) + if layer_type == "linear": + output_np = np.tile(per_channel[None, :], (batch, 1)) + else: + output_np = np.tile(per_channel[None, :, None, None], (batch, 1, 4, 4)) + + k_layer = KerasActivationPruning(cfg, layer_type) + k_layer.build(shape) + k_layer.post_pre_train_function() + + t_layer = TorchActivationPruning(cfg, layer_type) + t_layer.build(shape) + t_layer.post_pre_train_function() + + for _ in range(cfg["pruning_parameters"]["t_delta"]): + k_layer.collect_output(keras_tensor(output_np), training=True) + t_layer.collect_output(torch_tensor(output_np), training=True) + + k_layer.post_epoch_function(0, 1) + t_layer.post_epoch_function(0, 1) + + assert_close(k_layer.mask, t_layer.mask, msg=f"AP mask ({layer_type})") + + # Sanity check: the constructed per-channel outputs put ~1/3 of the + # channels at non-positive values, so their pct_active == 0 falls below + # the 0.3 threshold and they should be pruned. Channels with a positive + # value hit pct_active == 1 and survive. The expected pruned count is + # deterministic from per_channel — verifying we actually exercise both + # branches rather than matching a trivial all-ones mask. + mask_np = to_numpy(t_layer.mask) + pruned_fraction = float((mask_np == 0).sum()) / mask_np.size + # linspace(-0.5, 1.0, 16) → 6 values <= 0 → 6/16 = 0.375 pruned channels. + assert pruned_fraction == pytest.approx(0.375), f"AP mask ({layer_type}) pruned fraction {pruned_fraction} != 0.375" + + k_weight = keras_tensor(weight_np) + t_weight = torch_tensor(weight_np, requires_grad=True) + k_out = k_layer(k_weight) + t_out = t_layer(t_weight) + assert_close(k_out, t_out, msg=f"AP forward ({layer_type})") + + k_grad = keras_grad(k_layer, weight_np) + t_out.sum().backward() + assert_close(k_grad, t_weight.grad, msg=f"AP backward ({layer_type})") + + +def pdp_config(sparsity=0.75, structured=False): + return { + "pruning_parameters": { + "pruning_method": "pdp", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "epsilon": 1.0, + "sparsity": sparsity, + "temperature": 1e-5, + "threshold_decay": 0.0, + "structured_pruning": structured, + }, + } + + +@pytest.mark.parametrize( + "layer_type,shape,structured", + [ + ("linear", (16, 8), False), + ("linear", (16, 8), True), + ("conv", (16, 8, 3, 3), False), + ("conv", (16, 8, 3, 3), True), + ], +) +def test_pdp_matches_keras(layer_type, shape, structured): + cfg = pdp_config(structured=structured) + target_sparsity = cfg["pruning_parameters"]["sparsity"] + + reset_seed() + weight_np = np.random.randn(*shape).astype(np.float32) + + k_layer = KerasPDP(cfg, layer_type) + k_layer.build(shape) + k_layer.post_pre_train_function() + + t_layer = TorchPDP(cfg, layer_type) + t_layer.build(shape) + t_layer.post_pre_train_function() + + # Force the sparsity ramp to be fully complete. pre_epoch_function sets + # r = min(1, epsilon * (epoch + 1)) * init_r; with epsilon=1.0 that already + # puts the ramp multiplier at 1.0 on epoch 0, so r = init_r = 0.75. + k_layer.pre_epoch_function(0, None) + t_layer.pre_epoch_function(0, None) + + k_weight = keras_tensor(weight_np) + t_weight = torch_tensor(weight_np, requires_grad=True) + k_out = k_layer(k_weight) + t_out = t_layer(t_weight) + assert_close(k_out, t_out, msg=f"PDP forward ({layer_type}, structured={structured})") + + k_grad = keras_grad(k_layer, weight_np) + t_out.sum().backward() + assert_close(k_grad, t_weight.grad, msg=f"PDP backward ({layer_type}, structured={structured})") + + k_layer.update_mask(keras_tensor(weight_np)) + t_layer.update_mask(torch_tensor(weight_np)) + assert_close(k_layer.mask, t_layer.mask, msg="PDP mask after update_mask") + + # Verify the produced mask hits the configured target sparsity. + # For structured pruning the mask has shape (C, 1, ...) and encodes + # per-channel keep/prune; its sparsity directly equals the channel-level + # pruning fraction. For unstructured it's per-element. With temperature + # 1e-5 the soft mask is effectively binary, so use >= 0.5 to discretize. + t_mask_np = to_numpy(t_layer.mask) + actual_sparsity = float((t_mask_np < 0.5).sum()) / t_mask_np.size + assert actual_sparsity == pytest.approx(target_sparsity, abs=1e-6), ( + f"PDP {layer_type} (structured={structured}) mask sparsity {actual_sparsity} != target {target_sparsity}" + ) + + +def cs_config(threshold_decay=1e-4): + return { + "pruning_parameters": { + "pruning_method": "cs", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "threshold_init": -0.1, + "final_temp": 200, + "threshold_decay": threshold_decay, + }, + } + + +@pytest.mark.parametrize( + "shape", + [(16, 8), (16, 8, 3, 3)], +) +def test_cs_matches_keras(shape): + cfg = cs_config() + layer_type = "linear" if len(shape) == 2 else "conv" + + reset_seed() + s_override_np = (np.random.randn(*shape) * 0.5).astype(np.float32) + weight_np = np.random.randn(*shape).astype(np.float32) + + k_layer = KerasContinuousSparsification(cfg, layer_type) + k_layer.build(shape) + k_layer.post_pre_train_function() + k_layer.s.assign(keras_tensor(s_override_np)) + + t_layer = TorchContinuousSparsification(cfg, layer_type) + t_layer.build(shape) + t_layer.post_pre_train_function() + with torch.no_grad(): + t_layer.s.data.copy_(torch_tensor(s_override_np)) + + k_weight = keras_tensor(weight_np) + t_weight = torch_tensor(weight_np, requires_grad=True) + k_out = k_layer(k_weight) + t_out = t_layer(t_weight) + assert_close(k_out, t_out, msg=f"CS forward ({layer_type})") + + k_grad = keras_grad(k_layer, weight_np) + t_out.sum().backward() + assert_close(k_grad, t_weight.grad, msg=f"CS backward ({layer_type})") + + assert_close(k_layer.get_hard_mask(), t_layer.get_hard_mask(), msg="CS hard mask") + assert_close(k_layer.calculate_additional_loss(), t_layer.calculate_additional_loss(), msg="CS additional loss") + + # post_epoch_function updates beta — trajectories should match. + k_layer.post_epoch_function(0, 5) + t_layer.post_epoch_function(0, 5) + assert_close(k_layer.beta, t_layer.beta, msg="CS beta after post_epoch_function") + + +def dst_config(threshold_type="channelwise"): + return { + "pruning_parameters": { + "pruning_method": "dst", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "alpha": 5e-6, + "max_pruning_pct": 0.99, + "threshold_init": 0.0, + "threshold_type": threshold_type, + "threshold_decay": 0.0, + }, + } + + +@pytest.mark.parametrize( + "layer_type,shape,threshold_type", + [ + ("linear", (16, 8), "layerwise"), + ("linear", (16, 8), "channelwise"), + ("linear", (16, 8), "weightwise"), + ("conv", (16, 8, 3, 3), "channelwise"), + ], +) +def test_dst_matches_keras(layer_type, shape, threshold_type): + cfg = dst_config(threshold_type=threshold_type) + + reset_seed() + weight_np = (np.random.randn(*shape) * 0.5).astype(np.float32) + if threshold_type == "layerwise": + thr_np = np.array([[0.1]], dtype=np.float32) + elif threshold_type == "channelwise": + thr_np = (np.random.rand(shape[0], 1) * 0.2).astype(np.float32) + else: # weightwise + thr_np = (np.random.rand(shape[0], int(np.prod(shape[1:]))) * 0.2).astype(np.float32) + + k_layer = KerasDST(cfg, layer_type) + k_layer.build(shape) + k_layer.post_pre_train_function() + k_layer.threshold.assign(keras_tensor(thr_np)) + + t_layer = TorchDST(cfg, layer_type) + t_layer.build(shape) + t_layer.post_pre_train_function() + with torch.no_grad(): + t_layer.threshold.data.copy_(torch_tensor(thr_np)) + + k_weight = keras_tensor(weight_np) + t_weight = torch_tensor(weight_np, requires_grad=True) + k_out = k_layer(k_weight) + t_out = t_layer(t_weight) + assert_close(k_out, t_out, msg=f"DST forward ({layer_type}, {threshold_type})") + + k_grad = keras_grad(k_layer, weight_np) + t_out.sum().backward() + assert_close(k_grad, t_weight.grad, msg=f"DST backward ({layer_type}, {threshold_type})") + + assert_close( + k_layer.get_mask(keras_tensor(weight_np)), + t_layer.get_mask(torch_tensor(weight_np)), + msg=f"DST get_mask ({threshold_type})", + ) + assert_close(k_layer.calculate_additional_loss(), t_layer.calculate_additional_loss(), msg="DST additional loss") + + +def wanda_config(sparsity=0.75, N=None, M=None): + return { + "pruning_parameters": { + "pruning_method": "wanda", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "sparsity": sparsity, + "t_delta": 2, + "t_start_collecting_batch": 0, + "N": N, + "M": M, + "threshold_decay": 0.0, + "calculate_pruning_budget": True, + }, + } + + +@pytest.mark.parametrize( + "layer_type,shape,N,M", + [ + ("linear", (16, 8), None, None), + ("conv", (16, 8, 3, 3), None, None), + ("linear", (4, 8), 4, 8), + ("conv", (4, 8, 3, 3), 4, 8), + ], +) +def test_wanda_matches_keras(layer_type, shape, N, M): + cfg = wanda_config(N=N, M=M) + + reset_seed() + if layer_type == "linear": + x_np = np.random.randn(32, shape[1]).astype(np.float32) + else: + x_np = np.random.randn(32, shape[1], shape[2], shape[3]).astype(np.float32) + w_np = np.random.randn(*shape).astype(np.float32) + + k_layer = KerasWanda(cfg, layer_type) + k_layer.build(shape) + k_layer.post_pre_train_function() + + t_layer = TorchWanda(cfg, layer_type) + t_layer.build(shape) + t_layer.post_pre_train_function() + + for _ in range(cfg["pruning_parameters"]["t_delta"]): + k_layer.collect_input(keras_tensor(x_np), keras_tensor(w_np), training=True) + t_layer.collect_input(torch_tensor(x_np), torch_tensor(w_np), training=True) + + assert_close(k_layer.mask, t_layer.mask, msg=f"Wanda mask ({layer_type}, N={N}, M={M})") + + # Verify the mask hits the configured target sparsity. For N:M pruning + # Wanda internally uses N/M as the sparsity target; for unstructured it + # uses the configured sparsity directly. Mask values are strictly {0, 1} + # (produced by topk + scatter), so `== 0` counts pruned entries. + target_sparsity = (N / M) if (N is not None and M is not None) else cfg["pruning_parameters"]["sparsity"] + mask_np = to_numpy(t_layer.mask) + pruned_fraction = float((mask_np == 0).sum()) / mask_np.size + assert pruned_fraction == pytest.approx(target_sparsity), ( + f"Wanda {layer_type} (N={N}, M={M}) pruned fraction {pruned_fraction} != target {target_sparsity}" + ) + + k_weight = keras_tensor(w_np) + t_weight = torch_tensor(w_np, requires_grad=True) + k_out = k_layer(k_weight) + t_out = t_layer(t_weight) + assert_close(k_out, t_out, msg=f"Wanda forward ({layer_type}, N={N}, M={M})") + + k_grad = keras_grad(k_layer, w_np) + t_out.sum().backward() + assert_close(k_grad, t_weight.grad, msg=f"Wanda backward ({layer_type}, N={N}, M={M})") + + +def autosparse_config(threshold_type="channelwise", threshold_init=-2.0): + return { + "pruning_parameters": { + "pruning_method": "autosparse", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "alpha": 0.5, + "alpha_reset_epoch": 100, + "autotune_epochs": 10, + "backward_sparsity": False, + "threshold_init": threshold_init, + "threshold_type": threshold_type, + "threshold_decay": 0.0, + }, + } + + +@pytest.mark.parametrize( + "layer_type,shape,threshold_type", + [ + ("linear", (16, 8), "layerwise"), + ("linear", (16, 8), "channelwise"), + ("conv", (16, 8, 3, 3), "channelwise"), + ], +) +def test_autosparse_matches_keras(layer_type, shape, threshold_type): + # Keras AutoSparse.call() unconditionally routes through ops.custom_gradient, + # which under the torch backend errors because self.alpha is a keras Variable + # (save_for_backward rejects it). The path works fine under the tensorflow + # backend, which is the intended target for the keras implementation, so we + # skip this test when keras is on torch to avoid spurious failures. + import keras as _keras + + if _keras.backend.backend() == "torch": + pytest.skip("Keras AutoSparse forward is incompatible with the torch backend.") + cfg = autosparse_config(threshold_type=threshold_type) + + reset_seed() + weight_np = np.random.randn(*shape).astype(np.float32) + + k_layer = KerasAutoSparse(cfg, layer_type) + k_layer.build(shape) + k_layer.post_pre_train_function() + + t_layer = TorchAutoSparse(cfg, layer_type) + t_layer.build(shape) + t_layer.post_pre_train_function() + with torch.no_grad(): + t_layer.threshold.data.copy_(torch_tensor(to_numpy(k_layer.threshold))) + + assert_close( + k_layer.get_mask(keras_tensor(weight_np)), + t_layer.get_mask(torch_tensor(weight_np)), + msg=f"AutoSparse get_mask ({layer_type}, {threshold_type})", + ) + + k_weight = keras_tensor(weight_np) + t_weight = torch_tensor(weight_np, requires_grad=True) + k_out = k_layer(k_weight) + t_out = t_layer(t_weight) + assert_close(k_out, t_out, msg=f"AutoSparse forward ({layer_type}, {threshold_type})") + + k_grad = keras_grad(k_layer, weight_np) + t_out.sum().backward() + assert_close(k_grad, t_weight.grad, msg=f"AutoSparse backward ({layer_type}, {threshold_type})") + + # post_epoch_function updates alpha via decay; trajectories should match. + k_layer.post_epoch_function(3, 10) + t_layer.post_epoch_function(3, 10) + assert_close(k_layer.alpha, t_layer.alpha, msg="AutoSparse alpha after post_epoch_function") + + +def mdmm_config( + constraint_type="Equality", + metric_type="UnstructuredSparsity", + target_value=0.5, + use_grad=True, +): + return { + "pruning_parameters": { + "pruning_method": "mdmm", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "constraint_type": constraint_type, + "target_value": target_value, + "metric_type": metric_type, + "target_sparsity": 0.8, + "rf": 1, + "epsilon": 1e-3, + "scale": 1.0, + "damping": 1.0, + "use_grad": use_grad, + "l0_mode": "coarse", + "scale_mode": "mean", + "constraint_lr": 1e-3, + "threshold_decay": 0.0, + }, + } + + +@pytest.mark.parametrize( + "constraint_type,metric_type", + [ + ("Equality", "UnstructuredSparsity"), + ("LessThanOrEqual", "UnstructuredSparsity"), + ("GreaterThanOrEqual", "UnstructuredSparsity"), + ("Equality", "StructuredSparsity"), + ], +) +def test_mdmm_matches_keras(constraint_type, metric_type): + cfg = mdmm_config(constraint_type=constraint_type, metric_type=metric_type) + shape = (16, 8) + + reset_seed() + weight_np = (np.random.randn(*shape) * 0.2).astype(np.float32) + + k_layer = KerasMDMM(cfg, "linear") + k_layer.build(shape) + k_layer.post_pre_train_function() + + t_layer = TorchMDMM(cfg, "linear") + t_layer.build(shape) + t_layer.post_pre_train_function() + + k_weight = keras_tensor(weight_np) + t_weight = torch_tensor(weight_np, requires_grad=True) + k_out = k_layer(k_weight) + t_out = t_layer(t_weight) + assert_close(k_out, t_out, msg=f"MDMM forward ({constraint_type}, {metric_type})") + + k_grad = keras_grad(k_layer, weight_np) + t_out.sum().backward() + assert_close(k_grad, t_weight.grad, msg=f"MDMM backward ({constraint_type}, {metric_type})") + + assert_close( + k_layer.get_hard_mask(keras_tensor(weight_np)), + t_layer.get_hard_mask(torch_tensor(weight_np)), + msg="MDMM hard_mask", + ) + + # Constraint penalty: read directly from the constraint layer to avoid + # differences in how keras/torch surface accumulated losses. + k_penalty = ops.sum(k_layer.constraint_layer(keras_tensor(weight_np))) + t_penalty = t_layer.constraint_layer(torch_tensor(weight_np)).sum() + assert_close(k_penalty, t_penalty, msg="MDMM constraint penalty") + + +def test_mdmm_finetune_returns_masked_weight(): + """In finetuning mode both layers should return weight * hard_mask.""" + cfg = mdmm_config() + shape = (8, 6) + + reset_seed() + weight_np = (np.random.randn(*shape) * 0.2).astype(np.float32) + + k_layer = KerasMDMM(cfg, "linear") + k_layer.build(shape) + k_layer.post_pre_train_function() + k_layer.pre_finetune_function() + + t_layer = TorchMDMM(cfg, "linear") + t_layer.build(shape) + t_layer.post_pre_train_function() + t_layer.pre_finetune_function() + + k_out = k_layer(keras_tensor(weight_np)) + t_out = t_layer(torch_tensor(weight_np)) + assert_close(k_out, t_out, msg="MDMM finetune forward") + + +@pytest.mark.parametrize("l0_mode", ["coarse", "smooth"]) +@pytest.mark.parametrize("scale_mode", ["mean", "sum"]) +def test_unstructured_sparsity_metric_matches_keras(l0_mode, scale_mode): + k_layer = KerasUnstructuredSparsityMetric(l0_mode=l0_mode, scale_mode=scale_mode, target_sparsity=0.7, epsilon=1e-3) + t_layer = TorchUnstructuredSparsityMetric(l0_mode=l0_mode, scale_mode=scale_mode, target_sparsity=0.7, epsilon=1e-3) + + reset_seed() + w_np = (np.random.randn(16, 8) * 0.1).astype(np.float32) + + assert_close(k_layer(keras_tensor(w_np)), t_layer(torch_tensor(w_np)), msg=f"Unstructured({l0_mode},{scale_mode})") + + +@pytest.mark.parametrize("rf", [1, 4, 5]) +def test_structured_sparsity_metric_matches_keras(rf): + k_layer = KerasStructuredSparsityMetric(rf=rf, epsilon=1e-3) + t_layer = TorchStructuredSparsityMetric(rf=rf, epsilon=1e-3) + + reset_seed() + w_np = (np.random.randn(12, 7) * 0.05).astype(np.float32) + + assert_close(k_layer(keras_tensor(w_np)), t_layer(torch_tensor(w_np)), msg=f"Structured(rf={rf})") diff --git a/tests/test_wanda.py b/tests/test_wanda.py index 609572d..3af1026 100644 --- a/tests/test_wanda.py +++ b/tests/test_wanda.py @@ -2,7 +2,7 @@ import pytest from keras import ops -from pquant.pruning_methods.wanda import Wanda +from pquant.core.keras.pruning_methods.wanda import Wanda @pytest.fixture