Fixed issues with models larger than 1B. Added tests.

Author: AlexKoff88

optimum/exporters/openvino/convert.py

@@ -95,6 +95,7 @@ def _save_model(model, path: str, compression_option: Optional[str] = None, comp
                 "ratio": compression_ratio,
             },
         }
+
         model = nncf.compress_weights(model, **COMPRESSION_OPTIONS[compression_option])
 
     compress_to_fp16 = compression_option == "fp16"

optimum/intel/openvino/modeling_base.py

@@ -287,7 +287,7 @@ def _from_transformers(
 
         compression_option = None
         if load_in_8bit is not None:
-            compression_option = "int8" if load_in_8bit else "fp32"
+            compression_option = "fp32"
 
         main_export(
             model_name_or_path=model_id,
@@ -304,7 +304,7 @@ def _from_transformers(
         )
 
         config.save_pretrained(save_dir_path)
-        return cls._from_pretrained(model_id=save_dir_path, config=config, load_in_8bit=False, **kwargs)
+        return cls._from_pretrained(model_id=save_dir_path, config=config, load_in_8bit=load_in_8bit, **kwargs)
 
     @classmethod
     def _to_load(

optimum/intel/openvino/modeling_decoder.py

@@ -264,8 +264,8 @@ def _from_transformers(
                 task = task + "-with-past"
 
         compression_option = None
-        if load_in_8bit is not None and not load_in_4bit:
-            compression_option = "int8" if load_in_8bit else "fp32"
+        if load_in_8bit is not None or load_in_4bit is not None:
+            compression_option = "fp32"
         stateful = kwargs.pop("stateful", ensure_stateful_is_available(warn=False) and use_cache)
         main_export(
             model_name_or_path=model_id,
@@ -574,7 +574,7 @@ def _from_pretrained(
             local_files_only=local_files_only,
         )
 
-        model = cls.load_model(model_cache_path, load_in_8bit=load_in_8bit)
+        model = cls.load_model(model_cache_path, load_in_8bit=False if load_in_4bit else load_in_8bit)
 
         model_type = config.model_type.replace("_", "-")
         if model_type == "bloom":

optimum/intel/openvino/weight_quantization.py

@@ -141,7 +141,7 @@ def compress_decoder_weights(model, quantization_config: Union[OVWeightQuantizat
 
             from optimum.gptq.data import get_dataset, prepare_dataset
 
-            dataset = get_dataset(config.dataset, tokenizer)
+            dataset = get_dataset(config.dataset, tokenizer, seqlen=32)
             dataset = prepare_dataset(dataset)
             dataset = nncf.Dataset(dataset, lambda x: model.prepare_forward_inputs(**x))
 

tests/openvino/test_quantization.py

@@ -22,6 +22,7 @@
 import numpy as np
 from datasets import load_dataset
 from parameterized import parameterized
+import openvino.runtime as ov
 import nncf
 from transformers import (
     AutoModelForQuestionAnswering,
@@ -154,7 +155,8 @@ class OVWeightCompressionTest(unittest.TestCase):
     )
 
     SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_COMPRESSED_MATMULS = ((OVModelForCausalLM, "opt125m", 64, 365),)
-    SUPPORTED_ARCHITECTURES_STATEFUL_WITH_EXPECTED_4BIT_COMPRESSED_MATMULS = (
+    SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_AUTO_COMPRESSED_MATMULS = ((OVModelForCausalLM, "hf-internal-testing/tiny-random-OPTForCausalLM", 16, 136),)
+    SUPPORTED_ARCHITECTURES_STATEFUL_WITH_EXPECTED_8BIT_COMPRESSED_MATMULS = (
         (OVModelForCausalLM, "hf-internal-testing/tiny-random-gpt2", 44, 46),
     )
 
@@ -170,7 +172,7 @@ class OVWeightCompressionTest(unittest.TestCase):
             "hf-internal-testing/tiny-random-gpt2",
             dict(
                 mode=nncf.CompressWeightsMode.INT4_ASYM,
-                group_size=-1,
+                group_size=32,
                 ignored_scope=nncf.IgnoredScope(names=["__module.model.transformer.h.2.mlp.c_fc/aten::addmm/MatMul"]),
             ),
             6,
@@ -297,7 +299,7 @@ def test_ovmodel_4bit_weight_compression(self, model_cls, model_name, expected_i
             outputs = model(**tokens)
             self.assertTrue("logits" in outputs)
 
-    @parameterized.expand(SUPPORTED_ARCHITECTURES_STATEFUL_WITH_EXPECTED_4BIT_COMPRESSED_MATMULS)
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_STATEFUL_WITH_EXPECTED_8BIT_COMPRESSED_MATMULS)
     @unittest.skipIf(not IS_SUPPORT_STATEFUL, "Stateful models supported only in 2023.3 and above")
     def test_ovmodel_8bit_weight_compression_stateful(self, model_cls, model_name, expected_pt_int8, expected_ov_int8):
         task = model_cls.export_feature
@@ -351,6 +353,35 @@ 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)
+            
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_AUTO_COMPRESSED_MATMULS)
+    def test_ovmodel_4bit_auto_compression_with_custom_dataset(self, model_cls, model_id, expected_int8, expected_int4):
+        task = model_cls.export_feature
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        dataset_name, dataset_config_name, column = _TASK_TO_DATASET[task]
+        dataset = load_dataset(dataset_name, dataset_config_name, split="test")
+        
+        def transform_fn(data, tokenizer):
+            tokenized_text = tokenizer(data[column], return_tensors="np")
+            input_ids = tokenized_text["input_ids"]
+            attention_mask = tokenized_text["attention_mask"]
+            inputs = {}
+            inputs["input_ids"] = input_ids
+            inputs["attention_mask"] = attention_mask
+            batch_size = input_ids.shape[0]
+            inputs["beam_idx"] = np.arange(batch_size, dtype=int)
+            return inputs
+
+        quantization_dataset = nncf.Dataset(dataset, partial(transform_fn, tokenizer=tokenizer))
+        model = model_cls.from_pretrained(model_id, export=True, load_in_4bit=True, quantization_config=OVWeightQuantizationConfig(mode=nncf.CompressWeightsMode.INT4_SYM, group_size=-1, ratio=0.8, dataset=quantization_dataset))
+
+        _, num_int8, num_int4 = get_num_quantized_nodes(model)
+        self.assertEqual(expected_int8, num_int8)
+        self.assertEqual(expected_int4, num_int4)
 
     @parameterized.expand(((OVModelForCausalLM, "gpt2"),))
     @unittest.skipIf(not IS_SUPPORT_STATEFUL, "Stateful models supported only in 2023.3 and above")