Repository: openvinotoolkit/nncf
Pull request #3409: [FX][AWQ][Scale Estimation][Mixed Precision] Add Data Aware Algorithm Support for FX Backend
Author: login: anzr299
Labels: name: NNCF PT
description: Pull requests that updates NNCF PyTorch
name: NNCF PTQ
description: Pull requests that updates NNCF PTQ
Created: 2025-04-07T14:41:06Z | Merged: 2025-04-25T10:23:46Z

### Changes

The following algorithms were added for Torch FX backend:
1. AWQ
2. Scale Estimation
3. Data-Aware Mixed Precision

### Tests

`TemplateWeightCompression` was extended for torch FX backend for the above algorithms in `tests/torch/fx/test_compress_weights.py`

Category | Job | Status | Job Number | Notes
-- | -- | -- | -- | --
GH WC | WC test on GHA | Running | <a href="https://github.com/openvinotoolkit/nncf/actions/runs/14635657909">70</a> |

### Conformance Test Results

backend | INT8| INT4 | data_aware | scale_estimation_per_channel | data_aware_awq_scale_estimation |
-- | -- | -- | -- | -- | --
FX | 0.95624 | 0.73873 | 0.85811 | 0.80873 | 0.85502
PT | 0.95624 | 0.73873 | 0.85811 | 0.80873 | 0.85502
OpenVINO | 0.95624 | 0.73873 | 0.85811 | 0.80873 | 0.85502

### Related Tickets

165013



Diff:
diff --git a/.ci/cspell_dict.txt b/.ci/cspell_dict.txt
index fbe898113f8..738a11a18c6 100644
--- a/.ci/cspell_dict.txt
+++ b/.ci/cspell_dict.txt
@@ -148,6 +148,7 @@ frobenius
 fsolve
 funcs
 fval
+fxawq
 fxsq
 gacts
 gelsy
