[WC] Introduce flexible group size value search

nncf/__init__.py

@@ -54,6 +54,7 @@
 from nncf.quantization.advanced_parameters import AdvancedBiasCorrectionParameters as AdvancedBiasCorrectionParameters
 from nncf.quantization.advanced_parameters import AdvancedCompressionParameters as AdvancedCompressionParameters
 from nncf.quantization.advanced_parameters import AdvancedGPTQParameters as AdvancedGPTQParameters
+from nncf.quantization.advanced_parameters import AdvancedGroupSizeParameters as AdvancedGroupSizeParameters
 from nncf.quantization.advanced_parameters import AdvancedLoraCorrectionParameters as AdvancedLoraCorrectionParameters
 from nncf.quantization.advanced_parameters import AdvancedQuantizationParameters as AdvancedQuantizationParameters
 from nncf.quantization.advanced_parameters import AdvancedScaleEstimationParameters as AdvancedScaleEstimationParameters

nncf/quantization/advanced_parameters.py

@@ -359,6 +359,32 @@ class AdvancedLoraCorrectionParameters:
     use_int8_adapters: bool = True
 
 
+@api()
+@dataclass
+class AdvancedGroupSizeParameters:
+    """
+    Contains advanced parameters for flexible group size searching logic. When enabled, each weight for which the
+    channel size is not divisible by the general group size value will be compressed to a newly calculated group size.
+    The new group size value is the maximal power of two (i.e., 2^k) such that:
+        - channel size is divisible by it;
+        - it is less than the originally specified group size value;
+        - it is greater than or equal to `min_flexible_group_size`.
+
+    If it's not possible to find a value satisfying these requirements, such weight is compressed to the backup
+    precision. If ratio < 1.0 and some weights have to be compressed to the backup precision because of group size
+    issues, then these weights also contribute to the ratio of backup mode group.
+
+    :param enable_flexible_group_size: Whether to enable flexible group size searching.
+    :type enable_flexible_group_size: bool
+    :param min_flexible_group_size: Minimum group size for flexible group size searching. Defaults to 16. The reason
+        behind this argument is to avoid too small group size values, which may lead to performance issues.
+    :type min_flexible_group_size: int
+    """
+
+    enable_flexible_group_size: bool = False
+    min_flexible_group_size: int = 16
+
+
 @api()
 @dataclass
 class AdvancedCompressionParameters:
@@ -390,6 +416,7 @@ class AdvancedCompressionParameters:
     lora_correction_params: AdvancedLoraCorrectionParameters = field(default_factory=AdvancedLoraCorrectionParameters)
     lora_adapter_rank: int = 256
     backend_params: dict[str, Any] = field(default_factory=dict)
+    group_size_params: AdvancedGroupSizeParameters = field(default_factory=AdvancedGroupSizeParameters)
 
 
 @api()

nncf/quantization/algorithms/weight_compression/algorithm.py

@@ -44,6 +44,7 @@
 from nncf.quantization.algorithms.weight_compression.mixed_precision import MIXED_PRECISION_CRITERIA
 from nncf.quantization.algorithms.weight_compression.scale_estimation import ScaleEstimation
 from nncf.quantization.algorithms.weight_compression.weight_lowering import WeightCompressionConfig
+from nncf.quantization.algorithms.weight_compression.weight_lowering import get_reduction_channel_size
 from nncf.scopes import IgnoredScope
 from nncf.scopes import get_ignored_node_names_from_ignored_scope
 from nncf.tensor.definitions import TensorDataType
