Comments

Author: daniil-lyakhov

tests/openvino/native/models.py

@@ -1084,6 +1084,42 @@ def _create_ov_model(self, is_int8=False, with_multiply=False, n_layers=8):
         return model
 
 
+class AWQModel(OVReferenceModel):
+    OUTPUT_DIM = 32
+    HIDDEN_DIM = 16
+    INPUT_SHAPE = [1, 24, HIDDEN_DIM]  # [B, SeqLen, HiddenDim]
+
+    def _create_ov_model(
+        self,
+        transpose_b: bool = True,
+        transpose_a: bool = False,
+        input_shape: Optional[list[int]] = None,
+        is_int8=False,
+    ):
+        self._input_shape = self.INPUT_SHAPE if input_shape is None else input_shape
+        hdim_axis = -2 if transpose_a else -1
+        self._hidden_dim = self._input_shape[hdim_axis]
+        input_1 = opset.parameter(self._input_shape, name="Input")
+        weight_shape = self.get_weight_shape(transpose_b)
+        data = self._rng.random(weight_shape).astype(np.float32)
+
+        weights = AWQMatmulModel.get_weights(data, is_int8=is_int8, name="weights_1")
+
+        matmul = opset.matmul(input_1, weights, transpose_a=transpose_a, transpose_b=transpose_b, name="MatMul")
+
+        result = opset.result(matmul, name="Result")
+        result.get_output_tensor(0).set_names(set(["Result"]))
+        model = ov.Model([result], [input_1])
+        return model
+
+    @property
+    def hidden_dim(self):
+        return self._hidden_dim
+
+    def get_weight_shape(self, transpose_b: bool = True):
+        return [self.OUTPUT_DIM, self.hidden_dim] if transpose_b else [self.hidden_dim, self.OUTPUT_DIM]
+
+
 class AWQModel_fp16_overlow(OVReferenceModel):
     """
     Model for testing AWQ algorithm with fp16 overflow fix.

tests/openvino/native/quantization/test_weights_compression.py

@@ -65,6 +65,7 @@
 from tests.openvino.native.common import get_actual_reference_for_current_openvino
 from tests.openvino.native.models import AWQActMatmulModel
 from tests.openvino.native.models import AWQMatmulModel
+from tests.openvino.native.models import AWQModel
 from tests.openvino.native.models import AWQModel_fp16_overlow
 from tests.openvino.native.models import DifferentChannelSizeMatmulModel
 from tests.openvino.native.models import GatherAndMatmulShareData
@@ -1943,6 +1944,14 @@ def test_compression_with_different_algo_combinations(input_shape, kwargs):
     )
 
 
+@pytest.mark.parametrize(
+    "model_cls",
+    [
+        (LMLinearModel),
+        (AWQModel),
+    ],
+    ids=["lm_linear", "awq_model"],
+)
 @pytest.mark.parametrize(
     ("transpose_a", "transpose_b", "raises_error"),
     [
@@ -1967,9 +1976,9 @@ def test_compression_with_different_algo_combinations(input_shape, kwargs):
     ],
     ids=["se", "lora", "gptq_se_awq"],
 )
-def test_compression_with_transpose(transpose_a, transpose_b, raises_error, kwargs):
+def test_compression_with_transpose(model_cls, transpose_a, transpose_b, raises_error, kwargs):
     dataset_size = 4
-    model = LMLinearModel(transpose_a=transpose_a, transpose_b=transpose_b).ov_model
+    model = model_cls(transpose_a=transpose_a, transpose_b=transpose_b).ov_model
     input_data = [np.ones(inp.shape) for inp in model.inputs] * dataset_size
     dataset = Dataset(input_data)