[Torch FX] Compress PT2E Support #3663

anzr299 · 2025-09-22T14:43:32Z

Changes

Introduced a new API to offer weights compression algorithm for quantizers defined in torch.ao.
Currently only supports OpenVINO Quantizer.

Reason for changes

To support Quantizers defined in torch ao.

Related tickets

169342

src/nncf/experimental/torch/fx/quantization/quantize_pt2e.py

…match in signatures in prepare_pt2e.

…_algo

src/nncf/experimental/quantization/algorithms/weight_compression/algorithm.py

src/nncf/experimental/torch/fx/quantization/quantizer/__init__.py

tests/executorch/__init__.py

tests/executorch/test_quantizer.py

tests/executorch/observers.py

daniil-lyakhov

Can I see the PR with OpenVINOQuantizer?

daniil-lyakhov · 2025-09-23T15:28:47Z

src/nncf/experimental/quantization/algorithms/weight_compression/algorithm.py

+from nncf.quantization.algorithms.weight_compression.algorithm import WeightCompression
+
+
+class WeightsCompressionPT2E(Algorithm):


This algorithm is not designed for the PT2E, this is experimental WC algorithm which could be implemented in any backend

Suggested change

class WeightsCompressionPT2E(Algorithm):

class WeightCompression(Algorithm):

Should I rename it to ExperimentalWeightCompression instead? since it could be confused with the original

It is inside the experimental directory, that should be descriptive enough. I suggest the WeightCompression name

daniil-lyakhov · 2025-09-23T15:29:23Z

src/nncf/experimental/quantization/algorithms/weight_compression/algorithm.py

+
+import torch
+
+import nncf  # type: ignore[import-untyped]


Why # type: ignore[import-untyped] here?

I need to update the typehint ignores since I copied them off my scripts

daniil-lyakhov · 2025-09-23T15:29:53Z

src/nncf/experimental/quantization/algorithms/weight_compression/algorithm.py

+    ) -> torch.fx.GraphModule:
+        self._quantizer = quantizer


typehints an docstring are missing

daniil-lyakhov · 2025-09-23T15:34:51Z

