Address comments

Author: nikita-savelyevv

optimum/intel/openvino/configuration.py

@@ -18,7 +18,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 import torch
 from transformers.utils.quantization_config import QuantizationConfigMixin
@@ -571,9 +571,7 @@ def to_nncf_dict(self) -> Dict[str, Any]:
             mode = "e2m1"
         mode = nncf.CompressWeightsMode(mode)
 
-        awq = None
-        if self.quant_method == "awq" or self.quant_method == OVQuantizationMethod.AWQ:
-            awq = True
+        awq = True if self.quant_method == OVQuantizationMethod.AWQ else None
         sensitivity_metric = nncf.SensitivityMetric(self.sensitivity_metric) if self.sensitivity_metric else None
         backup_mode = nncf.BackupMode(self.backup_precision) if self.backup_precision else None
         result = {
@@ -896,21 +894,22 @@ def __init__(
                 machine arbitrary code present in the model repository.
             **kwargs:
         """
-        if isinstance(weight_quantization_config, dict):
-            weight_quantization_config = OVWeightQuantizationConfig.from_dict(weight_quantization_config)
-        else:
-            weight_quantization_config = weight_quantization_config.clone()
-        self.weight_quantization_config = weight_quantization_config
+        self.weight_quantization_config = self._initialize_quantization_config(
+            weight_quantization_config, OVWeightQuantizationConfig
+        )
         wqc = self.weight_quantization_config
 
-        if isinstance(full_quantization_config, dict):
-            full_quantization_config = OVQuantizationConfig.from_dict(full_quantization_config)
-        else:
-            full_quantization_config = full_quantization_config.clone()
-        self.full_quantization_config = full_quantization_config
+        self.full_quantization_config = self._initialize_quantization_config(
+            full_quantization_config, OVQuantizationConfig
+        )
         fqc = self.full_quantization_config
 
         if fqc.dtype in ["f8e4m3", "f8e5m2"] and wqc.backup_precision is None:
+            # Here we simulate FP8 backup weight compression precision through full quantization: during weight
+            # compression step some weighted layers are kept in original precision and later are compressed to FP8
+            # during full precision quantization step.
+            # The issue with current approach is that if one provides an ignored scope for the full quantization step,
+            # then the weights of the layers under this ignored scope won't be compressed to FP8.
             # TODO: remove once there is support for FP8 weight compression in NNCF
             wqc.backup_precision = "none"
 
@@ -932,6 +931,21 @@ def __init__(
 
         self.post_init()
 
+    @staticmethod
+    def _initialize_quantization_config(
+        config: Union[dict, OVWeightQuantizationConfig, OVQuantizationConfig],
+        config_type: Type[Union[OVWeightQuantizationConfig, OVQuantizationConfig]],
+    ):
+        if isinstance(config, dict):
+            return config_type.from_dict(config)
+        elif isinstance(config, config_type):
+            return config.clone()
+        else:
+            raise ValueError(
+                f"Unsupported type of quantization config. Expected either a dictionary or an instance of "
+                f"{config_type}, but found: {type(config)}."
+            )
+
     def to_dict(self):
         result = super().to_dict()
         result["weight_quantization_config"] = self.weight_quantization_config.to_dict()

optimum/intel/openvino/quantization.py

@@ -1014,7 +1014,6 @@ def _weight_only_quantization(
     model: openvino.runtime.Model,
     quantization_config: Union[OVWeightQuantizationConfig, Dict],
     calibration_dataset: Optional[Union[nncf.Dataset, Iterable]] = None,
-    remove_kv_cache_precision_flag: Optional[bool] = True,
     **kwargs,
 ) -> openvino.runtime.Model:
     _verify_not_optimized(model)
@@ -1043,13 +1042,7 @@ def _weight_only_quantization(
         **wc_kwargs,
     )
 
-    if remove_kv_cache_precision_flag:
-        # Remove the KV cache compression disabling flag from the model
-        if compressed_model.has_rt_info(["runtime_options", "KV_CACHE_PRECISION"]):
-            prev_rt_info = compressed_model.get_rt_info("runtime_options").value
-            if prev_rt_info["KV_CACHE_PRECISION"] == "f16":
-                prev_rt_info.pop("KV_CACHE_PRECISION")
-                compressed_model.set_rt_info(prev_rt_info, "runtime_options")
+    _remove_f16_kv_cache_precision_flag(compressed_model)
 
     return compressed_model
 
@@ -1065,11 +1058,11 @@ def _full_quantization(
         _verify_not_optimized(model)
     q_kwargs = copy.deepcopy(kwargs)
     q_kwargs.update(quantization_config.to_nncf_dict())
-    return nncf.quantize(
-        model,
-        calibration_dataset=calibration_dataset,
-        **q_kwargs,
-    )
+    quantized_model = nncf.quantize(model, calibration_dataset=calibration_dataset, **q_kwargs)
+
+    _remove_f16_kv_cache_precision_flag(quantized_model)
+
+    return quantized_model
 
 
 def _get_operation_const_op(operation, const_port_id: int):
@@ -1201,9 +1194,7 @@ def merge_ignored_scopes(
     wc_config = quantization_config.weight_quantization_config.clone()
     wc_config.ignored_scope = merge_ignored_scopes(wc_config.ignored_scope, quantization_config.ignored_scope)
     wc_dataset = dataset if wc_config.bits != 8 else None
-    compressed_model = _weight_only_quantization(
-        model, wc_config, wc_dataset, remove_kv_cache_precision_flag=False, **kwargs
-    )
+    compressed_model = _weight_only_quantization(model, wc_config, wc_dataset, **kwargs)
 
     q_config = quantization_config.full_quantization_config.clone()
     q_config.ignored_scope = merge_ignored_scopes(q_config.ignored_scope, quantization_config.ignored_scope)
@@ -1227,3 +1218,13 @@ def _verify_not_optimized(ov_model):
             raise RuntimeError(message_template.format(model_weight_compression_config))
         elif model_quantization_config is not None:
             raise RuntimeError(message_template.format(model_quantization_config))
+
+
+def _remove_f16_kv_cache_precision_flag(model: openvino.Model) -> openvino.Model:
+    # Remove the KV cache compression disabling flag from the model
+    if model.has_rt_info(["runtime_options", "KV_CACHE_PRECISION"]):
+        prev_rt_info = model.get_rt_info("runtime_options").value
+        if prev_rt_info["KV_CACHE_PRECISION"] == "f16":
+            prev_rt_info.pop("KV_CACHE_PRECISION")
+            model.set_rt_info(prev_rt_info, "runtime_options")
+    return model

tests/openvino/test_exporters_cli.py

@@ -507,14 +507,11 @@ def test_exporters_cli_full_quantization(
                 submodels = [model.encoder, model.decoder]
                 if model.decoder_with_past is not None:
                     submodels.append(model.decoder_with_past)
-                    expected_kv_cache_precision_per_model = [None, None, None]
                 else:
                     expected_num_weight_nodes_per_model = expected_num_weight_nodes_per_model[:-1]
                     expected_fake_nodes_per_model = expected_fake_nodes_per_model[:-1]
-                    expected_kv_cache_precision_per_model = [None, "f16"]
             elif "text-generation" in task:
                 submodels = [model]
-                expected_kv_cache_precision_per_model = ["f16"]
             else:
                 raise Exception("Unexpected task.")
 
@@ -523,7 +520,6 @@ def test_exporters_cli_full_quantization(
                 submodels,
                 expected_num_weight_nodes_per_model,
                 expected_fake_nodes_per_model,
-                expected_kv_cache_precision_per_model,
             )
 
     def test_exporters_cli_int4_with_local_model_and_default_config(self):

tests/openvino/test_quantization.py

@@ -342,17 +342,14 @@ def test_ov_model_static_quantization_with_auto_dataset(
                 submodels = [ov_model.encoder.model, ov_model.decoder.model]
                 if ov_model.decoder_with_past is not None:
                     submodels.append(ov_model.decoder_with_past.model)
-                    expected_kv_cache_precision_per_model = [None, None, None]
                 else:
                     expected_num_weight_nodes_per_model = expected_num_weight_nodes_per_model[:-1]
                     expected_fake_nodes_per_model = expected_fake_nodes_per_model[:-1]
-                    expected_kv_cache_precision_per_model = [None, "f16"]
 
                 input_features = torch.randn((1, 128, 3000), dtype=torch.float32)
                 ov_model.generate(input_features)
             elif model_cls == OVModelForCausalLM:
                 submodels = [ov_model]
-                expected_kv_cache_precision_per_model = ["f16"]
 
                 tokenizer = AutoTokenizer.from_pretrained(model_id)
                 if tokenizer.pad_token is None:
@@ -368,7 +365,6 @@ def test_ov_model_static_quantization_with_auto_dataset(
                 submodels,
                 expected_num_weight_nodes_per_model,
                 expected_fake_nodes_per_model,
-                expected_kv_cache_precision_per_model,
             )
 
 

tests/openvino/utils_tests.py

@@ -295,12 +295,10 @@ def check_compression_state_per_model(
     models: List[Union[ov.Model, OVBaseModel]],
     expected_num_weight_nodes_per_model: List[Dict[str, int]],
     expected_num_fake_nodes_per_model: Optional[List[int]] = None,
-    expected_kv_cache_precision_per_model: Optional[List[Union[str, None]]] = None,
 ):
     test_case.assertEqual(len(models), len(expected_num_weight_nodes_per_model))
     actual_num_weights_per_model = [{}] * len(models)
     actual_num_fake_nodes_per_model = [0] * len(models)
-    actual_kv_cache_precision_per_model = [None] * len(models)
     for i, (submodel, expected_num_weight_nodes) in enumerate(zip(models, expected_num_weight_nodes_per_model)):
         ov_model = submodel if isinstance(submodel, ov.Model) else submodel.model
         num_fake_nodes, num_weight_nodes = get_num_quantized_nodes(ov_model)
@@ -309,19 +307,11 @@ def check_compression_state_per_model(
         actual_num_weights_per_model[i] = num_weight_nodes
         actual_num_fake_nodes_per_model[i] = num_fake_nodes
 
-        if ov_model.has_rt_info(["runtime_options", "KV_CACHE_PRECISION"]):
-            actual_kv_cache_precision = ov_model.get_rt_info(["runtime_options", "KV_CACHE_PRECISION"]).value
-        else:
-            actual_kv_cache_precision = None
-        actual_kv_cache_precision_per_model[i] = actual_kv_cache_precision
+        test_case.assertFalse(ov_model.has_rt_info(["runtime_options", "KV_CACHE_PRECISION"]))
 
     # Check weight nodes
     test_case.assertEqual(expected_num_weight_nodes_per_model, actual_num_weights_per_model)
 
     # Check fake nodes
     if expected_num_fake_nodes_per_model is not None:
         test_case.assertEqual(expected_num_fake_nodes_per_model, actual_num_fake_nodes_per_model)
-
-    # Check KV cache precision
-    expected_kv_cache_precision_per_model = expected_kv_cache_precision_per_model or ([None] * len(models))
-    test_case.assertEqual(expected_kv_cache_precision_per_model, actual_kv_cache_precision_per_model)