Data-aware compression fixes (#3019)

Fixes after #3003 .

### Changes

1. Convert raw activations to WC statistics for GPTQ + SE scenario.
2. Allow 2D tensor inputs for data-aware mixed precision. 2D activations
arise in `opt`-like models, e.g. `opt-125m`. There, LayerNorm reshapes
activations from [B, L, D] to [B*L, D].

### Tests

1. Added a unit test for GPTQ + SE.
2. Modified a test for 2D activations and mixed precision.
3. Compressed tiny-llama to int4_asym with SQ + GPTQ before #3003 and
for this PR. Got the same PPL value of 15.739704794594019 .

Author: nikita-savelyevv

nncf/common/tensor_statistics/statistics.py

@@ -126,7 +126,7 @@ class WCTensorStatistic(TensorStatistic):
     MEAN_STAT = "mean_values"
     SHAPE_STAT = "shape_values"
 
-    def __init__(self, mean_values: List[Tensor], shapes: List[Tuple[int]]):
+    def __init__(self, mean_values: List[Tensor], shapes: List[Tuple[int, ...]]):
         """
         :param mean_values: List of N tensors of shape [HiddenDim] obtained by reducing activations along batch and
             sequence length dimensions.

nncf/quantization/algorithms/weight_compression/gptq.py

@@ -18,6 +18,7 @@
 from nncf.common.graph import NNCFNode
 from nncf.common.logging.track_progress import track
 from nncf.common.tensor_statistics.statistic_point import StatisticPointsContainer
+from nncf.common.tensor_statistics.statistics import WCTensorStatistic
 from nncf.common.utils.backend import BackendType
 from nncf.common.utils.backend import get_backend
 from nncf.parameters import CompressWeightsMode
@@ -265,9 +266,10 @@ def _quantize_weights(
                     else:
                         if self._scale_estimation and block_compression_config.num_bits == 4:
                             activations = [inp.squeeze()[:, (i1 + i) : (i1 + i + group_size)] for inp in inputs]
+                            wc_statistics = self._activations_to_wc_statistics(activations)
                             scale, zero_point = ScaleEstimation.calculate_quantization_params(
                                 self._backend_entity,
-                                activations,
+                                wc_statistics,
                                 weight_tensor[:, (i1 + i) : (i1 + i + group_size)],
                                 reduction_axes,
                                 wc_params.compression_config,
@@ -325,3 +327,15 @@ def _quantize_weights(
         else:
             zero_points = None
         return scales, zero_points
+
+    @staticmethod
+    def _activations_to_wc_statistics(activations: List[Tensor]) -> WCTensorStatistic:
+        # The code below mimics the logic from WeightCompression.get_statistic_points
+        mean_values = []
+        shapes = []
+        for act in activations:
+            shapes.append(act.shape)
+            reduction_shape = tuple(range(act.ndim - 1))
+            mean_values.append(fns.mean(act, axis=reduction_shape))
+        wc_statistics = WCTensorStatistic(mean_values, shapes)
+        return wc_statistics

nncf/quantization/algorithms/weight_compression/mixed_precision.py

@@ -237,9 +237,9 @@ def get_statistic_points(
         statistic_container = StatisticPointsContainer()
         for act_node, output_port_id in nodes_and_port_ids:
             n_dims = len(graph.get_output_edges_by_port_id(act_node, output_port_id)[0].tensor_shape)
-            if n_dims < 3:
+            if n_dims < 2:
                 raise RuntimeError(
-                    f"Data-aware mixed precision criteria are not supported for MatMuls with 1D/2D activations. "
+                    f"Data-aware mixed precision criteria are not supported for MatMuls with 1D inputs. "
                     f"Node: {act_node.node_name}, number of dimensions: {n_dims}."
                 )
             statistic_point = self._backend_entity.target_point(

tests/openvino/native/quantization/test_weights_compression.py

@@ -30,6 +30,7 @@
 from nncf.parameters import BackupMode
 from nncf.quantization import compress_weights
 from nncf.quantization.advanced_parameters import AdvancedCompressionParameters as CompressionParams
+from nncf.quantization.advanced_parameters import AdvancedGPTQParameters as GPTQParams
 from nncf.quantization.advanced_parameters import AdvancedLoraCorrectionParameters as LoraParams
 from nncf.quantization.algorithms.weight_compression.config import WeightCompressionConfig
 from nncf.quantization.algorithms.weight_compression.config import WeightCompressionParameters
@@ -1377,29 +1378,42 @@ def test_data_aware_algo_with_different_activation_dimensions(n_extra_dims):
         group_size=-1,
         dataset=dataset,
         awq=True,
+        ratio=0.5,
+        sensitivity_metric=SensitivityMetric.MEAN_ACTIVATION_MAGNITUDE,
     )
 
 
-@pytest.mark.parametrize("n_extra_dims,raises", ([0, True], (1, False), (2, False)))
-def test_data_aware_mixed_precision_with_different_activation_dimensions(n_extra_dims, raises):
-    model = AWQMatmulModel(n_extra_dims=n_extra_dims).ov_model
-    dataset = Dataset([np.ones([1] * n_extra_dims + [8, 8])])
-
-    def call_compression():
-        compress_weights(
-            model,
-            mode=CompressWeightsMode.INT4_ASYM,
-            ratio=0.5,
-            sensitivity_metric=SensitivityMetric.MEAN_ACTIVATION_MAGNITUDE,
-            group_size=-1,
-            dataset=dataset,
-        )
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        dict(scale_estimation=True),
+        dict(lora_correction=True),
+        dict(
+            gptq=True,
+            scale_estimation=True,
+            advanced_parameters=CompressionParams(gptq_params=GPTQParams(subset_size=2)),
+        ),
+        dict(
+            awq=True,
+            gptq=True,
+            scale_estimation=True,
+            advanced_parameters=CompressionParams(gptq_params=GPTQParams(subset_size=2)),
+        ),
+    ],
+)
+def test_compression_with_different_algo_combinations(kwargs):
+    dataset_size = 4
+    model = LMLinearModel().ov_model
+    input_data = [np.ones(inp.shape) for inp in model.inputs] * dataset_size
+    dataset = Dataset(input_data)
 
-    if raises:
-        with pytest.raises(RuntimeError) as exc_info:
-            call_compression()
-        assert "Data-aware mixed precision criteria are not supported for MatMuls with 1D/2D activations." in str(
-            exc_info.value
-        )
-    else:
-        call_compression()
+    compress_weights(
+        model,
+        mode=CompressWeightsMode.INT4_SYM,
+        ratio=1.0,
+        group_size=8,
+        subset_size=2,
+        dataset=dataset,
+        all_layers=True,
+        **kwargs,
+    )