diff --git a/nncf/quantization/algorithms/weight_compression/awq.py b/nncf/quantization/algorithms/weight_compression/awq.py
index 9e0b8ab85a2..3c614d228d9 100644
--- a/nncf/quantization/algorithms/weight_compression/awq.py
+++ b/nncf/quantization/algorithms/weight_compression/awq.py
@@ -106,7 +106,10 @@ def _set_backend_entity(
             from nncf.quantization.algorithms.weight_compression.torch_backend import PTAWQAlgoAlgoBackend
 
             self._backend_entity = PTAWQAlgoAlgoBackend()
+        elif model_backend == BackendType.TORCH_FX:
+            from nncf.quantization.algorithms.weight_compression.torch_fx_backend import FXAWQAlgoAlgoBackend
 
+            self._backend_entity = FXAWQAlgoAlgoBackend()
         else:
             msg = f"Cannot return backend-specific AWQ entity because {model_backend.value} is not supported!"
             raise nncf.UnsupportedBackendError(msg)
diff --git a/nncf/quantization/algorithms/weight_compression/mixed_precision.py b/nncf/quantization/algorithms/weight_compression/mixed_precision.py
index 38459c9d605..6bb9391d0cb 100644
--- a/nncf/quantization/algorithms/weight_compression/mixed_precision.py
+++ b/nncf/quantization/algorithms/weight_compression/mixed_precision.py
@@ -232,6 +232,10 @@ def _set_backend_entity(self, model: TModel) -> None:
             from nncf.quantization.algorithms.weight_compression.torch_backend import PTMixedPrecisionAlgoBackend
 
             self._backend_entity = PTMixedPrecisionAlgoBackend()
+        elif model_backend == BackendType.TORCH_FX:
+            from nncf.quantization.algorithms.weight_compression.torch_fx_backend import FXMixedPrecisionAlgoBackend
+
+            self._backend_entity = FXMixedPrecisionAlgoBackend()
         else:
             msg = f"Cannot return backend-specific entity because {model_backend.value} is not supported!"
             raise nncf.UnsupportedBackendError(msg)
diff --git a/nncf/quantization/algorithms/weight_compression/scale_estimation.py b/nncf/quantization/algorithms/weight_compression/scale_estimation.py
index 4792637b6c6..4aea4633ebb 100644
--- a/nncf/quantization/algorithms/weight_compression/scale_estimation.py
+++ b/nncf/quantization/algorithms/weight_compression/scale_estimation.py
@@ -81,6 +81,10 @@ def _set_backend_entity(self, model: TModel) -> None:
             from nncf.quantization.algorithms.weight_compression.torch_backend import PTWeightCompressionAlgoBackend
 
             self._backend_entity = PTWeightCompressionAlgoBackend()
+        elif model_backend == BackendType.TORCH_FX:
+            from nncf.quantization.algorithms.weight_compression.torch_fx_backend import FXWeightCompressionAlgoBackend
+
+            self._backend_entity = FXWeightCompressionAlgoBackend()
         else:
             msg = (
                 "Cannot return backend-specific Scale Estimation entity because"
diff --git a/nncf/quantization/algorithms/weight_compression/torch_fx_backend.py b/nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
index 156c1c76708..2650f16600c 100644
--- a/nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
+++ b/nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
@@ -37,9 +37,14 @@
 from nncf.parameters import CompressionFormat
 from nncf.parameters import CompressWeightsMode
 from nncf.quantization.advanced_parameters import AdvancedCompressionParameters
+from nncf.quantization.algorithms.weight_compression.backend import AWQAlgoBackend
+from nncf.quantization.algorithms.weight_compression.backend import MixedPrecisionAlgoBackend
 from nncf.quantization.algorithms.weight_compression.backend import WeightCompressionAlgoBackend
 from nncf.quantization.algorithms.weight_compression.config import WeightCompressionParameters
+from nncf.quantization.algorithms.weight_compression.handle_errors import handle_invalid_group_size_error
 from nncf.quantization.algorithms.weight_compression.lora_correction import LoraCorrectionAlgorithm
+from nncf.quantization.algorithms.weight_compression.torch_backend import PTAWQAlgoAlgoBackend
+from nncf.quantization.algorithms.weight_compression.torch_backend import PTMixedPrecisionAlgoBackend
 from nncf.quantization.algorithms.weight_compression.torch_backend import PTWeightCompressionAlgoBackend
 from nncf.quantization.algorithms.weight_compression.weight_lowering import compress_weight
 from nncf.tensor import Tensor
@@ -133,12 +138,11 @@ def get_activation_port_id(node: NNCFNode, graph: NNCFGraph) -> int:
     def get_weight(
         self, node_with_weight: NNCFNode, weight_port_id: int, model: torch.fx.GraphModule, graph: NNCFGraph
     ) -> Tensor:
-        weight_edge = graph.get_input_edge_by_port_id(node_with_weight, weight_port_id)
-        weight_node = weight_edge.from_node
-        graph_weight_node = get_graph_node_by_name(model.graph, weight_node.node_name)
+        graph_node_with_weight = get_graph_node_by_name(model.graph, node_with_weight.node_name)
+        graph_weight_node = graph_node_with_weight.all_input_nodes[weight_port_id]
         weight = get_tensor_constant_from_node(graph_weight_node, model).data
         if weight is None:
-            msg = f"Could not find a node in the model by name {weight_node}."
+            msg = f"Could not find a node in the model by name {graph_weight_node}."
             raise nncf.InternalError(msg)
 
         return Tensor(weight)
@@ -192,7 +196,8 @@ def transform_model(
         advanced_parameters: AdvancedCompressionParameters = AdvancedCompressionParameters(),
     ) -> torch.fx.GraphModule:
         transformation_layout = TransformationLayout()
-
+        invalid_node_names = []
+        first_caught_error = None
         for wc_params in weight_compression_parameters:
             compression_config = wc_params.compression_config
             if compression_config.mode in [
@@ -207,15 +212,19 @@ def transform_model(
             if weight is None or not isinstance(weight, Tensor):
                 msg = f"Could not find a nncf.tensor in the model by name {weight_name}."
                 raise nncf.InternalError(msg)
-
-            # calculates compressed weights and decompression parameters
-            compressed_weight = compress_weight(
-                weight,
-                wc_params.reduction_axes,
-                compression_config,
-                None if precomputed_scales is None else precomputed_scales.get(wc_params.weight_name),
-                None if precomputed_zero_points is None else precomputed_zero_points.get(wc_params.weight_name),
-            )
+            try:
+                # calculates compressed weights and decompression parameters
+                compressed_weight = compress_weight(
+                    weight,
+                    wc_params.reduction_axes,
+                    compression_config,
+                    None if precomputed_scales is None else precomputed_scales.get(wc_params.weight_name),
+                    None if precomputed_zero_points is None else precomputed_zero_points.get(wc_params.weight_name),
+                )
+            except nncf.InvalidGroupSizeError as error:
+                first_caught_error = error
+                invalid_node_names.append(wc_params.node_with_weight.node_name)
+                continue
 
             # creates weight decompressor
             if compression_config.mode == CompressWeightsMode.INT8_SYM:
@@ -265,8 +274,73 @@ def transform_model(
                     )
                 )
             )
-
+        if first_caught_error:
+            handle_invalid_group_size_error(first_caught_error, invalid_node_names)
         # apply transformations
         transformed_model = FXModelTransformer(model).transform(transformation_layout)
 
         return transformed_model
+
+
+class FXMixedPrecisionAlgoBackend(MixedPrecisionAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
+    def mean_variance_statistic_collector(
+        reduction_axes: tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.mean_variance_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+    @staticmethod
+    def max_variance_statistic_collector(
+        reduction_axes: tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.max_variance_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+    @staticmethod
+    def mean_abs_max_statistic_collector(
+        reduction_axes: tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.mean_abs_max_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+
+class FXAWQMultiply(torch.nn.Module):
+    def __init__(self, scale: torch.Tensor):
+        super().__init__()
+        self.register_buffer("_scale_value", scale)
+        self._scale_value: torch.Tensor
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.mul(x, self._scale_value)
+
+
+class FXAWQAlgoAlgoBackend(AWQAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
+    def get_awq_patterns():
+        return PTAWQAlgoAlgoBackend.get_awq_patterns()
+
+    @staticmethod
+    def scale_insertion_command(source_node, next_nodes, source_node_output_port, scale):
+        input_port_id = 0
+        target_points = []
+        for node in next_nodes:
+            target_points.append(
+                PTTargetPoint(
+                    TargetType.OPERATOR_PRE_HOOK,
+                    node.node_name,
+                    input_port_id=input_port_id,
+                )
+            )
+        awq_multiply = FXAWQMultiply(scale)
+        awq_node_name = f"{source_node.node_name}/awq_mul"
+        return FXApplyTransformationCommand(
+            module_insertion_transformation_builder(
+                awq_multiply,
+                target_points,
+                awq_node_name,
+            )
+        )
diff --git a/nncf/quantization/quantize_model.py b/nncf/quantization/quantize_model.py
index c5ec1caa026..41a90b212cd 100644
--- a/nncf/quantization/quantize_model.py
+++ b/nncf/quantization/quantize_model.py
@@ -8,6 +8,7 @@
 # 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.
+import warnings
 from typing import Any, Callable, Iterable, Optional, TypedDict, TypeVar, Union
 
 import nncf
@@ -564,8 +565,6 @@ def compress_weights(
             raise nncf.ParameterNotSupportedError(msg)
 
         options = {
-            "awq": awq,
-            "scale_estimation": scale_estimation,
             "gptq": gptq,
             "lora_correction": lora_correction,
         }
@@ -574,16 +573,6 @@ def compress_weights(
             msg = f"TorchFX backend does not support {', '.join(unsupported_options)} option(s). Set them to None."
             raise nncf.ParameterNotSupportedError(msg)
 
-        if sensitivity_metric not in [None, SensitivityMetric.WEIGHT_QUANTIZATION_ERROR]:
-            msg = (
-                "TorchFX backend only supports data-free sensitivity metric. "
-                "Set None or SensitivityMetric.WEIGHT_QUANTIZATION_ERROR."
-            )
-            raise nncf.ParameterNotSupportedError(msg)
-
-        if dataset:
-            msg = "TorchFX only supports data-free weights compression. Set the 'dataset' option to None"
-            raise nncf.ParameterNotSupportedError(msg)
         if advanced_parameters and advanced_parameters.statistics_path:
             msg = "TorchFX does not supports statistics caching."
             raise nncf.ParameterNotSupportedError(msg)
@@ -592,6 +581,13 @@ def compress_weights(
             msg = "Torch FX backend does not support FQ and FQ_LORA compression formats."
             raise nncf.ParameterNotSupportedError(msg)
 
+        if (
+            mode in (CompressWeightsMode.INT8, CompressWeightsMode.INT8_ASYM, CompressWeightsMode.INT8_SYM)
+            and dataset is not None
+        ):
+            warnings.warn("data-aware methods are not supported in INT8 modes. dataset is set to None")
+            dataset = None
+
         compression_weights_impl = fx_compression_weights_impl
 
     elif backend == BackendType.OPENVINO:
diff --git a/tests/post_training/data/wc_reference_data.yaml b/tests/post_training/data/wc_reference_data.yaml
index 8ea75b4d236..1a3e74e7d00 100644
--- a/tests/post_training/data/wc_reference_data.yaml
+++ b/tests/post_training/data/wc_reference_data.yaml
@@ -61,3 +61,44 @@ tinyllama_int4_data_free_backend_TORCH:
   metric_value: 0.73873
   num_int4: 114
   num_int8: 84
+tinyllama_int8_data_free_backend_FX_TORCH:
+  metric_value: 0.95624
+  num_int4: 0
+  num_int8: 312
+  exception_xfail_reason:
+    type: "TypeError"
+    error_message: "'function' object is not iterable"
+    message: "Issue-165013"
+tinyllama_int4_data_free_backend_FX_TORCH:
+  metric_value: 0.73873
+  num_int4: 114
+  num_int8: 84
+  exception_xfail_reason:
+    type: "TypeError"
+    error_message: "'function' object is not iterable"
+    message: "Issue-165013"
+tinyllama_data_aware_backend_FX_TORCH:
+  metric_value: 0.85811
+  num_int4: 94
+  num_int8: 124
+  exception_xfail_reason:
+    type: "TypeError"
+    error_message: "'function' object is not iterable"
+    message: "Issue-165013"
+tinyllama_scale_estimation_per_channel_backend_FX_TORCH:
+  metric_value: 0.80873
+  num_int4: 188
+  num_int8: 124
+  exception_xfail_reason:
+    type: "TypeError"
+    error_message: "'function' object is not iterable"
+    message: "Issue-165013"
+  atol: 0.006 # difference across devices: 0.80873 vs 0.81389
+tinyllama_data_aware_awq_scale_estimation_backend_FX_TORCH:
+  metric_value: 0.85502
+  num_int4: 94
+  num_int8: 124
+  exception_xfail_reason:
+    type: "TypeError"
+    error_message: "'function' object is not iterable"
+    message: "Issue-165013"
diff --git a/tests/post_training/model_scope.py b/tests/post_training/model_scope.py
index 92d08f067ad..e25912c18e0 100644
--- a/tests/post_training/model_scope.py
+++ b/tests/post_training/model_scope.py
@@ -409,7 +409,7 @@
         "model_id": "tinyllama/tinyllama-1.1b-step-50k-105b",
         "pipeline_cls": LMWeightCompression,
         "compression_params": {"group_size": 64, "ratio": 0.8, "mode": CompressWeightsMode.INT4_SYM},
-        "backends": [BackendType.OV, BackendType.TORCH],
+        "backends": [BackendType.OV, BackendType.TORCH, BackendType.FX_TORCH],
     },
     {
         "reported_name": "tinyllama_data_aware_awq_stateful",
@@ -433,7 +433,7 @@
                 scale_estimation_params=AdvancedScaleEstimationParameters(32, 5, 10, 1.0)
             ),
         },
-        "backends": [BackendType.OV, BackendType.TORCH],
+        "backends": [BackendType.OV, BackendType.TORCH, BackendType.FX_TORCH],
     },
     {
         "reported_name": "tinyllama_data_aware_awq_scale_estimation_stateful",
@@ -459,7 +459,7 @@
         "compression_params": {
             "mode": CompressWeightsMode.INT8_ASYM,
         },
-        "backends": [BackendType.TORCH],
+        "backends": [BackendType.TORCH, BackendType.FX_TORCH],
     },
     {
         "reported_name": "tinyllama_int4_data_free",
@@ -471,7 +471,7 @@
             "mode": CompressWeightsMode.INT4_SYM,
             "sensitivity_metric": SensitivityMetric.WEIGHT_QUANTIZATION_ERROR,
         },
-        "backends": [BackendType.TORCH],
+        "backends": [BackendType.TORCH, BackendType.FX_TORCH],
     },
     {
         "reported_name": "tinyllama_data_aware_gptq_scale_estimation_stateful",
@@ -516,7 +516,7 @@
             "mode": CompressWeightsMode.INT4_ASYM,
             "scale_estimation": True,
         },
-        "backends": [BackendType.OV, BackendType.TORCH],
+        "backends": [BackendType.OV, BackendType.TORCH, BackendType.FX_TORCH],
     },
     {
         "reported_name": "tinyllama_data_aware_lora_stateful",
diff --git a/tests/post_training/pipelines/fx_utils.py b/tests/post_training/pipelines/fx_utils.py
new file mode 100644
index 00000000000..224a050fd15
--- /dev/null
+++ b/tests/post_training/pipelines/fx_utils.py
@@ -0,0 +1,170 @@
+# 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.
+
+import torch
+import torch.fx
+from optimum.modeling_base import OptimizedModel
+from transformers import GenerationConfig
+from transformers import GenerationMixin
+from transformers import PretrainedConfig
+from transformers import PreTrainedModel
+from transformers.cache_utils import StaticCacheConfig
+from transformers.integrations.executorch import TorchExportableModuleWithStaticCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class FXAutoModelForCausalLM(OptimizedModel, GenerationMixin):
+    def __init__(
+        self,
+        model: torch.fx.GraphModule,
+        config: PretrainedConfig,
+        device: str = "cpu",
+        compile: bool = True,
+        backend: str = "openvino",
+        dtype=torch.float32,
+    ):
+        super().__init__(model, config)
+        self.generation_config = GenerationConfig.from_model_config(self.config)
+        self.main_input_name = "input_ids"
+        self.compile = compile
+        self.backend = backend
+        self._device = device.upper()
+        self._dtype = dtype
+        self._cached_prefill_input_ids = None
+        self._cached_cache_position = None
+        self.compiled_decode = False
+        self.compiled_prefill = False
+        if self.compile:
+            self.prefill = None
+            self.decode_one_token = None
+
+    def get_openvino_backend_options(self) -> dict:
+        return {
+            "device": self._device,
+            "aot_autograd": True,
+        }
+
+    def get_prefill(self, input_ids: torch.Tensor, cache_position: torch.Tensor):
+        if (
+            self.prefill is None
+            or self._cached_prefill_input_ids.shape != input_ids.shape
+            or self._cached_cache_position.shape != cache_position.shape
+        ):
+            self._cached_prefill_input_ids = input_ids
+            self._cached_cache_position = cache_position
+            self.prefill = self.model
+            self.prefill.forward = torch.compile(
+                self.prefill.forward,
+                backend="openvino",
+                options=self.get_openvino_backend_options(),
+            )
+        return self.prefill
+
+    def get_decode_one_token(self, input_ids: torch.Tensor, cache_position: torch.Tensor):
+        if self.decode_one_token is None:
+            self.decode_one_token = self.model
+            self.decode_one_token.forward = torch.compile(
+                self.decode_one_token.forward,
+                backend="openvino",
+                options=self.get_openvino_backend_options(),
+            )
+        return self.decode_one_token
+
+    def infer_prefill(self, input_ids: torch.Tensor, cache_position: torch.Tensor):
+        if self.compile:
+            _ = self.get_prefill(input_ids, cache_position)(input_ids, cache_position)
+        else:
+            self.model(input_ids, cache_position)
+
+    def infer_decode_one_token(self, input_ids: torch.Tensor, cache_position: torch.Tensor):
+        if self.compile:
+            _ = self.get_decode_one_token(input_ids, cache_position)(input_ids, cache_position)
+        else:
+            self.model(input_ids, cache_position)
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device(self._device.lower())
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        cache_position = kwargs["cache_position"]
+        past_len = cache_position[0]
+        if past_len < input_ids.shape[1]:
+            input_ids = input_ids[:, past_len:]
+
+        return {"input_ids": input_ids, "cache_position": cache_position}
+
+    def _save_pretrained(self, save_directory):
+        pass
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        cache_position: torch.Tensor,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        if self.compile:
+            if input_ids.shape[1] == 1:
+                logits = self.get_decode_one_token(input_ids, cache_position)(input_ids, cache_position)
+            else:
+                logits = self.get_prefill(input_ids, cache_position)(input_ids, cache_position)
+        else:
+            logits = self.model(input_ids, cache_position)
+        return CausalLMOutputWithPast(logits=logits)
+
+    def can_generate(self):
+        return True
+
+    def _supports_default_dynamic_cache(self) -> bool:
+        return False
+
+
+class TorchExportableModuleWithStaticCacheDynamicShape(TorchExportableModuleWithStaticCache):
+    def forward(self, input_ids: torch.Tensor, cache_position: torch.Tensor):
+        outs = self.model(
+            input_ids=input_ids,
+            position_ids=cache_position.unsqueeze(0),
+            cache_position=cache_position,
+            past_key_values=self.static_cache,
+            use_cache=True,
+        )
+        return outs.logits
+
+
+def convert_and_export_with_cache(model: PreTrainedModel, re_export=False):
+    """
+    Convert a `PreTrainedModel` into an exportable module and export it using `torch.export`
+    or `torch._export.capture_pre_autograd_graph`.
+    """
+
+    with torch.no_grad():
+        example_input_ids = torch.ones(1, 8, dtype=torch.long)
+        example_cache_position = torch.arange(0, 8, dtype=torch.long)
+        model_config = None
+        if not re_export:
+            model.generation_config.cache_implementation = "static"
+            model.generation_config.cache_config = StaticCacheConfig(batch_size=1, max_cache_len=512)
+            model(example_input_ids)
+            model_config = model.config
+            model = TorchExportableModuleWithStaticCacheDynamicShape(model)
+
+        sequence_length = torch.export.Dim("sequence_length", min=1, max=512)
+        dynamic_shapes = {"input_ids": {1: sequence_length}, "cache_position": {0: sequence_length}}
+
+        exported_program = torch.export.export_for_training(
+            model,
+            args=(
+                example_input_ids,
+                example_cache_position,
+            ),
+            dynamic_shapes=dynamic_shapes,
+        ).run_decompositions(decomp_table={})
+        return exported_program, model_config
diff --git a/tests/post_training/pipelines/lm_weight_compression.py b/tests/post_training/pipelines/lm_weight_compression.py
index 0c7a89ae45a..dd8440d9fb7 100644
--- a/tests/post_training/pipelines/lm_weight_compression.py
+++ b/tests/post_training/pipelines/lm_weight_compression.py
@@ -38,6 +38,8 @@
 from tests.post_training.pipelines.base import RunInfo
 from tests.post_training.pipelines.base import StatsFromOutput
 from tests.post_training.pipelines.base import get_num_fq_int4_int8
+from tests.post_training.pipelines.fx_utils import FXAutoModelForCausalLM
+from tests.post_training.pipelines.fx_utils import convert_and_export_with_cache
 from tools.memory_monitor import MemoryType
 from tools.memory_monitor import MemoryUnit
 from tools.memory_monitor import memory_monitor_context
@@ -127,9 +129,9 @@ def prepare_model(self) -> None:
         is_stateful = self.params.get("is_stateful", False)
 
         # load model
-        if self.backend == BackendType.TORCH:
+        if self.backend in [BackendType.TORCH, BackendType.FX_TORCH]:
             if is_stateful:
-                msg = f"is_stateful={is_stateful} is not supported for PyTorch backend."
+                msg = f"is_stateful={is_stateful} is not supported for {self.backend.value} backend."
                 raise RuntimeError(msg)
 
             self.model_hf = AutoModelForCausalLM.from_pretrained(
@@ -138,6 +140,9 @@ def prepare_model(self) -> None:
                 device_map="cpu",  # TODO (kshpv): add support of 'cuda', when supported
             )
             self.model = self.model_hf
+            if self.backend == BackendType.FX_TORCH:
+                self.model, self.model_config = convert_and_export_with_cache(self.model)
+                self.model = self.model.module()
         elif self.backend == BackendType.OV:
             if is_stateful:
                 self.fp32_model_dir = self.fp32_model_dir.parent / (self.fp32_model_dir.name + "_sf")
@@ -209,6 +214,11 @@ def transform_fn(data, max_tokens=128, filter_bad_tokens=True):
             if self.backend == BackendType.TORCH:
                 for input_name in inputs:
                     inputs[input_name] = torch.from_numpy(inputs[input_name]).to(self.model_hf.device)
+            elif self.backend == BackendType.FX_TORCH:
+                fx_inputs = ()
+                fx_inputs += (torch.from_numpy(inputs["input_ids"]).to(self.model_hf.device),)
+                fx_inputs += (torch.from_numpy(inputs["position_ids"]).to(self.model_hf.device).squeeze(0),)
+                inputs = fx_inputs
             return inputs
 
         return transform_fn
@@ -270,6 +280,25 @@ def save_compressed_model(self) -> None:
                 compression_option="fp32",
                 device=self.model_hf.device,
             )
+        elif self.backend == BackendType.FX_TORCH:
+            with torch.no_grad():
+                example_input_ids = torch.ones(1, 8, dtype=torch.long)
+                example_cache_position = torch.arange(0, 8, dtype=torch.long)
+                torch.compile(
+                    self.compressed_model,
+                    backend="openvino",
+                    options={"aot_autograd": True, "model_caching": True, "cache_dir": str(self.output_model_dir)},
+                )(example_input_ids, example_cache_position)
+
+                # Get the OV *.xml files in torch compile cache directory
+                cached_ov_model_files = list(Path(self.output_model_dir / "model").glob("*.xml"))
+                if len(cached_ov_model_files) > 1:
+                    msg = "Graph break encountered in torch compile!"
+                    raise nncf.InternalError(msg)
+                elif len(cached_ov_model_files) == 0:
+                    msg = "Openvino Model Files Not Found!"
+                    raise FileNotFoundError(msg)
+                self.path_compressed_ir = cached_ov_model_files[0]
 
     def run_bench(self) -> None:
         pass
@@ -330,7 +359,7 @@ def _validate(self) -> None:
             )
 
         compressed_model_hf = self.model_hf
-        if self.backend != BackendType.FP32:
+        if self.backend != BackendType.FP32 and self.backend != BackendType.FX_TORCH:
             compressed_model_hf = OVModelForCausalLM.from_pretrained(
                 self.output_model_dir,
                 trust_remote_code=True,
@@ -339,6 +368,8 @@ def _validate(self) -> None:
                 stateful=is_stateful,
                 ov_config={"DYNAMIC_QUANTIZATION_GROUP_SIZE": "0", "KV_CACHE_PRECISION": "f16"},
             )
+        if self.backend == BackendType.FX_TORCH:
+            compressed_model_hf = FXAutoModelForCausalLM(self.model, self.model_config)
         print("Evaluation of the target model")
         _, all_metrics = evaluator.score(compressed_model_hf)
         similarity = all_metrics["similarity"][0]
diff --git a/tests/torch/fx/test_compress_weights.py b/tests/torch/fx/test_compress_weights.py
index 4b3b5bd2c29..8b3b7c11227 100644
--- a/tests/torch/fx/test_compress_weights.py
+++ b/tests/torch/fx/test_compress_weights.py
@@ -18,24 +18,34 @@
 from nncf import CompressWeightsMode
 from nncf import SensitivityMetric
 from nncf.common.factory import NNCFGraphFactory
-from nncf.data.dataset import Dataset
 from nncf.experimental.torch.fx.node_utils import get_tensor_constant_from_node
+from nncf.experimental.torch.fx.transformations import get_graph_node_by_name
 from nncf.parameters import CompressionFormat
 from nncf.quantization import compress_weights
 from nncf.quantization.advanced_parameters import AdvancedCompressionParameters
+from nncf.quantization.algorithms.weight_compression.torch_fx_backend import FXAWQMultiply
+from nncf.tensor import Tensor
+from nncf.tensor import TensorDataType
 from nncf.torch.dynamic_graph.patch_pytorch import disable_patching
+from nncf.torch.quantization.layers import INT4SymmetricWeightsDecompressor
+from tests.cross_fw.test_templates.template_test_weights_compression import TemplateWeightCompression
 from tests.torch.fx.helpers import get_torch_fx_model
 from tests.torch.test_models.synthetic import ShortTransformer
+from tests.torch.test_tensor import cast_to
 from tests.torch2.function_hook.quantization.test_weights_compression import ALL_SENSITIVITY_METRICS
 from tests.torch2.function_hook.quantization.test_weights_compression import INT4_MODES
 from tests.torch2.function_hook.quantization.test_weights_compression import INT8_MODES
 from tests.torch2.function_hook.quantization.test_weights_compression import SUPPORTED_MODES
 from tests.torch2.function_hook.quantization.test_weights_compression import UNSUPPORTED_MODES
+from tests.torch2.function_hook.quantization.test_weights_compression import AWQActLinearModel
+from tests.torch2.function_hook.quantization.test_weights_compression import AWQLinearModel
 from tests.torch2.function_hook.quantization.test_weights_compression import ConvolutionModel
 from tests.torch2.function_hook.quantization.test_weights_compression import DTypeModel
 from tests.torch2.function_hook.quantization.test_weights_compression import EmptyModel
 from tests.torch2.function_hook.quantization.test_weights_compression import FunctionalModel
+from tests.torch2.function_hook.quantization.test_weights_compression import LinearModel
 from tests.torch2.function_hook.quantization.test_weights_compression import MatMulModel
+from tests.torch2.function_hook.quantization.test_weights_compression import SequentialMatmulModel
 
 DATA_BASED_SENSITIVITY_METRICS = (
     SensitivityMetric.HESSIAN_INPUT_ACTIVATION,
@@ -221,10 +231,8 @@ def test_compress_weights_functional_model(mode):
         {"all_layers": False},
         *({"sensitivity_metric": metric} for metric in ALL_SENSITIVITY_METRICS),
         {"gptq": True},
-        {"awq": True},
         {"scale_estimation": True},
         {"lora_correction": True},
-        {"dataset": Dataset([1])},
         {"backup_mode": BackupMode.NONE},
         {"backup_mode": BackupMode.INT8_ASYM},
         {"backup_mode": BackupMode.INT8_SYM},
@@ -245,12 +253,8 @@ def test_raise_error_with_unsupported_params_for_int8(mode, params):
 @pytest.mark.parametrize(
     "params",
     (
-        *({"sensitivity_metric": metric} for metric in DATA_BASED_SENSITIVITY_METRICS),
         {"gptq": True},
-        {"awq": True},
-        {"scale_estimation": True},
         {"lora_correction": True},
-        {"dataset": Dataset([1])},
         {"compression_format": CompressionFormat.FQ},
         {"compression_format": CompressionFormat.FQ_LORA},
     ),
@@ -307,3 +311,134 @@ def test_model_devices_and_precisions(use_cuda, dtype):
     assert compressed_model.state_dict()["asymmetric_weights_decompressor_w._scale"].dtype == torch.float16
     # Result should be in the precision of the model
     assert result.dtype == dtype
+
+
+class TestFXTemplateWeightCompression(TemplateWeightCompression):
+    @staticmethod
+    def get_matmul_model() -> torch.fx.GraphModule:
+        model = MatMulModel(255 * torch.eye(3, dtype=torch.float32))
+        ex_input = torch.ones([1, 3, 3], dtype=torch.float32)
+        exported_model = get_torch_fx_model(model, ex_input)
+        return exported_model
+
+    @staticmethod
+    def get_sequential_matmul_model() -> torch.fx.GraphModule:
+        model = SequentialMatmulModel()
+        ex_input = torch.ones([1, 4, 4], dtype=torch.float32)
+        exported_model = get_torch_fx_model(model, ex_input)
+        return exported_model
+
+    @staticmethod
+    def get_model_for_test_scale_estimation():
+        model = LinearModel(torch.arange(0, 8 * 16, dtype=torch.float32).reshape(16, 8))
+        ex_input = torch.ones([1, 4, 8], dtype=torch.float32)
+        exported_model = get_torch_fx_model(model, ex_input)
+        return exported_model
+
+    @staticmethod
+    def get_awq_model() -> torch.fx.GraphModule:
+        model = AWQLinearModel()
+        dynamic_shapes = [[None, torch.export.Dim("dynamic_shape"), None]]
+        ex_input = torch.ones([1, 4, 8], dtype=torch.float32)
+        exported_model = get_torch_fx_model(model, ex_input, dynamic_shapes=dynamic_shapes)
+        return exported_model
+
+    @staticmethod
+    def get_awq_act_model(with_multiply, n_layers):
+        model = AWQActLinearModel(with_multiply=with_multiply, n_layers=n_layers)
+        ex_input = torch.ones([1, 8, 8], dtype=torch.float32)
+        exported_model = get_torch_fx_model(model, ex_input)
+        return exported_model
+
+    @staticmethod
+    def to_tensor(t) -> torch.Tensor:
+        return torch.tensor(t)
+
+    @staticmethod
+    def cast_to(x: torch.Tensor, dtype: TensorDataType) -> torch.Tensor:
+        return cast_to(x, dtype)
+
+    @staticmethod
+    def check_weights(model: torch.fx.GraphModule, ref_ids: list[int]) -> None:
+        all_names = list(model.graph.nodes)
+        low_precision_nodes = list(map(lambda i: all_names[i].name, ref_ids))
+        for node in model.graph.nodes:
+            for name in low_precision_nodes:
+                if name in node.name and node.op == "call_module":
+                    assert isinstance(getattr(model, node.target), INT4SymmetricWeightsDecompressor)
+
+    @staticmethod
+    def get_not_supported_algorithms() -> list[str]:
+        return ["lora_correction", "gptq"]
+
+    @staticmethod
+    def get_scale_estimation_ref():
+        return torch.tensor(
+            [
+                [[0.473328]],
+                [[0.929023]],
+                [[1.446527]],
+                [[1.920595]],
+                [[2.517054]],
+                [[3.030102]],
+                [[3.584279]],
+                [[4.043509]],
+                [[4.620008]],
+                [[5.165322]],
+                [[5.710637]],
+                [[6.122581]],
+                [[6.655914]],
+                [[7.237174]],
+                [[7.722580]],
+                [[8.255914]],
+            ]
+        )
+
+    @staticmethod
+    def get_orig_weight(model: torch.fx.GraphModule) -> Tensor:
+        return Tensor(model.linear.weight.data.detach())
+
+    @staticmethod
+    def get_decompressed_weight(compressed_model: torch.fx.GraphModule, input: torch.Tensor) -> Tensor:
+        for node in compressed_model.graph.nodes:
+            print(node.name)
+        model_graph = compressed_model.graph
+        weight_node = get_graph_node_by_name(model_graph, "linear_weight_updated_constant0")
+        decompression_node = get_graph_node_by_name(model_graph, "asymmetric_weights_decompressor_linear_weight_0")
+        weight = get_tensor_constant_from_node(weight_node, compressed_model)
+        decompress_module = getattr(compressed_model, decompression_node.target)
+        unpacked_w = decompress_module(weight)
+        return Tensor(unpacked_w)
+
+    @staticmethod
+    def get_ignored_scope_name() -> str:
+        return "linear_5"
+
+    @staticmethod
+    def get_num_int4_nodes(model: torch.fx.GraphModule) -> int:
+        num = 0
+        for node in model.graph.nodes:
+            if node.op != "call_module":
+                continue
+            num += isinstance(getattr(model, node.target), INT4SymmetricWeightsDecompressor)
+        return num
+
+    @pytest.fixture(params=INT4_MODES)
+    def int4_mode(self, request):
+        return request.param
+
+    @staticmethod
+    def get_num_multiply_from_awq(model):
+        awq_num = 0
+        for node in model.graph.nodes:
+            if "awq" in node.name and isinstance(getattr(model, node.target), FXAWQMultiply):
+                awq_num += 1
+        return awq_num
+
+    @staticmethod
+    def get_reference_for_test_awq_scale_reference() -> dict[str, Tensor]:
+        return {
+            "linear_2": Tensor(
+                torch.tensor([[1.226455, 1.205499, 1.141340, 1.097436, 1.064355, 1.037971, 1.016118, 0.997526]])
+            )
+        }
diff --git a/tests/torch/fx/test_weights_compression_backends.py b/tests/torch/fx/test_weights_compression_backends.py
index fccdd0e2c01..29fcec214eb 100644
--- a/tests/torch/fx/test_weights_compression_backends.py
+++ b/tests/torch/fx/test_weights_compression_backends.py
@@ -12,27 +12,27 @@
 from nncf.quantization.algorithms.weight_compression.mixed_precision import MaxVarianceCriterion
 from nncf.quantization.algorithms.weight_compression.mixed_precision import MeanMaxCriterion
 from nncf.quantization.algorithms.weight_compression.mixed_precision import MeanVarianceCriterion
-from nncf.quantization.algorithms.weight_compression.torch_backend import PTMixedPrecisionAlgoBackend
+from nncf.quantization.algorithms.weight_compression.torch_fx_backend import FXMixedPrecisionAlgoBackend
 from tests.cross_fw.test_templates.test_weights_compression_backends import TemplateTestMixedPrecisionAlgoBackend
 
 
-class TestPTMixedPrecisionAlgoBackend(TemplateTestMixedPrecisionAlgoBackend):
+class TestFXMixedPrecisionAlgoBackend(TemplateTestMixedPrecisionAlgoBackend):
     def get_hawq_with_backend(self, subset_size):
         hawq = HAWQCriterion(None, None, subset_size=subset_size)
-        hawq._backend_entity = PTMixedPrecisionAlgoBackend()
+        hawq._backend_entity = FXMixedPrecisionAlgoBackend()
         return hawq
 
     def get_mean_variance_with_backend(self, subset_size: int):
         mean_variance = MeanVarianceCriterion(None, None, subset_size=subset_size)
-        mean_variance._backend_entity = PTMixedPrecisionAlgoBackend()
+        mean_variance._backend_entity = FXMixedPrecisionAlgoBackend()
         return mean_variance
 
     def get_max_variance_with_backend(self, subset_size: int):
         max_variance = MaxVarianceCriterion(None, None, subset_size=subset_size)
-        max_variance._backend_entity = PTMixedPrecisionAlgoBackend()
+        max_variance._backend_entity = FXMixedPrecisionAlgoBackend()
         return max_variance
 
     def get_mean_max_with_backend(self, subset_size: int):
         mean_max_variance = MeanMaxCriterion(None, None, subset_size=subset_size)
-        mean_max_variance._backend_entity = PTMixedPrecisionAlgoBackend()
+        mean_max_variance._backend_entity = FXMixedPrecisionAlgoBackend()
         return mean_max_variance


Reviews:
user: login: daniil-lyakhov
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T12:39:18Z
user: login: daniil-lyakhov
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T12:39:23Z
user: login: anzr299
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T12:42:51Z
user: login: daniil-lyakhov
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T12:46:16Z
user: login: anzr299
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T12:48:08Z
user: login: anzr299
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T12:49:55Z
user: login: daniil-lyakhov
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-23T12:52:50Z
user: login: anzr299
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-23T13:00:30Z
user: login: daniil-lyakhov
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-24T16:02:49Z

Review comments:
diff_hunk: @@ -265,8 +274,73 @@ def transform_model(
                     )
                 )
             )
-
+        if first_caught_error:
+            handle_invalid_group_size_error(first_caught_error, invalid_node_names)
         # apply transformations
         transformed_model = FXModelTransformer(model).transform(transformation_layout)
 
         return transformed_model
+
+
+class FXMixedPrecisionAlgoBackend(MixedPrecisionAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
+    def mean_variance_statistic_collector(
+        reduction_axes: Tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.mean_variance_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+    @staticmethod
+    def max_variance_statistic_collector(
+        reduction_axes: Tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.max_variance_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+    @staticmethod
+    def mean_abs_max_statistic_collector(
+        reduction_axes: Tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.mean_abs_max_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+
+class FXAWQMultiply(torch.nn.Module):
+    def __init__(self, scale: torch.Tensor):
+        super().__init__()
+        self.register_buffer("_scale_value", scale)
+        self._scale_value: torch.Tensor
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.mul(x, self._scale_value)
+
+
+class FXAWQAlgoAlgoBackend(AWQAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
+    def get_awq_patterns():
+        return PTAWQAlgoAlgoBackend.get_awq_patterns()
+
+    @staticmethod
+    def scale_insertion_command(source_node, next_nodes, source_node_output_port, scale):
+        input_port_id = 0
+        target_points = []
+        for node in next_nodes:
+            target_points.append(
+                PTTargetPoint(
+                    PTWeightCompressionAlgoBackend.TARGET_TYPE_TO_PT_INS_TYPE_MAP[TargetType.PRE_LAYER_OPERATION],
path: nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
user: login: daniil-lyakhov
body: Could you please specify the actual target type without the dict?
created_at: 2025-04-11T12:39:18Z
author_association: COLLABORATOR
diff_hunk: @@ -265,8 +274,73 @@ def transform_model(
                     )
                 )
             )
-
+        if first_caught_error:
+            handle_invalid_group_size_error(first_caught_error, invalid_node_names)
         # apply transformations
         transformed_model = FXModelTransformer(model).transform(transformation_layout)
 
         return transformed_model
+
+
+class FXMixedPrecisionAlgoBackend(MixedPrecisionAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
path: nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
user: login: daniil-lyakhov
body: Do we have a plan on how to reuse the same classes across all the PyTorch backends? If not, could you please create a ticket for that?
created_at: 2025-04-11T12:39:22Z
author_association: COLLABORATOR
diff_hunk: @@ -265,8 +274,73 @@ def transform_model(
                     )
                 )
             )
-
+        if first_caught_error:
+            handle_invalid_group_size_error(first_caught_error, invalid_node_names)
         # apply transformations
         transformed_model = FXModelTransformer(model).transform(transformation_layout)
 
         return transformed_model
+
+
+class FXMixedPrecisionAlgoBackend(MixedPrecisionAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
path: nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
user: login: anzr299
body: Maybe we can set a flag in a common backend class (such as FX=False) when initializing it.
created_at: 2025-04-11T12:42:51Z
author_association: COLLABORATOR
diff_hunk: @@ -265,8 +274,73 @@ def transform_model(
                     )
                 )
             )
-
+        if first_caught_error:
+            handle_invalid_group_size_error(first_caught_error, invalid_node_names)
         # apply transformations
         transformed_model = FXModelTransformer(model).transform(transformation_layout)
 
         return transformed_model
+
+
+class FXMixedPrecisionAlgoBackend(MixedPrecisionAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
path: nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
user: login: daniil-lyakhov
body: @anzr299, It would be great if you would open a ticket and implement a make a proposal which we could discuss with the team 
created_at: 2025-04-11T12:46:16Z
author_association: COLLABORATOR
diff_hunk: @@ -265,8 +274,73 @@ def transform_model(
                     )
                 )
             )
-
+        if first_caught_error:
+            handle_invalid_group_size_error(first_caught_error, invalid_node_names)
         # apply transformations
         transformed_model = FXModelTransformer(model).transform(transformation_layout)
 
         return transformed_model
+
+
+class FXMixedPrecisionAlgoBackend(MixedPrecisionAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
path: nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
user: login: anzr299
body: Alright, I will do that
created_at: 2025-04-11T12:48:08Z
author_association: COLLABORATOR
diff_hunk: @@ -265,8 +274,73 @@ def transform_model(
                     )
                 )
             )
-
+        if first_caught_error:
+            handle_invalid_group_size_error(first_caught_error, invalid_node_names)
         # apply transformations
         transformed_model = FXModelTransformer(model).transform(transformation_layout)
 
         return transformed_model
+
+
+class FXMixedPrecisionAlgoBackend(MixedPrecisionAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
+    def mean_variance_statistic_collector(
+        reduction_axes: Tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.mean_variance_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+    @staticmethod
+    def max_variance_statistic_collector(
+        reduction_axes: Tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.max_variance_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+    @staticmethod
+    def mean_abs_max_statistic_collector(
+        reduction_axes: Tuple[int], subset_size: Optional[int] = None
+    ) -> TensorCollector:
+        return PTMixedPrecisionAlgoBackend.mean_abs_max_statistic_collector(
+            reduction_axes=reduction_axes, subset_size=subset_size
+        )
+
+
+class FXAWQMultiply(torch.nn.Module):
+    def __init__(self, scale: torch.Tensor):
+        super().__init__()
+        self.register_buffer("_scale_value", scale)
+        self._scale_value: torch.Tensor
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.mul(x, self._scale_value)
+
+
+class FXAWQAlgoAlgoBackend(AWQAlgoBackend, FXWeightCompressionAlgoBackend):
+    @staticmethod
+    def get_awq_patterns():
+        return PTAWQAlgoAlgoBackend.get_awq_patterns()
+
+    @staticmethod
+    def scale_insertion_command(source_node, next_nodes, source_node_output_port, scale):
+        input_port_id = 0
+        target_points = []
+        for node in next_nodes:
+            target_points.append(
+                PTTargetPoint(
+                    PTWeightCompressionAlgoBackend.TARGET_TYPE_TO_PT_INS_TYPE_MAP[TargetType.PRE_LAYER_OPERATION],
path: nncf/quantization/algorithms/weight_compression/torch_fx_backend.py
user: login: anzr299
body: Done, my reasoning behind this approach was that it will be less modular if we directly define it in case of changes to the target type in the future
created_at: 2025-04-11T12:49:54Z
author_association: COLLABORATOR
diff_hunk: @@ -256,6 +266,30 @@ def collect_data_from_stdout(self, stdout: str):
         stats.fill(stdout)
         self.run_info.stats_from_output = stats
 
+    def _rename_files(self, folder_path, new_name):
path: tests/post_training/pipelines/lm_weight_compression.py
user: login: daniil-lyakhov
body: Please refactor the same way it done here: https://github.com/openvinotoolkit/nncf/pull/3321/files#diff-5d01260f55336ab44f9ae6583b583333bdbd46271c48cf70862bb17c33d5c191R513-R521
created_at: 2025-04-23T12:49:51Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,171 @@
+# 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 Dict
+
+import torch
+import torch.fx
+from optimum.modeling_base import OptimizedModel
+from transformers import GenerationConfig
+from transformers import GenerationMixin
+from transformers import PretrainedConfig
+from transformers import PreTrainedModel
+from transformers.cache_utils import StaticCacheConfig
+from transformers.integrations.executorch import TorchExportableModuleWithStaticCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class FXAutoModelForCausalLM(OptimizedModel, GenerationMixin):
path: tests/post_training/pipelines/fx_utils.py
user: login: daniil-lyakhov
body: Can we import this classes instead of copy-paste?
created_at: 2025-04-23T12:50:53Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,171 @@
+# 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 Dict
+
+import torch
+import torch.fx
+from optimum.modeling_base import OptimizedModel
+from transformers import GenerationConfig
+from transformers import GenerationMixin
+from transformers import PretrainedConfig
+from transformers import PreTrainedModel
+from transformers.cache_utils import StaticCacheConfig
+from transformers.integrations.executorch import TorchExportableModuleWithStaticCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class FXAutoModelForCausalLM(OptimizedModel, GenerationMixin):
path: tests/post_training/pipelines/fx_utils.py
user: login: anzr299
body: For converting and exporting with cache, I think yes. I am checking this right now. 
As for the `FXAutoModelForCausalLM` for generation. The class we can import does not allow to `torch.compile` the model as 2 models (prefill and decode one token).
https://github.com/huggingface/transformers/blob/cc9a245e6df408711e840ad778c73e05ead18930/src/transformers/integrations/executorch.py#L181
created_at: 2025-04-23T13:00:29Z
author_association: COLLABORATOR

PR comments:
user: login: MaximProshin
created_at: 2025-04-08T11:49:41Z
body: @anzr299 , please share conformance test results for these methods
author_association: COLLABORATOR
user: login: daniil-lyakhov
created_at: 2025-04-23T09:22:23Z
body: @anzr299, please attach link to the issue, fix the pre-commit, and rebase
author_association: COLLABORATOR
user: login: anzr299
created_at: 2025-04-23T09:45:08Z
body: > @anzr299, please attach link to the issue, fix the pre-commit, and rebase

Done
Issue created at CVS-166664
author_association: COLLABORATOR