[Torch] Add the GroupedMatMul operation for data-free weight compression

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

@@ -1033,6 +1033,18 @@ def get_weight_compression_parameters(
                         )
 
                     model_backend = get_backend(model)
+                    if model_backend == BackendType.TORCH and (
+                        self._data_aware_compression or self._data_aware_mixed_precision
+                    ):
+                        from nncf.torch.graph.operator_metatypes import PTGroupedMatMulMetatype
+
+                        if node.metatype == PTGroupedMatMulMetatype:
+                            msg = (
+                                "Data-aware weight compression is not supported for GroupedMatMul in the Torch "
+                                f"backend. Node with weight: {node.node_name}."
+                            )
+                            raise nncf.ParameterNotSupportedError(msg)
+
                     ov_version = self._maybe_get_ov_version()
                     if (
                         model_backend == BackendType.OPENVINO

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

@@ -51,8 +51,10 @@
 from nncf.torch.graph.operator_metatypes import EMBEDDING_METATYPES
 from nncf.torch.graph.operator_metatypes import MATMUL_METATYPES
 from nncf.torch.graph.operator_metatypes import PTDropoutMetatype
+from nncf.torch.graph.operator_metatypes import PTGroupedMatMulMetatype
 from nncf.torch.graph.operator_metatypes import PTMulMetatype
 from nncf.torch.graph.operator_metatypes import PTNoopMetatype
+from nncf.torch.graph.operator_metatypes import PTTransposeMetatype
 from nncf.torch.graph.pattern_operations import ATOMIC_ACTIVATIONS_OPERATIONS
 from nncf.torch.graph.transformations.commands import PTSharedFnInsertionCommand
 from nncf.torch.graph.transformations.commands import PTTransformationCommand
@@ -134,6 +136,10 @@ def get_reduction_axes(node_with_weight: NNCFNode, weight_port_id: int, graph: N
 
         ndims = len(weight_node.layer_attributes.shape)
         reduction_axes = get_weight_compression_reduction_axes(node_with_weight_metatype, weight_port_id, ndims)
+        weight_input_metatype = graph.get_input_edge_by_port_id(node_with_weight, weight_port_id).from_node.metatype
+        if node_with_weight_metatype == PTGroupedMatMulMetatype and weight_input_metatype != PTTransposeMetatype:
+            # A weight that is not transposed is consumed with the [num_groups, C_IN, C_OUT] layout.
+            reduction_axes = [ndims - 2]
         return tuple(reduction_axes)
 
     @staticmethod
@@ -220,14 +226,14 @@ def init_lora_adapters(svd_residual: torch.Tensor, rank: int | None = None) -> t
         :param rank: The rank for the decomposition. If None, the full rank is used.
         :return: A tuple containing the U and V matrices from the SVD.
         """
-        # O stands for output dimension, H - input dimension or hidden size, R - rank.
+        # O stands for output dimension, H - input dimension or hidden size, R - rank, B - batch dimension.
         U_full, S_full, V_full = torch.linalg.svd(svd_residual, full_matrices=False)
-        U = U_full[:, :rank]  # [H, R]
+        U = U_full[..., :rank]  # [H, R] or [B, H, R]
         S_sqrt = torch.sqrt(S_full)
-        S = torch.diag(S_sqrt[:rank])  # [R, R]
-        V = V_full[:rank, :]  # [R, O]
-        V = S @ V  # [R, O]
-        U = U @ S  # [H, R]
+        S = torch.diag_embed(S_sqrt[..., :rank])  # [R, R] or [B, R, R]
+        V = V_full[..., :rank, :]  # [R, O] or [B, R, O]
+        V = S @ V  # [R, O] or [B, R, O]
+        U = U @ S  # [H, R] or [B, H, R]
         return U, V
 
     @staticmethod

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

@@ -50,6 +50,7 @@
 from nncf.torch.graph.operator_metatypes import CONVOLUTION_METATYPES
 from nncf.torch.graph.operator_metatypes import EMBEDDING_METATYPES
 from nncf.torch.graph.operator_metatypes import MATMUL_METATYPES
+from nncf.torch.graph.operator_metatypes import PTGroupedMatMulMetatype
 from nncf.torch.graph.transformations.commands import PTTargetPoint
 from nncf.torch.model_graph_manager import get_const_node
 from nncf.torch.model_graph_manager import get_weight_compression_reduction_axes
@@ -66,7 +67,8 @@
 class FXWeightCompressionAlgoBackend(WeightCompressionAlgoBackend):
     @property
     def matmul_metatypes(self) -> list[OperatorMetatype]:
-        return MATMUL_METATYPES
+        # TODO(anzr299): add groupedmatmul support in torch fx
+        return [metatype for metatype in MATMUL_METATYPES if metatype != PTGroupedMatMulMetatype]
 
     @property
     def embedding_metatypes(self) -> list[OperatorMetatype]:

src/nncf/torch/graph/operator_metatypes.py

@@ -522,6 +522,16 @@ class PTMatMulMetatype(PTOperatorMetatype):
     weight_port_ids = [0, 1]
 
 
+@PT_OPERATOR_METATYPES.register()
+class PTGroupedMatMulMetatype(PTOperatorMetatype):
+    name = "GroupedMatMulOp"
+    module_to_function_names = {
+        NamespaceTarget.TORCH: ["_grouped_mm"],
+        NamespaceTarget.TORCH_NN_FUNCTIONAL: ["grouped_mm"],
+    }
+    weight_port_ids = [1]
+
+
 @PT_OPERATOR_METATYPES.register()
 class PTAddmmMetatype(PTOperatorMetatype):
     name = "MatMulOp"
@@ -1041,7 +1051,7 @@ def get_operator_metatypes() -> list[type[OperatorMetatype]]:
     PTAtenEmbeddingBagMetatype,
 ]
 
-MATMUL_METATYPES = [PTLinearMetatype, PTMatMulMetatype, PTAddmmMetatype]
+MATMUL_METATYPES = [PTLinearMetatype, PTMatMulMetatype, PTAddmmMetatype, PTGroupedMatMulMetatype]
 
 EMBEDDING_METATYPES = [PTEmbeddingMetatype, PTAtenEmbeddingMetatype]
 

src/nncf/torch/model_graph_manager.py

@@ -36,7 +36,9 @@
 
 OPERATORS_WITH_BIAS_METATYPES = CONV_META_TYPES + [om.PTLinearMetatype]
 CONV_FUSED_META_TYPES = [om.PTBatchNormMetatype]
-AVAILABLE_NODE_TYPES_FOR_WEIGHT_SUBGRAPH = om.QUANTIZE_NODE_TYPES + om.PRUNING_NODE_TYPES
+AVAILABLE_NODE_TYPES_FOR_WEIGHT_SUBGRAPH = (
+    om.QUANTIZE_NODE_TYPES + om.PRUNING_NODE_TYPES + ["transpose", "permute", "transpose_"]
+)
 
 
 def find_const_node_in_constant_subgraph(node: NNCFNode, graph: NNCFGraph) -> NNCFNode | None:
@@ -359,7 +361,7 @@ def get_weight_compression_reduction_axes(metatype: OperatorMetatype, weight_por
     """
     if metatype in [om.PTAtenEmbeddingMetatype, om.PTEmbeddingMetatype]:
         return [1]
-    if metatype == om.PTLinearMetatype:
+    if metatype in [om.PTLinearMetatype, om.PTGroupedMatMulMetatype]:
         return [ndims - 1]
     if metatype == om.PTMatMulMetatype:
         if weight_port_id == 0:

src/nncf/torch/node_utils.py

@@ -28,7 +28,7 @@ def get_activation_channel_axis(node: NNCFNode, port_id: int) -> int:
         msg = f"Activation channel axis retrieval from node with metatype {node.metatype} is not supported"
         raise nncf.InternalError(msg)
 
-    if node.metatype not in [PTMatMulMetatype, PTAddmmMetatype]:
+    if node.metatype not in [PTMatMulMetatype, PTAddmmMetatype, op.PTGroupedMatMulMetatype]:
         if not isinstance(node.metatype.output_channel_axis, int):
             msg = f"Node metatype {node.metatype} does not have defined output channel axis"
             raise nncf.InternalError(msg)

src/nncf/torch/quantization/layers.py

@@ -987,16 +987,18 @@ class LoraMixin:
     def init_lora(self, lspec: PTLoraSpec):
         self._lspec = lspec
         default_lora_dtype = torch.bfloat16
-        out_features, in_features = lspec.orig_weight_shape
+        # 3D weights have leading batch dimension. A separate pair of adapters
+        # is created for each group.
+        *group_dims, out_features, in_features = lspec.orig_weight_shape
         rank = lspec.lora_rank
         if rank > out_features or rank > in_features:
             msg = (
                 f"Specified LoRA rank={rank} cannot exceed any dimension of the weight tensor: "
                 f"[{out_features}, {in_features}]"
             )
             raise nncf.ValidationError(msg)
-        self.lora_A = torch.nn.Parameter(torch.ones((rank, in_features), dtype=default_lora_dtype))
-        self.lora_B = torch.nn.Parameter(torch.zeros((out_features, rank), dtype=default_lora_dtype))
+        self.lora_A = torch.nn.Parameter(torch.ones((*group_dims, rank, in_features), dtype=default_lora_dtype))
+        self.lora_B = torch.nn.Parameter(torch.zeros((*group_dims, out_features, rank), dtype=default_lora_dtype))
 
     def enable_gradients(self):
         self.lora_A.requires_grad = True
@@ -1042,8 +1044,8 @@ def get_active_adapters(self) -> tuple[torch.Tensor, torch.Tensor]:
 
         :return: A dictionary containing the active LoRA adapters.
         """
-        lora_A = self.lora_A[: self.active_lora_rank, :]
-        lora_B = self.lora_B[:, : self.active_lora_rank]
+        lora_A = self.lora_A[..., : self.active_lora_rank, :]
+        lora_B = self.lora_B[..., : self.active_lora_rank]
         return lora_A, lora_B
 
 

tests/torch/function_hook/quantization/test_fq_lora.py

@@ -33,6 +33,7 @@
 from nncf.torch.quantization.layers import SymmetricQuantizer as SQ
 from tests.cross_fw.shared.paths import TEST_ROOT
 from tests.torch.function_hook.quantization.test_weights_compression import AWQLinearModel
+from tests.torch.test_models.synthetic import GroupedMatMulModel
 from tests.torch.test_models.synthetic import LinearModel
 from tests.torch.test_models.synthetic import ShortTransformer
 from tests.torch.utils import compare_with_reference_file
@@ -199,6 +200,29 @@ def test_invalid_lora_rank():
         )
 
 
+@pytest.mark.parametrize("transposed_weight", (False, True), ids=("native_weight", "transposed_weight"))
+def test_fq_lora_grouped_matmul(transposed_weight):
+    rank = 2
+    model = GroupedMatMulModel(transposed_weight=transposed_weight)
+    example_input = torch.ones(model.INPUT_SHAPE, dtype=model.INPUT_DTYPE)
+    wrapped_model = wrap_model(model, example_input=example_input, trace_parameters=True)
+
+    compressed_model = compress_weights(
+        wrapped_model,
+        mode=CompressWeightsMode.INT4_SYM,
+        group_size=4,
+        all_layers=True,
+        compression_format=CompressionFormat.FQ_LORA,
+        advanced_parameters=AdvancedCompressionParameters(lora_adapter_rank=rank),
+    )
+
+    num_groups, c_in, c_out = GroupedMatMulModel.NUM_GROUPS, GroupedMatMulModel.C_IN, GroupedMatMulModel.C_OUT
+    out_features, in_features = (c_out, c_in) if transposed_weight else (c_in, c_out)
+    quantizer = next(module for module in compressed_model.modules() if isinstance(module, LoraMixin))
+    assert quantizer.lora_A.shape == (num_groups, rank, in_features)
+    assert quantizer.lora_B.shape == (num_groups, out_features, rank)
+
+
 @pytest.mark.parametrize("all_layers", [True, False])
 def test_compress_shared_weights(all_layers, regen_ref_data):
     model = ShortTransformer(8, 16, share_weights=True)

tests/torch/function_hook/quantization/test_weights_compression.py

@@ -48,6 +48,7 @@
 from tests.cross_fw.test_templates.helpers import RoPEModel
 from tests.cross_fw.test_templates.helpers import SAMPEModel
 from tests.cross_fw.test_templates.template_test_weights_compression import TemplateWeightCompression
+from tests.torch.test_models.synthetic import GroupedMatMulModel
 from tests.torch.test_models.synthetic import ShortTransformer
 from tests.torch.test_tensor import cast_to
 
@@ -412,6 +413,39 @@ def test_compress_weights_conv():
     assert n_compressed_weights == n_target_modules
 
 
+@pytest.mark.parametrize("transposed_weight", (False, True), ids=("native_weight", "transposed_weight"))
+@pytest.mark.parametrize("mode", SUPPORTED_MODES)
+def test_compress_weights_grouped_matmul(mode, transposed_weight):
+    model = GroupedMatMulModel(transposed_weight=transposed_weight)
+    example_input = torch.ones(model.INPUT_SHAPE, dtype=model.INPUT_DTYPE)
+    wrapped_model = GraphModelWrapper(wrap_model(model), example_input=example_input)
+
+    kwargs = {"group_size": 4, "all_layers": True} if mode in INT4_MODES else {}
+    compressed_model = compress_weights(wrapped_model, mode=mode, **kwargs)
+
+    n_compressed_weights = 0
+    for module in compressed_model.modules():
+        if isinstance(module, BaseWeightsDecompressor):
+            n_compressed_weights += 1
+    assert n_compressed_weights == 1
+
+
+@pytest.mark.parametrize("algorithm", ("awq", "scale_estimation"))
+def test_raise_error_for_data_aware_grouped_matmul(algorithm):
+    model = GroupedMatMulModel()
+    example_input = torch.ones(model.INPUT_SHAPE, dtype=model.INPUT_DTYPE)
+    wrapped_model = GraphModelWrapper(wrap_model(model), example_input=example_input)
+    with pytest.raises(nncf.ParameterNotSupportedError, match="GroupedMatMul"):
+        compress_weights(
+            wrapped_model,
+            mode=CompressWeightsMode.INT4_SYM,
+            group_size=4,
+            all_layers=True,
+            dataset=nncf.Dataset([example_input]),
+            **{algorithm: True},
+        )
+
+
 @pytest.mark.parametrize("mode", SUPPORTED_MODES)
 def test_compress_shared_weights(mocker, mode):
     model = wrap_model(ShortTransformer(8, 16, share_weights=True))

tests/torch/test_model_graph_manager.py

@@ -350,6 +350,7 @@ def test_update_fused_bias(model_cls):
         (om.PTAddmmMetatype, 1, 3, [2]),
         (om.PTAddmmMetatype, 2, 3, [1]),
         (om.PTAddmmMetatype, 2, 1, [0]),
+        (om.PTGroupedMatMulMetatype, 1, 3, [2]),
         (om.PTConv1dMetatype, 0, 3, [1, 2]),
         (om.PTConv2dMetatype, 0, 4, [1, 2, 3]),
         (om.PTConv3dMetatype, 0, 5, [1, 2, 3, 4]),

tests/torch/test_models/synthetic.py

@@ -591,6 +591,29 @@ def forward(self, input):
         return self.linear(input)
 
 
+class GroupedMatMulModel(torch.nn.Module):
+    NUM_GROUPS = 2
+    C_IN = 8
+    C_OUT = 16
+    NUM_TOKENS = 6
+    INPUT_SHAPE = (NUM_TOKENS, C_IN)
+    INPUT_DTYPE = torch.bfloat16
+
+    def __init__(self, transposed_weight: bool = False):
+        super().__init__()
+        self.transposed_weight = transposed_weight
+        shape = (
+            (self.NUM_GROUPS, self.C_OUT, self.C_IN) if transposed_weight else (self.NUM_GROUPS, self.C_IN, self.C_OUT)
+        )
+        weight = torch.arange(1, self.NUM_GROUPS * self.C_IN * self.C_OUT + 1, dtype=self.INPUT_DTYPE)
+        self.weight = torch.nn.Parameter(weight.reshape(shape))
+        self.register_buffer("offs", torch.tensor([self.NUM_TOKENS // 2, self.NUM_TOKENS], dtype=torch.int32))
+
+    def forward(self, x):
+        weight = self.weight.transpose(-2, -1) if self.transposed_weight else self.weight
+        return torch._grouped_mm(x, weight, offs=self.offs)
+
+
 class YOLO11N_SDPABlock(torch.nn.Module):
     INPUT_SIZE = (1, 2, 4)