src/nncf/quantization/algorithms/weight_compression/algorithm.py

            model,
            parameters={
-                "mode": self._mode.value,
+                "mode": self._mode.value if not isinstance(self._mode, str) else self._mode,


What is the str mode here? Can we force self._mode to always be an ENUM param?

src/nncf/experimental/quantization/algorithms/weight_compression/algorithm.py

daniil-lyakhov

Can we just pass the dataset param to the quantizer.get_nncf_weight_compression_parameters and simplify the pipeline? With that, we don't need get_nodes_to_compress and collect_weight_compression_statistics methods in WC algorithm

daniil-lyakhov · 2025-09-24T13:03:24Z

src/nncf/experimental/torch/fx/quantization/quantizer/openvino_adapter.py

    def get_quantization_setup(self, model: torch.fx.GraphModule, nncf_graph: NNCFGraph) -> SingleConfigQuantizerSetup:
        return self._quantizer.get_nncf_quantization_setup(model, nncf_graph)
+
+    def get_weight_compression_setup(


Please do not use the word setup in context of WC

Suggested change

def get_weight_compression_setup(

def get_weight_compression_params(

tests/executorch/test_quantizer.py

src/nncf/experimental/quantization/algorithms/weight_compression/algorithm.py

src/nncf/quantization/algorithms/weight_compression/algorithm.py

…e static. Accepts data_aware_mixed_precision and data_aware_algo flags and mixed precision algorithm as input

src/nncf/quantization/algorithms/weight_compression/algorithm.py

daniil-lyakhov · 2025-10-07T12:39:40Z

src/nncf/quantization/algorithms/weight_compression/algorithm.py

-        for weight_param in primary_precision_weight_params:
-            weight_param.compression_config = self._get_primary_config(group_size_values[weight_param.weight_name])
-
+            for weight_param in ratio_defining_params:


Suggested change

for weight_param in ratio_defining_params:

# ratio_defining_params are all in primary precision. Update parameters which set to backup precision

# by the mixed precision algorithm.

for weight_param in ratio_defining_params:

daniil-lyakhov · 2025-10-07T12:40:24Z

src/nncf/quantization/algorithms/weight_compression/algorithm.py

+    ):
+        # Collect statistics for the weights compression


Description

daniil-lyakhov · 2025-10-07T12:41:04Z

src/nncf/quantization/algorithms/weight_compression/algorithm.py

-        :return: A tuple consisting of a list of weight compression parameters, based on the Weight
-            Compression algorithm configuration, and a mapping of target node names to the
-            collected statistics.
+        :return: A tuple consisting of a list of all weight compression parameters, based on the Weight


Please update the docstirng as discussed offline

daniil-lyakhov · 2025-10-09T08:41:12Z

src/nncf/experimental/quantization/algorithms/weight_compression/algorithm.py

+from nncf.common.utils.backend import BackendType
+from nncf.experimental.quantization.quantizer import Quantizer
+from nncf.quantization.algorithms.algorithm import Algorithm
+from nncf.quantization.algorithms.weight_compression.algorithm import WeightCompression


Suggested change

from nncf.quantization.algorithms.weight_compression.algorithm import WeightCompression

from nncf.quantization.algorithms.weight_compression.algorithm import WeightCompression as OriginalWeightCompression

Ah yes, good catch! I will change it.

…is is because mixed precision algo has clashes in ranking layer sensitivities with very few samples.

…on_params

anzr299 added 3 commits September 22, 2025 17:22

init

190f9d5

fixes

c52fcca

add message for unsupported external quantizers

4e56cb5

anzr299 requested a review from a team as a code owner September 22, 2025 14:43

github-actions bot added the API Public API-impacting changes label Sep 22, 2025

anzr299 marked this pull request as draft September 22, 2025 14:56

daniil-lyakhov self-requested a review September 22, 2025 15:03

daniil-lyakhov reviewed Sep 22, 2025

View reviewed changes

src/nncf/experimental/torch/fx/quantization/quantize_pt2e.py Outdated Show resolved Hide resolved

anzr299 added 19 commits September 22, 2025 19:27

add algorithm

9651ceb

impotr openvino quantizer from nncf instead of executorch

14daeb5

Add observers and openvino quantizer to nncf

3746815

fix

0815dc5

minor fix

1b8d940

fix

7d35374

fix some more bugs; observers was importing from torchao. causing mis…

427ebc2

…match in signatures in prepare_pt2e.

add compress pt2e to init

24dbfb6

fix quantizer init file. Remove extra code.

4bb8c1a

small fix for the big problem:)

8902842

fix quantizer preset definition

3842538

fix openvino quantizer for ptq. call _algo instead of legacy _min_max…

2e70c2e

…_algo

fix quantizer defaults

b1c9aad

microfix

33fe01c

precommit fix

d8e1006

revert openvino quantizer to old

88a8472

create ovquantizer in executorch dir

7a8e51a

update executorch quantizer location.

fed5052

check if openvino quantizer has weight compression in openvino adapter

2866473

daniil-lyakhov requested changes Sep 23, 2025

View reviewed changes

daniil-lyakhov reviewed Sep 23, 2025

View reviewed changes

src/nncf/experimental/quantization/algorithms/weight_compression/algorithm.py Outdated Show resolved Hide resolved

anzr299 added 5 commits September 24, 2025 15:26

fix last precommit

d12225a

remove unused mypy ignore

9870ee2

get the mode as struct

8015629

fix algorithm

0804218

remove quantizer and observers from nncf. Instead import from executorch

1f1fda3

daniil-lyakhov requested changes Sep 24, 2025

View reviewed changes

daniil-lyakhov reviewed Sep 24, 2025

View reviewed changes

src/nncf/quantization/algorithms/weight_compression/algorithm.py Outdated Show resolved Hide resolved

anzr299 added 4 commits October 1, 2025 12:12

rework wc algorithm so that get_weight_comrpession_params becomes mor…

623ce46

…e static. Accepts data_aware_mixed_precision and data_aware_algo flags and mixed precision algorithm as input

fix bugs; use sensitivity metric instead of mixed precision algo

d14a6eb

update algorithm with new reworking

e91b455

changes

448bf84

daniil-lyakhov reviewed Oct 6, 2025

View reviewed changes

src/nncf/quantization/algorithms/weight_compression/algorithm.py Outdated Show resolved Hide resolved

anzr299 added 12 commits October 6, 2025 21:27

review changes

8e23572

change WeightsCompressionPT2E to ExperimentalWeightsCompression

36ddf53

change ExperimentalWeightsCompression to WeightsCompression

07b730b

add comments

d5dd422

add typehints

076a76b

add docstrings

2ce9eec

add typehint for quantize pt2e

1bebf3e

Merge branch 'openvinotoolkit:develop' into an/fx/compress_pt2e

ea81cfd

return original develop branch changes

e82920f

update typehints and docs

82cc10b

format

beae508

update type hinting of openvino adapter

8bd95df

daniil-lyakhov reviewed Oct 7, 2025

View reviewed changes

daniil-lyakhov reviewed Oct 9, 2025

View reviewed changes

add test

aac9d3f

github-actions bot added the NNCF PT Pull requests that updates NNCF PyTorch label Oct 10, 2025

anzr299 added 2 commits October 10, 2025 18:25

update reference graphs; use more samples for calibration dataset. Th…

4278cfd

…is is because mixed precision algo has clashes in ranking layer sensitivities with very few samples.

remove groupsize values as return statement from get_weight_compressi…

6fd5216

…on_params

		from nncf.quantization.algorithms.weight_compression.algorithm import WeightCompression


		class WeightsCompressionPT2E(Algorithm):

	class WeightsCompressionPT2E(Algorithm):
	class WeightCompression(Algorithm):

	def get_weight_compression_setup(
	def get_weight_compression_params(

[Torch FX] Compress PT2E Support #3663

Are you sure you want to change the base?

[Torch FX] Compress PT2E Support #3663

Uh oh!

Conversation

anzr299 commented Sep 22, 2025

Changes

Reason for changes

Related tickets

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

daniil-lyakhov left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anzr299 Sep 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

daniil-lyakhov Sep 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

daniil-lyakhov left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

anzr299 Sep 24, 2025 •

edited

Loading

daniil-lyakhov Sep 24, 2025 •

edited

Loading

daniil-lyakhov left a comment •

edited

Loading