@@ -295,7 +296,9 @@ def __init__(
 
         primary_config = WeightCompressionConfig(mode=self._mode, group_size=self._group_size)
         criterion_cls = MIXED_PRECISION_CRITERIA.get(self._sensitivity_metric)
-        self._mixed_precision_algo = criterion_cls(primary_config, self._ratio, self._subset_size)
+        self._mixed_precision_algo = criterion_cls(
+            primary_config, self._ratio, self._subset_size, self._advanced_parameters.group_size_params
+        )
         self._statistics_path = self._advanced_parameters.statistics_path
 
         if self._awq:
@@ -445,12 +448,26 @@ def _set_weight_compression_config(
         :param graph: The model graph associated with the model.
         :param statistics_points: Statistics points.
         """
-        primary_config = WeightCompressionConfig(mode=self._mode, group_size=self._group_size)
-        if self._ratio == 1:
-            for weight_param in ratio_defining_params:
-                weight_param.compression_config = primary_config
-        else:
-            self._mixed_precision_algo.apply(model, graph, statistics_points, weight_params=ratio_defining_params)
+        self._mixed_precision_algo.apply(model, graph, statistics_points, weight_params=ratio_defining_params)
+
+        # Check if group size is valid for each weight in ratio_defining_params
+        failed_nodes = []
+        for w_params in ratio_defining_params:
+            if w_params.compression_config is None or w_params.compression_config.group_size == -1:
+                continue
+            reduction_channel_size, _ = get_reduction_channel_size(w_params.weight_shape, w_params.reduction_axes)
+            if reduction_channel_size % w_params.compression_config.group_size != 0:
+                failed_nodes.append((w_params.node_with_weight.node_name, reduction_channel_size))
+        if len(failed_nodes) > 0:
+            names = ",".join(f'"{name}"' for name, _ in failed_nodes)
+            msg = (
+                "Failed to apply group-wise quantization with "
+                f"group size value {self._group_size} and channel size value {failed_nodes[0][1]}.\n"
+                "Ensure that the group size is divisible by the channel size, "
+                "or include this node and others with similar issues in the ignored scope:\n"
+                f"nncf.compress_weight(\n\t..., \n\tignored_scope=IgnoredScope(names=[{names}]\n\t)\n)"
+            )
+            raise nncf.InvalidGroupSizeError(msg)
 
     @staticmethod
     def _proportion_str(num_weights_list: list[int], total_num_weights: int, total_num_params: int) -> str:
@@ -586,7 +603,6 @@ def apply(
                 if weight_dtype not in SUPPORTED_DATA_TYPES:
                     continue
                 weight_shape = self._backend_entity.get_weight_shape(node, weight_port_id, graph)
-                weight_size = reduce(operator.mul, weight_shape, 1)
                 reduction_axes = self._backend_entity.get_reduction_axes(node, weight_port_id, graph)
                 if (
                     self._group_size != -1
@@ -615,7 +631,7 @@ def apply(
                     )
                     wc_config = WeightCompressionConfig(mode=mode)
                 weight_params = WeightCompressionParameters(
-                    weight_name, node, weight_port_id, weight_size, reduction_axes, wc_config
+                    weight_name, node, weight_port_id, weight_shape, reduction_axes, wc_config
                 )
                 all_weight_params.append(weight_params)
                 weight_names.add(weight_name)

nncf/quantization/algorithms/weight_compression/config.py

@@ -8,8 +8,10 @@
 # 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 operator
 from dataclasses import dataclass
 from dataclasses import field
+from functools import reduce
 from typing import Optional, TypeVar
 
 import numpy as np
@@ -66,19 +68,20 @@ class WeightCompressionParameters:
     :param weight_name: Unique weight name.
     :param node_with_weight: Node with weight in the NNCF graph.
     :param weight_port_id: Number of elements in the weight array.
-    :param num_weights: Number of elements in the weight array.
+    :param weight_shape: Shape of the weight array.
     :param reduction_axes: Axes, along which to reduce (collect) different statistics (e.g. min, max).
     :param compression_config: Configuration of weight compression for the weight node.
     """
 
     weight_name: str
     node_with_weight: NNCFNode
     weight_port_id: int
-    num_weights: np.uint64
+    weight_shape: tuple[int, ...]
     reduction_axes: tuple[int, ...]
     compression_config: Optional[WeightCompressionConfig] = field(default_factory=WeightCompressionConfig)
 
-    def __post_init__(self):
-        # Explicitly cast num_weights to avoid overflow on finding total number of weights.
-        # The issue happens on Windows, because np.ndarray.size() returns np.int32 and sum of weights is more than 2^32.
-        self.num_weights = np.uint64(self.num_weights)
+    @property
+    def num_weights(self) -> np.uint64:
+        if not hasattr(self, "_num_weights"):
+            self._num_weights = np.uint64(reduce(operator.mul, self.weight_shape, 1))
+        return self._num_weights

nncf/quantization/algorithms/weight_compression/mixed_precision.py

@@ -8,13 +8,14 @@
 # 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 dataclasses
 from abc import ABC
 from abc import abstractmethod
 from typing import Iterable, Optional, TypeVar
 
 import nncf
 from nncf import Dataset
+from nncf import nncf_logger
 from nncf.common.graph import NNCFGraph
 from nncf.common.graph import NNCFNode
 from nncf.common.graph.transformations.commands import TargetType
@@ -25,10 +26,12 @@
 from nncf.common.utils.backend import get_backend
 from nncf.common.utils.registry import Registry
 from nncf.parameters import SensitivityMetric
+from nncf.quantization.advanced_parameters import AdvancedGroupSizeParameters
 from nncf.quantization.algorithms.algorithm import Algorithm
 from nncf.quantization.algorithms.weight_compression.config import WeightCompressionConfig
 from nncf.quantization.algorithms.weight_compression.config import WeightCompressionParameters
 from nncf.quantization.algorithms.weight_compression.weight_lowering import get_integer_quantization_error
+from nncf.quantization.algorithms.weight_compression.weight_lowering import get_reduction_channel_size
 from nncf.quantization.algorithms.weight_compression.weight_lowering import integer_quantize_dequantize_weight
 from nncf.tensor import Tensor
 from nncf.tensor import functions as fns
@@ -45,16 +48,27 @@ class MixedPrecisionCriterion(Algorithm):
     for weights based on some criteria.
     """
 
-    def __init__(self, primary_config: WeightCompressionConfig, ratio: float, subset_size: Optional[int] = None):
+    def __init__(
+        self,
+        primary_config: WeightCompressionConfig,
+        ratio: float,
+        subset_size: Optional[int] = None,
+        group_size_parameters: Optional[AdvancedGroupSizeParameters] = None,
+    ):
         """
         :param primary_config: Configuration on how to compress (quantize) weights to primary precision.
         :param ratio: The ratio between primary and backup precisions (e.g. 0.9 means 90% of layers quantized to NF4
             and the rest to INT8_ASYM).
         :param subset_size: Size of dataset subset for statistics.
         """
+        if group_size_parameters is None:
+            group_size_parameters = AdvancedGroupSizeParameters()
+
         self._primary_config = primary_config
         self._ratio = ratio
         self._subset_size = subset_size
+        self._enable_flexible_group_size = group_size_parameters.enable_flexible_group_size
+        self._min_flexible_group_size = group_size_parameters.min_flexible_group_size
         self._algorithm_key = f"MPC_{hash(self)}"
         self._backend_entity = None
 
@@ -85,20 +99,44 @@ def apply(
         """
         self._set_backend_entity(model)
 
-        scores = self._calc_sensitivity(model, graph, weight_params, statistic_points)
-        num_all_weights = sum(wp.num_weights for wp in weight_params)
-
-        indexes_of_layers_in_ascending_order_of_scores = [
-            i[0] for i in sorted(enumerate(scores), reverse=False, key=lambda x: x[1])
-        ]
-        num_weights_in_4bit = 0
-        for index in indexes_of_layers_in_ascending_order_of_scores:
-            weight_param = weight_params[index]
-            current_ratio = (num_weights_in_4bit + weight_param.num_weights) / num_all_weights
-            if current_ratio >= self._ratio:
-                break
-            weight_param.compression_config = self._primary_config
-            num_weights_in_4bit += weight_param.num_weights
+        flexible_group_size_values = {}
+        valid_weight_params = weight_params
+        if self._enable_flexible_group_size and self._primary_config.group_size != -1:
+            group_size_data = self._get_flexible_group_size_data(weight_params)
+            flexible_group_size_values = {w_param.weight_name: group_size for w_param, group_size in group_size_data}
+            valid_weight_params = [w_param for w_param, _ in group_size_data]
+
+        if self._ratio == 1.0:
+            for weight_param in valid_weight_params:
+                weight_param.compression_config = self._primary_config
+                if weight_param.weight_name in flexible_group_size_values:
+                    weight_param.compression_config = dataclasses.replace(
+                        weight_param.compression_config,
+                        group_size=flexible_group_size_values[weight_param.weight_name],
+                    )
+        else:
+            scores = self._calc_sensitivity(model, graph, valid_weight_params, statistic_points)
+
+            # Sum all weights to calculate the ratio. This way the weights for which we weren't able to find a flexible
+            # group size value will contribute to the backup group as well.
+            num_all_weights = sum(wp.num_weights for wp in weight_params)
+
+            indexes_of_layers_in_ascending_order_of_scores = [
+                i[0] for i in sorted(enumerate(scores), reverse=False, key=lambda x: x[1])
+            ]
+            num_weights_in_4bit = 0
+            for index in indexes_of_layers_in_ascending_order_of_scores:
+                weight_param = valid_weight_params[index]
+                current_ratio = (num_weights_in_4bit + weight_param.num_weights) / num_all_weights
+                if current_ratio >= self._ratio:
+                    break
+                weight_param.compression_config = self._primary_config
+                if weight_param.weight_name in flexible_group_size_values:
+                    weight_param.compression_config = dataclasses.replace(
+                        weight_param.compression_config,
+                        group_size=flexible_group_size_values[weight_param.weight_name],
+                    )
+                num_weights_in_4bit += weight_param.num_weights
 
     @abstractmethod
     def _set_backend_entity(self, model: TModel) -> None:
@@ -124,6 +162,59 @@ def get_statistic_points(
         :return: Statistic points, for which StatisticsCollector should collect statistics.
         """
 
+    def _get_flexible_group_size_data(
+        self, weight_params: list[WeightCompressionParameters]
+    ) -> list[tuple[WeightCompressionParameters, int]]:
+        primary_group_size = self._primary_config.group_size
+        flexible_group_size_not_found_weight_params = []
+        group_size_data = []
+        for w_params in weight_params:
+            reduction_channel_size, _ = get_reduction_channel_size(w_params.weight_shape, w_params.reduction_axes)
+            if reduction_channel_size % primary_group_size == 0:
+                # The weight can be compressed with the given group size, nothing else to do
+                group_size_data.append((w_params, primary_group_size))
+                continue
+
+            # Find the maximal power of two that divides reduction_channel_size
+            new_group_size = 2
+            while reduction_channel_size % new_group_size == 0 and new_group_size < primary_group_size:
+                new_group_size *= 2
+            new_group_size //= 2
+
+            if new_group_size < self._min_flexible_group_size:
+                flexible_group_size_not_found_weight_params.append(w_params)
+            else:
+                group_size_data.append((w_params, new_group_size))
+
+        node_strings = []
+        for i, (w_params, new_group_size) in enumerate(group_size_data):
+            if new_group_size == primary_group_size:
+                continue
+            weight_shape = w_params.weight_shape
+            reduction_channel_size, _ = get_reduction_channel_size(weight_shape, w_params.reduction_axes)
+            node_strings.append(
+                f"{w_params.node_with_weight.node_name} "
+                f"(weight shape: {weight_shape}, adjusted group size: {new_group_size})"
+            )
+        if len(node_strings) > 0:
+            nncf_logger.info(
+                f"Wasn't able to set the specified group size value ({primary_group_size}) to some nodes. These nodes "
+                f"will have an adjusted group size value:\n\t" + "\n\t".join(node_strings)
+            )
+
+        if len(flexible_group_size_not_found_weight_params) > 0:
+            node_strings = [""] * len(flexible_group_size_not_found_weight_params)
+            for i, w_params in enumerate(flexible_group_size_not_found_weight_params):
+                weight_shape = w_params.weight_shape
+                reduction_channel_size, _ = get_reduction_channel_size(weight_shape, w_params.reduction_axes)
+                node_strings[i] = f"{w_params.node_with_weight.node_name} (weight shape: {weight_shape})"
+            nncf_logger.warning(
+                "Large enough flexible group size value cannot be found for some nodes. They will be compressed to the "
+                "backup mode. Nodes:\n\t" + "\n\t".join(node_strings)
+            )
+
+        return group_size_data
+
 
 @MIXED_PRECISION_CRITERIA.register(SensitivityMetric.WEIGHT_QUANTIZATION_ERROR)
 class DataFreeCriterion(MixedPrecisionCriterion):