[WC] Introduce flexible group size value search (#3556)

### Changes

Introduce flexible group size search logic as a part of mixed precision
algorithm. 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` (16 by
default).

If it's not possible to find a value satisfying these requirements, such
weight is compressed to 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 won't contribute to the ratio of backup
mode group.

This method is disabled by default.

### Reason for changes

Some models may have channel size values that are not divisible by the
default group size. In such case a user can now provide
`nncf.AdvancedCompressionParameters(enable_flexible_group_size=True)`
advanced parameter instead of an ignored scope.

Example models:
- `microsoft/Phi-4-multimodal-instruct` (lm_model and
vision_embeddings_model)
- `HuggingFaceH4/Qwen2.5-Math-1.5B-Instruct-PRM-0.2`

### Metrics

Results for phi4-multimodal are below.

| Language Model Precision | Vision Embed. Model Precision | WWB
Similarity | Time of image-to-text request (sec.) | Time of
audio-to-text request (sec.) |

|------------------------------------|-------------------------------------|----------------|---------------------------------------|----------------------------------------|
| FP16 | FP16 | 99.19% | 31.21 | 17.76 |
| Mixed precision: int4 or bf16 | Mixed precision: int4 or bf16 | 77.51%
| 22.37 | 10.93 |
| Mixed precision: int4 or int8 | Mixed precision: int4 or int8 | 79.03%
| 19.95 | 9.47 |
| int4 with mixed group size: 128 or 64 | int4 with mixed group size:
128 or 16 | 81.36% | 19.89 | 9.16 |

Last row corresponds to
`nncf.AdvancedCompressionParameters(enable_flexible_group_size=True)`.

Third row corresponds to
`nncf.AdvancedCompressionParameters(enable_flexible_group_size=True,
min_flexible_group_size=128)`

Second row corresponds to
`nncf.AdvancedCompressionParameters(enable_flexible_group_size=True,
min_flexible_group_size=128)` with `backup_mode="none"`.

Inference time results are expected. Similarity not so much, but still
no degradation for group size 16 case.

### Related tickets

167337

### Tests

Added test cases which assert that the expected log messages are
printed.

https://github.com/openvinotoolkit/nncf/actions/runs/15852358755

Author: Nikita Savelyev

src/nncf/quantization/advanced_parameters.py

@@ -369,6 +369,22 @@ class AdvancedCompressionParameters:
 
     :param statistics_path: Directory path to dump statistics.
     :type statistics_path: str
+    :param lora_adapter_rank: Rank of lora adapters for FQ_LORA format. Defaults to 256.
+    :type lora_adapter_rank: int
+    :param enable_flexible_group_size: Whether to enable flexible group size searching. 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 won't contribute to the ratio of backup mode group.
+    :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
     :param awq_params: Advanced parameters for AWQ algorithm.
     :type awq_params: AdvancedAWQParameters
     :param scale_estimation_params: Advanced parameters for Scale Estimation algorithm.
@@ -377,8 +393,6 @@ class AdvancedCompressionParameters:
     :type gptq_params: AdvancedGPTQParameters
     :param lora_correction_params: Advanced parameters for Lora Correction algorithm.
     :type lora_correction_params: AdvancedLoraCorrectionParameters
-    :param lora_adapter_rank: Rank of lora adapters for FQ_LORA format. Defaults to 256.
-    :type lora_adapter_rank: int
     :param backend_params: Backend-specific parameters.
     :type backend_params: dict[str, Any]
     :param codebook: The codebook (LUT) for the weight compression.
@@ -387,13 +401,15 @@ class AdvancedCompressionParameters:
     """
 
     statistics_path: Optional[str] = None
+    lora_adapter_rank: int = 256
+    enable_flexible_group_size: bool = False
+    min_flexible_group_size: int = 16
     awq_params: AdvancedAWQParameters = field(default_factory=AdvancedAWQParameters)
     scale_estimation_params: AdvancedScaleEstimationParameters = field(
         default_factory=AdvancedScaleEstimationParameters
     )
     gptq_params: AdvancedGPTQParameters = field(default_factory=AdvancedGPTQParameters)
     lora_correction_params: AdvancedLoraCorrectionParameters = field(default_factory=AdvancedLoraCorrectionParameters)
-    lora_adapter_rank: int = 256
     backend_params: dict[str, Any] = field(default_factory=dict)
     codebook: Optional[TTensor] = None
 

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

@@ -45,6 +45,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 import Tensor
@@ -318,11 +319,13 @@ def __init__(
             advanced_parameters if advanced_parameters is not None else AdvancedCompressionParameters()
         )
 
-        primary_config = self._get_primary_config()
         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(self._ratio, self._subset_size)
         self._statistics_path = self._advanced_parameters.statistics_path
 
+        self._enable_flexible_group_size = self._advanced_parameters.enable_flexible_group_size
+        self._min_flexible_group_size = self._advanced_parameters.min_flexible_group_size
+
         if self._awq:
             awq_params = self._advanced_parameters.awq_params
             self.awq_algo = AWQ(
@@ -454,7 +457,7 @@ def _get_ratio_defining_params(
 
         return ratio_defining_params
 
-    def _get_primary_config(self):
+    def _get_primary_config(self, group_size: int) -> WeightCompressionConfig:
         codebook_values = None
 
         if self._mode == CompressWeightsMode.CB4_F8E4M3:
@@ -464,7 +467,7 @@ def _get_primary_config(self):
 
         return WeightCompressionConfig(
             mode=self._mode,
-            group_size=self._group_size,
+            group_size=group_size,
             codebook_values=codebook_values,
         )
 
@@ -474,6 +477,7 @@ def _set_weight_compression_config(
         model: TModel,
         graph: NNCFGraph,
         statistics_points: StatisticPointsContainer,
+        group_size_values: dict[str, int],
     ) -> None:
         """
         Sets the appropriate compression configuration for weights based on some criteria.
@@ -483,13 +487,92 @@ def _set_weight_compression_config(
         :param model: The model.
         :param graph: The model graph associated with the model.
         :param statistics_points: Statistics points.
+        :param group_size_values: A dictionary mapping weight names to their group size values.
         """
-        primary_config = self._get_primary_config()
-        if self._ratio == 1:
-            for weight_param in ratio_defining_params:
-                weight_param.compression_config = primary_config
+        if self._ratio < 1 and len(ratio_defining_params) > 0:
+            primary_precision_weight_params = self._mixed_precision_algo.apply(
+                model, graph, statistics_points, weight_params=ratio_defining_params
+            )
         else:
-            self._mixed_precision_algo.apply(model, graph, statistics_points, weight_params=ratio_defining_params)
+            primary_precision_weight_params = ratio_defining_params
+
+        for weight_param in primary_precision_weight_params:
+            weight_param.compression_config = self._get_primary_config(group_size_values[weight_param.weight_name])
+
+        # 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)
+
+    def _get_flexible_group_size_data(
+        self, weight_params: list[WeightCompressionParameters]
+    ) -> list[tuple[WeightCompressionParameters, int]]:
+        """
+        Compute flexible group size values.
+        :param weight_params: Weight parameters for which to compute flexible group size.
+        :return: A list of tuples, where each tuple pair contains a WeightCompressionParameters object and the
+            group size values associated with it. If group size can't be assigned to some weight parameter
+            it won't be included in the result.
+        """
+        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 % self._group_size == 0:
+                # The weight can be compressed with the given group size, nothing else to do
+                group_size_data.append((w_params, self._group_size))
+                continue
+
+            # Find the maximal power of two that divides reduction_channel_size
+            flexible_group_size = reduction_channel_size & (~reduction_channel_size + 1)
+
+            if flexible_group_size < self._min_flexible_group_size:
+                flexible_group_size_not_found_weight_params.append(w_params)
+            else:
+                group_size_data.append((w_params, flexible_group_size))
+
+        node_strings = []
+        for i, (w_params, new_group_size) in enumerate(group_size_data):
+            if new_group_size == self._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 ({self._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 "
+                "according to the backup mode. Nodes:\n\t" + "\n\t".join(node_strings)
+            )
+
+        return group_size_data
 
     @staticmethod
     def _proportion_str(num_weights_list: list[int], total_num_weights: int, total_num_params: int) -> str:
@@ -625,7 +708,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
@@ -654,13 +736,21 @@ 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)
 
         ratio_defining_params = self._get_ratio_defining_params(all_weight_params, is_last_layer_shared)
-        self._set_weight_compression_config(ratio_defining_params, model, graph, statistic_points)
+        if self._enable_flexible_group_size and self._group_size != -1:
+            # Compute flexible group size values if enabled
+            flexible_group_size_data = self._get_flexible_group_size_data(ratio_defining_params)
+            group_size_values = {w_param.weight_name: group_size for w_param, group_size in flexible_group_size_data}
+            # Select a subset of ratio_defining_params that can be compressed with some group size
+            ratio_defining_params = [w_param for w_param, _ in flexible_group_size_data]
+        else:
+            group_size_values = {w_param.weight_name: self._group_size for w_param in ratio_defining_params}
+        self._set_weight_compression_config(ratio_defining_params, model, graph, statistic_points, group_size_values)
         ignored_scope_weight_statistics = self._get_ignored_scope_weight_statistics(model, graph)
         nncf_logger.info(
             self._get_bitwidth_distribution_str(

src/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
@@ -86,19 +88,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

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

@@ -41,18 +41,16 @@
 
 class MixedPrecisionCriterion(Algorithm):
     """
-    Assigns mixed quantization scheme (e.g. uniform int8 or uniform int4/non-uniform fp4)
+    Computes mixed quantization scheme (e.g. uniform int8 or uniform int4/non-uniform fp4)
     for weights based on some criteria.
     """
 
-    def __init__(self, primary_config: WeightCompressionConfig, ratio: float, subset_size: Optional[int] = None):
+    def __init__(self, ratio: float, subset_size: Optional[int] = 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.
         """
-        self._primary_config = primary_config
         self._ratio = ratio
         self._subset_size = subset_size
         self._algorithm_key = f"MPC_{hash(self)}"
@@ -79,15 +77,17 @@ def apply(
         statistic_points: Optional[StatisticPointsContainer] = None,
         dataset: Optional[Dataset] = None,
         weight_params: list[WeightCompressionParameters] = None,
-    ) -> None:
+    ) -> list[WeightCompressionParameters]:
         """
-        Assigns quantization precision based on computed layers' sensitivities, ratio of parameters.
+        Selects which weights should be compressed to a primary (4 bit) precision based on computed layers'
+        sensitivities, ratio of parameters.
         """
         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)
 
+        primary_precision_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])
         ]
@@ -97,8 +97,9 @@ def apply(
             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
+            primary_precision_weight_params.append(weight_param)
             num_weights_in_4bit += weight_param.num_weights
+        return primary_precision_weight_params
 
     @abstractmethod
     def _set_backend_entity(self, model: TModel) -> None: