Changed the naming. Added additional tests

Author: AlexKoff88

optimum/intel/openvino/configuration.py

@@ -19,6 +19,8 @@
 
 from optimum.configuration_utils import BaseConfig
 
+from .weight_quantization import OVWeightQuantizationConfig
+
 
 DEFAULT_QUANTIZATION_CONFIG = {
     "algorithm": "quantization",
@@ -84,7 +86,7 @@ def __init__(
         compression: Union[List[Dict], Dict, None] = None,
         input_info: Optional[List] = None,
         save_onnx_model: bool = False,
-        weight_quantization_config: Optional[QuantizationConfigMixin] = None,
+        quantization_config: Optional[QuantizationConfigMixin] = None,
         **kwargs,
     ):
         super().__init__()
@@ -93,7 +95,7 @@ def __init__(
         self.save_onnx_model = save_onnx_model
         self._enable_standard_onnx_export_option()
         self.optimum_version = kwargs.pop("optimum_version", None)
-        self.weight_quantization_config = weight_quantization_config
+        self.quantization_config = quantization_config
 
     def add_input_info(self, model_inputs: Dict, force_batch_one: bool = False):
         self.input_info = [
@@ -105,6 +107,11 @@ def add_input_info(self, model_inputs: Dict, force_batch_one: bool = False):
             for name, value in model_inputs.items()
         ]
 
+    def save_pretrained(self, *args, **kwargs):
+        if self.quantization_config is None:
+            self.quantization_config = OVWeightQuantizationConfig()
+        super().save_pretrained(*args, **kwargs)
+
     def _enable_standard_onnx_export_option(self):
         # This method depends on self.save_onnx_model.
         # save_onnx_model is defaulted to false so that the final model output is

optimum/intel/openvino/modeling_base.py

@@ -290,6 +290,7 @@ def _from_transformers(
         save_dir = TemporaryDirectory()
         save_dir_path = Path(save_dir.name)
 
+        # If load_in_8bit is not specified then compression_option should be set to None and will be set by default in main_export depending on the model size
         compression_option = None
         if load_in_8bit is not None:
             compression_option = "fp32"

optimum/intel/openvino/modeling_decoder.py

@@ -262,6 +262,7 @@ def _from_transformers(
             if use_cache:
                 task = task + "-with-past"
 
+        # If load_in_8bit is not specified then compression_option should be set to None and will be set by default in main_export depending on the model size
         compression_option = None
         if load_in_8bit is not None or load_in_4bit is not None:
             compression_option = "fp32"

optimum/intel/openvino/quantization.py

@@ -49,7 +49,7 @@
     ONNX_WEIGHTS_NAME,
     OV_XML_FILE_NAME,
 )
-from .weight_quantization import compress_decoder_weights
+from .weight_quantization import OVWeightQuantizationConfig, compress_decoder_weights
 
 
 COMPRESSION_OPTIONS = {
@@ -318,12 +318,14 @@ def _quantize_ovcausallm(
         save_directory.mkdir(parents=True, exist_ok=True)
 
         if weights_only:
-            quantization_config = None if ov_config is None else ov_config.weight_quantization_config
+            quantization_config = None if ov_config is None else ov_config.quantization_config
             if quantization_config is None:
                 # Use default 8-bit compression
+                quantization_config = OVWeightQuantizationConfig(mode=nncf.CompressWeightsMode.INT8_SYM)
                 self.model.model = nncf.compress_weights(self.model.model)
             else:
                 compress_decoder_weights(self.model, quantization_config)
+
             self.model.save_pretrained(save_directory)
             return
 

optimum/intel/openvino/weight_quantization.py

@@ -59,7 +59,7 @@ class OVWeightQuantizationConfig(QuantizationConfigMixin):
 
     def __init__(
         self,
-        mode=nncf.CompressWeightsMode.INT4_ASYM,
+        mode=None,
         tokenizer: Any = None,
         dataset: Optional[Union[nncf.Dataset, str]] = None,
         ratio: Optional[float] = None,

tests/openvino/test_quantization.py

@@ -155,6 +155,7 @@ class OVWeightCompressionTest(unittest.TestCase):
     )
 
     SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_COMPRESSED_MATMULS = ((OVModelForCausalLM, "opt125m", 64, 365),)
+    SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_AUTOCOMPRESSED_MATMULS = ((OVModelForCausalLM, "opt125m", 6, 379),)
     SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_AUTO_COMPRESSED_MATMULS = (
         (OVModelForCausalLM, "hf-internal-testing/tiny-random-OPTForCausalLM", 16, 136),
     )
@@ -287,9 +288,7 @@ def test_ovmodel_4bit_weight_compression(self, model_cls, model_name, expected_i
 
             quantizer = OVQuantizer.from_pretrained(transformers_model, task=task)
             ov_config = OVConfig(
-                weight_quantization_config=OVWeightQuantizationConfig(
-                    mode=nncf.CompressWeightsMode.INT4_SYM, ratio=0.8
-                )
+                quantization_config=OVWeightQuantizationConfig(mode=nncf.CompressWeightsMode.INT4_SYM, ratio=0.8)
             )
             quantizer.quantize(
                 save_directory=tmp_dir,
@@ -330,25 +329,43 @@ def test_ovmodel_8bit_weight_compression_stateful(self, model_cls, model_id, exp
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_AUTO_COMPRESSION)
     def test_ovmodel_load_with_compressed_weights(self, model_cls, model_type):
-        model = model_cls.from_pretrained(MODEL_NAMES[model_type], export=True, load_in_8bit=True, stateful=False)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model = model_cls.from_pretrained(MODEL_NAMES[model_type], export=True, load_in_8bit=True, stateful=False)
+
+            if model.export_feature.startswith("text2text-generation"):
+                models = [model.encoder, model.decoder, model.decoder_with_past]
+            elif model.export_feature.startswith("stable-diffusion"):
+                models = [model.unet, model.vae_encoder, model.vae_decoder]
+                models.append(
+                    model.text_encoder if model.export_feature == "stable-diffusion" else model.text_encoder_2
+                )
+            else:
+                models = [model]
 
-        if model.export_feature.startswith("text2text-generation"):
-            models = [model.encoder, model.decoder, model.decoder_with_past]
-        elif model.export_feature.startswith("stable-diffusion"):
-            models = [model.unet, model.vae_encoder, model.vae_decoder]
-            models.append(model.text_encoder if model.export_feature == "stable-diffusion" else model.text_encoder_2)
-        else:
-            models = [model]
+            expected_ov_int8 = _ARCHITECTURES_TO_EXPECTED_INT8[model_type]
+            for i, model in enumerate(models):
+                _, num_int8, _ = get_num_quantized_nodes(model)
+                self.assertEqual(expected_ov_int8[i], num_int8)
+            model.save_pretrained(tmp_dir)
 
-        expected_ov_int8 = _ARCHITECTURES_TO_EXPECTED_INT8[model_type]
-        for i, model in enumerate(models):
-            _, num_int8, _ = get_num_quantized_nodes(model)
-            self.assertEqual(expected_ov_int8[i], num_int8)
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_AUTOCOMPRESSED_MATMULS)
+    def test_ovmodel_4bit_auto_compression(self, model_cls, model_type, expected_ov_int8, expected_ov_int4):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model_id = MODEL_NAMES[model_type]
+            model = model_cls.from_pretrained(model_id, export=True, load_in_4bit=True)
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
 
-    @parameterized.expand(LOAD_IN_4_BITS_SCOPE)
-    def test_ovmodel_4bit_auto_compression(self, model_cls, model_id, quantization_config, expected_ov_int4):
-        task = model_cls.export_feature
+            _, num_int8, num_int4 = get_num_quantized_nodes(model)
+            self.assertEqual(expected_ov_int4, num_int4)
+            self.assertEqual(expected_ov_int8, num_int8)
+            model.save_pretrained(tmp_dir)
 
+    @parameterized.expand(LOAD_IN_4_BITS_SCOPE)
+    def test_ovmodel_4bit_auto_compression_with_config(
+        self, model_cls, model_id, quantization_config, expected_ov_int4
+    ):
         with tempfile.TemporaryDirectory() as tmp_dir:
             model = model_cls.from_pretrained(
                 model_id, export=True, load_in_4bit=True, quantization_config=quantization_config
@@ -359,6 +376,7 @@ def test_ovmodel_4bit_auto_compression(self, model_cls, model_id, quantization_c
 
             _, num_int4, _ = get_num_quantized_nodes(model)
             self.assertEqual(expected_ov_int4, num_int4)
+            model.save_pretrained(tmp_dir)
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_AUTO_COMPRESSED_MATMULS)
     def test_ovmodel_4bit_auto_compression_with_custom_dataset(