[Torch FX] Compress PT2E Support

.github/workflows/call_precommit.yml

@@ -136,6 +136,48 @@ jobs:
         env:
           NUM_WORKERS: 4
 
+  executorch:
+    timeout-minutes: 40
+    runs-on: ubuntu-latest-8-cores
+    defaults:
+      run:
+        shell: bash
+    env:
+      DEBIAN_FRONTEND: noninteractive
+    steps:
+      - name: Install dependencies
+        run : |
+          sudo apt-get update
+          sudo apt-get --assume-yes install gcc g++ build-essential ninja-build libgl1-mesa-dev libglib2.0-0
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
+        with:
+          lfs: true
+      - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0
+        with:
+          python-version: ${{ inputs.python_version }}
+      - name: Runner info
+        continue-on-error: true
+        run: |
+          cat /etc/*release
+          cat /proc/cpuinfo
+      - name: Override constraints
+        if: ${{ inputs.override_requirements != '' }}
+        run: python .github/scripts/override_constraints.py "${{ inputs.override_requirements }}"
+        shell: bash
+      - name: Install NNCF and test requirements
+        run: |
+          pip install . -r tests/executorch/requirements.txt
+          # Executorch
+          # Editable install due to https://github.com/pytorch/executorch/issues/6475
+          pip install --no-build-isolation -e git+https://github.com/anzr299/executorch.git@an/quantizer_nncf_pt2e_support#egg=executorch
+      - name: Print installed modules
+        run: pip list
+      - name: Run PyTorch precommit test scope
+        run: |
+          make test-executorch
+        env:
+          NUM_WORKERS: 4
+
   pytorch-cuda:
     timeout-minutes: 40
     runs-on: aks-linux-4-cores-28gb-gpu-tesla-t4

Makefile

@@ -141,6 +141,9 @@ test-torch-cpu:
 test-torch-cuda:
 	pytest ${COVERAGE_ARGS} tests/torch -ra -m "cuda and not weekly and not nightly and not models_hub and not legacy" --junitxml ${JUNITXML_PATH}
 
+test-executorch:
+	pytest ${COVERAGE_ARGS} tests/executorch --junitxml ${JUNITXML_PATH}
+
 test-torch-nightly:
 	pytest ${COVERAGE_ARGS} tests/torch -m "nightly or legacy" --junitxml ${JUNITXML_PATH} $(DATA_ARG)
 

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

@@ -0,0 +1,10 @@
+# Copyright (c) 2025 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

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

@@ -0,0 +1,141 @@
+# Copyright (c) 2025 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Iterable, Optional
+
+import torch
+
+from nncf import AdvancedCompressionParameters
+from nncf import CompressionFormat
+from nncf import CompressWeightsMode
+from nncf import Dataset
+from nncf import SensitivityMetric
+from nncf.common.graph.graph import NNCFGraph
+from nncf.common.graph.graph import NNCFNode
+from nncf.common.logging import nncf_logger
+from nncf.common.tensor_statistics.statistic_point import StatisticPointsContainer
+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 as OriginalWeightCompression
+
+
+class WeightsCompression(Algorithm):
+    """
+    Post-training Weight Compression algorithm implementation.
+
+    Compresses weights of Linear and Embedding layers to 8-bit integer or
+    to 4-bit integer/float depending on mode, ratio and group size.
+    """
+
+    def __init__(
+        self,
+        quantizer: Quantizer,
+        ratio: float,
+        subset_size: int,
+        awq: bool,
+        scale_estimation: bool,
+        gptq: bool,
+        lora_correction: bool,
+        sensitivity_metric: SensitivityMetric,
+        compression_format: CompressionFormat,
+        advanced_parameters: AdvancedCompressionParameters,
+    ) -> torch.fx.GraphModule:
+        """
+        :param quantizer: Quantizer to use in WeightCompression algorithm.
+        :param ratio: the ratio between primary and backup precisions (e.g. 0.9 means 90% of layers specified as
+            `ratio_defining_params` by the quantizer are quantized to INT4
+        :param subset_size: Number of data samples to calculate activation statistics used for assigning different
+            quantization precision.
+        :param awq: determines whether to use or not modified AWQ algorithm.
+        :param scale_estimation: determines whether to use or not scale estimation for 4 bit layers.
+        :param gptq: determines whether to use or not GPTQ algorithm.
+        :param lora_correction: determines whether to use or not LoRA Correction algorithm.
+        :param sensitivity_metric: The sensitivity metric for assigning quantization precision to layers. In order to
+            preserve the accuracy of the model, the more sensitive layers receives a higher precision.
+        :param compression_format: Describes the format in which the model is saved after weight compression.
+        :param advanced_parameters: advanced parameters for algorithms in compression pipeline.
+        """
+        self._quantizer = quantizer
+        wc_config = quantizer.get_weight_compression_config()
+
+        mode = wc_config.get("mode", CompressWeightsMode.INT8_ASYM)
+
+        self._algo = OriginalWeightCompression(
+            mode=CompressWeightsMode(mode),
+            ratio=ratio,
+            group_size=wc_config.get("group_size", None),
+            ignored_scope=None,
+            all_layers=wc_config.get("all_layers", None),
+            sensitivity_metric=sensitivity_metric,
+            awq=awq,
+            subset_size=subset_size,
+            scale_estimation=scale_estimation,
+            gptq=gptq,
+            lora_correction=lora_correction,
+            backup_mode=wc_config.get("backup_mode", None),
+            compression_format=compression_format,
+            advanced_parameters=advanced_parameters,
+        )
+
+    def available_backends(self) -> list[BackendType]:
+        return [BackendType.TORCH_FX]
+
+    def apply(
+        self,
+        model: torch.fx.GraphModule,
+        graph: NNCFGraph,
+        statistic_points: Optional[StatisticPointsContainer] = None,
+        dataset: Optional[Dataset] = None,
+    ) -> torch.fx.GraphModule:
+        self._algo.set_backend_entity(model)
+
+        all_weight_params, ratio_defining_params, skipped_weight_params = (
+            self._quantizer.get_weight_compression_parameters(model, graph)
+        )
+        # Collect statistics for the weights compression
+        statistics, statistic_points = self._algo.collect_statistics_and_statistic_points(
+            model, graph, statistic_points, dataset, ratio_defining_params, all_weight_params
+        )
+        # Apply Mixed precision algorithm to ratio defining parameters
+        self._algo.apply_mixed_precision(ratio_defining_params, model, graph, statistic_points)
+        self._algo.validate_group_size(ratio_defining_params)
+
+        # Print statistics
+        nncf_logger.info(
+            self._algo.get_bitwidth_distribution_str(all_weight_params, ratio_defining_params, skipped_weight_params)
+        )
+
+        # Filter all_weight_params by excluding nodes that should remain in their original floating-point precision
+        all_weight_params = [w_params for w_params in all_weight_params if w_params.compression_config is not None]
+        return self._algo.apply_with_parameters(
+            model,
+            graph,
+            dataset,
+            statistics,
+            all_weight_params,
+        )
+
+    def get_statistic_points(
+        self,
+        model: torch.fx.GraphModule,
+        graph: NNCFGraph,
+        nodes_and_port_ids: Iterable[tuple[NNCFNode, int]],
+    ) -> StatisticPointsContainer:
+        """
+        Returns statistic points, for which StatisticsCollector should collect statistics.
+
+        :param model: Model for statistics collection.
+        :param graph: Model graph.
+        :param nodes_and_port_ids: Nodes and port ids for which statistics should be collected.
+        :return: Statistic points, for which StatisticsCollector should collect statistics.
+        """
+        return self._algo.get_statistic_points(model, graph, nodes_and_port_ids)

src/nncf/experimental/quantization/quantizer.py

@@ -11,10 +11,11 @@
 
 from abc import ABC
 from abc import abstractmethod
-from typing import TypeVar
+from typing import Any, TypeVar
 
 from nncf.common.graph.graph import NNCFGraph
 from nncf.common.quantization.quantizer_setup import SingleConfigQuantizerSetup
+from nncf.quantization.algorithms.weight_compression.config import WeightCompressionParameters
 
 TModel = TypeVar("TModel")
 
@@ -43,3 +44,34 @@ def get_quantization_setup(self, model: TModel, nncf_graph: NNCFGraph) -> Single
         :param nncf_graph: NNCFGraph instance.
         :return: SingleConfigQuantizerSetup for the given model.
         """
+
+    @abstractmethod
+    def get_weight_compression_parameters(
+        self,
+        model: TModel,
+        nncf_graph: NNCFGraph,
+    ) -> tuple[
+        list[WeightCompressionParameters],
+        list[WeightCompressionParameters],
+        list[WeightCompressionParameters],
+    ]:
+        """
+        Obtains the weight compression parameters from the quantizer which can be used to determine
+        weights to compress, weights to skip, weights to consider for mixed precision assignment.
+
+        :param model: Backend-specific model.
+        :param nncf_graph: NNCFGraph instance.
+        :return: Tuple of (all_weight_params, ratio_defining_params, skipped_weight_params) where:
+            1. all_weight_params: all compressible weight parameters in the model
+            2. ratio_defining_params: subset of weights used for mixed precision assignment
+            3. skipped_weight_params: weights that should be excluded from compression
+        """
+
+    @abstractmethod
+    def get_weight_compression_config(self) -> dict[str, Any]:
+        """
+        Returns the weight compression configuration as a dictionary.
+
+        :return: Dictionary containing compression configuration parameters obtained
+            from the quantizer.
+        """

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

@@ -9,5 +9,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from nncf.experimental.torch.fx.quantization.quantize_pt2e import compress_pt2e as compress_pt2e
 from nncf.experimental.torch.fx.quantization.quantize_pt2e import quantize_pt2e as quantize_pt2e
 from nncf.experimental.torch.fx.quantization.quantizer.openvino_quantizer import OpenVINOQuantizer as OpenVINOQuantizer

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

@@ -22,11 +22,14 @@
 from torch.fx.passes.infra.pass_manager import PassManager
 
 import nncf
+from nncf import AdvancedCompressionParameters
 from nncf import Dataset
+from nncf import SensitivityMetric
 from nncf.common.factory import NNCFGraphFactory
 from nncf.common.logging import nncf_logger
 from nncf.common.utils.api_marker import api
 from nncf.experimental.quantization.algorithms.post_training.algorithm import ExperimentalPostTrainingQuantization
+from nncf.experimental.quantization.algorithms.weight_compression.algorithm import WeightsCompression
 from nncf.experimental.torch.fx.constant_folding import constant_fold
 from nncf.experimental.torch.fx.quantization.quantizer.openvino_adapter import OpenVINOQuantizerAdapter
 from nncf.experimental.torch.fx.quantization.quantizer.openvino_quantizer import OpenVINOQuantizer
@@ -157,3 +160,72 @@ def _quant_node_constraint(n: torch.fx.Node) -> bool:
     related to quantization
     """
     return n.op == "call_function" and n.target in QUANTIZE_NODE_TARGETS
+
+
+@api(canonical_alias="nncf.experimental.torch.fx.compress_pt2e")
+def compress_pt2e(
+    model: torch.fx.GraphModule,
+    quantizer: Quantizer,
+    *,
+    dataset: Optional[nncf.Dataset] = None,
+    awq: bool = False,
+    scale_estimation: bool = False,
+    gptq: bool = False,
+    lora_correction: bool = False,
+    subset_size: int = 128,
+    ratio: int = 1,
+    sensitivity_metric: Optional[SensitivityMetric] = None,
+    advanced_parameters: Optional[AdvancedCompressionParameters] = None,
+) -> torch.fx.GraphModule:
+    """
+    Applies Weight Compression to the torch.fx.GraphModule model using provided torch.ao quantizer.
+
+    :param model: A torch.fx.GraphModule instance to be quantized.
+    :param quantizer: Torch ao quantizer to annotate nodes in the graph with quantization setups
+        to convey the desired way of quantization.
+    :param dataset: A representative dataset for the calibration process.
+    :param awq: Determines whether to use or not the modified AWQ algorithm.
+    :param scale_estimation: Determines whether to use or not scale estimation for 4-bit layers.
+    :param gptq: Determines whether to use or not GPTQ algorithm.
+    :param lora_correction: Determines whether to use or not LoRA Correction algorithm.
+    :param subset_size: Number of data samples to calculate activation statistics used for assigning different
+        quantization precision.
+    :param ratio: the ratio between baseline and backup precisions (e.g. 0.9 means 90% of layers quantized to NF4
+        and the rest to INT8_ASYM).
+    :param sensitivity_metric: The sensitivity metric for assigning quantization precision to layers. In order to
+        preserve the accuracy of the model, the more sensitive layers receive a higher precision.
+    :param advanced_parameters: Advanced parameters for algorithms in the compression pipeline.
+    """
+    if isinstance(quantizer, OpenVINOQuantizer) or hasattr(quantizer, "get_nncf_weight_compression_parameters"):
+        quantizer = OpenVINOQuantizerAdapter(quantizer)
+        compression_format = nncf.CompressionFormat.DQ
+    else:
+        # TODO Support Third party quantizers here.
+        msg = "Only OpenVINO Quantizer is supported currently."
+        raise nncf.InternalError(msg)
+
+    sensitivity_metric = (
+        (SensitivityMetric.WEIGHT_QUANTIZATION_ERROR if dataset is None else SensitivityMetric.MAX_ACTIVATION_VARIANCE)
+        if sensitivity_metric is None
+        else sensitivity_metric
+    )
+
+    quantization_algorithm = WeightsCompression(
+        quantizer=quantizer,
+        subset_size=subset_size,
+        compression_format=compression_format,
+        ratio=ratio,
+        awq=awq,
+        scale_estimation=scale_estimation,
+        gptq=gptq,
+        lora_correction=lora_correction,
+        sensitivity_metric=sensitivity_metric,
+        advanced_parameters=advanced_parameters,
+    )
+
+    # Here the model is annotated
+    transformed_model = quantizer.transform_prior_quantization(model)
+    nncf_graph = NNCFGraphFactory.create(transformed_model)
+    quantized_model = quantization_algorithm.apply(transformed_model, nncf_graph, dataset=dataset)
+    quantized_model = torch.fx.GraphModule(quantized_model, graph=quantized_model.graph)
+    return quantized_model

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

@@ -9,12 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any
+
 import torch.fx
 
 from nncf.common.graph.graph import NNCFGraph
 from nncf.common.quantization.quantizer_setup import SingleConfigQuantizerSetup
 from nncf.experimental.quantization.quantizer import Quantizer
 from nncf.experimental.torch.fx.quantization.quantizer.openvino_quantizer import OpenVINOQuantizer
+from nncf.quantization.algorithms.weight_compression.config import WeightCompressionParameters
 
 
 class OpenVINOQuantizerAdapter(Quantizer):
@@ -30,3 +33,17 @@ def transform_prior_quantization(self, model: torch.fx.GraphModule) -> torch.fx.
 
     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_parameters(
+        self,
+        model: torch.fx.GraphModule,
+        nncf_graph: NNCFGraph,
+    ) -> tuple[
+        list[WeightCompressionParameters],
+        list[WeightCompressionParameters],
+        list[WeightCompressionParameters],
+    ]:
+        return self._quantizer.get_nncf_weight_compression_parameters(model, nncf_graph)
+
+    def get_weight_compression_config(self) -> dict[str, Any]:
+        return self._quantizer.weight_compression_configuration