Add load_in_4bit option for OVModelForCausalLM (#538)

* Initial code for load_in_4_bit

* Dataset does not work

* Intermediate changes

* Make it working with dataset

* Style

* Fixed small issue

* Fixed failed tests

* Style

* Comment failed tests due to NNCF 2.8

* Commented failed tests until new NNCF release

* Added tests for load_in_4bit

* Added awq option. Included NNCF package into openvino extra.

* Rolled back including nncf into openvino extra

* Style

* Fixed tests

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

* Style

* Fixed issues. Applied comments.

* Removed unnecessary exception

* Applied more comments

* Fixed issue

* Make quantization_config a part of OVConfig in OVQuantizer

* Fixed issue with Transformers

* Fixed test

* Changed the naming. Added additional tests

* Fixed tests

* Fixed tests

* Applied more comments

* Style

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/__init__.py

@@ -62,9 +62,12 @@
         "OVQuantizer",
         "OVTrainer",
         "OVTrainingArguments",
+        "OVWeightQuantizationConfig",
     ]
 else:
-    _import_structure["openvino"].extend(["OVConfig", "OVQuantizer", "OVTrainer", "OVTrainingArguments"])
+    _import_structure["openvino"].extend(
+        ["OVConfig", "OVQuantizer", "OVTrainer", "OVTrainingArguments", "OVWeightQuantizationConfig"]
+    )
 
 try:
     if not (is_openvino_available() and is_diffusers_available()):
@@ -176,9 +179,15 @@
         if not (is_openvino_available() and is_nncf_available()):
             raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
-        from .utils.dummy_openvino_and_nncf_objects import OVConfig, OVQuantizer, OVTrainer, OVTrainingArguments
+        from .utils.dummy_openvino_and_nncf_objects import (
+            OVConfig,
+            OVQuantizer,
+            OVTrainer,
+            OVTrainingArguments,
+            OVWeightQuantizationConfig,
+        )
     else:
-        from .openvino import OVConfig, OVQuantizer, OVTrainer, OVTrainingArguments
+        from .openvino import OVConfig, OVQuantizer, OVTrainer, OVTrainingArguments, OVWeightQuantizationConfig
 
     try:
         if not (is_openvino_available() and is_diffusers_available()):

optimum/intel/openvino/__init__.py

@@ -32,6 +32,7 @@
     from .quantization import OVQuantizer
     from .trainer import OVTrainer
     from .training_args import OVTrainingArguments
+    from .weight_quantization import OVWeightQuantizationConfig
 
 from .modeling import (
     OVModelForAudioClassification,

optimum/intel/openvino/configuration.py

@@ -15,9 +15,12 @@
 from typing import Dict, List, Optional, Union
 
 import torch
+from transformers.utils.quantization_config import QuantizationConfigMixin
 
 from optimum.configuration_utils import BaseConfig
 
+from .weight_quantization import OVWeightQuantizationConfig
+
 
 DEFAULT_QUANTIZATION_CONFIG = {
     "algorithm": "quantization",
@@ -83,6 +86,7 @@ def __init__(
         compression: Union[List[Dict], Dict, None] = None,
         input_info: Optional[List] = None,
         save_onnx_model: bool = False,
+        quantization_config: Optional[QuantizationConfigMixin] = None,
         **kwargs,
     ):
         super().__init__()
@@ -91,6 +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.quantization_config = quantization_config
 
     def add_input_info(self, model_inputs: Dict, force_batch_one: bool = False):
         self.input_info = [
@@ -102,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

@@ -164,6 +164,7 @@ def _from_pretrained(
         from_onnx: bool = False,
         local_files_only: bool = False,
         load_in_8bit: bool = False,
+        load_in_4bit: bool = False,
         **kwargs,
     ):
         """
@@ -186,13 +187,18 @@ def _from_pretrained(
             force_download (`bool`, defaults to `False`):
                 Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                 cached versions if they exist.
-            file_name(`str`, *optional*):
+            file_name (`str`, *optional*):
                 The file name of the model to load. Overwrites the default file name and allows one to load the model
                 with a different name.
-            local_files_only(`bool`, *optional*, defaults to `False`):
+            local_files_only (`bool`, *optional*, defaults to `False`):
                 Whether or not to only look at local files (i.e., do not try to download the model).
+            load_in_8bit (`bool`, *optional*, defaults to `False`):
+                Whether or not to apply 8-bit weight quantization.
+            load_in_4bit (`bool`, *optional*, defaults to `False`):
+                Whether or not to apply 4-bit weight quantization.
         """
-
+        if load_in_4bit:
+            raise ValueError("load_in_4bit is available for OVModelForCausalLM only.")
         model_path = Path(model_id)
         default_file_name = ONNX_WEIGHTS_NAME if from_onnx else OV_XML_FILE_NAME
         file_name = file_name or default_file_name
@@ -260,6 +266,7 @@ def _from_transformers(
         task: Optional[str] = None,
         trust_remote_code: bool = False,
         load_in_8bit: Optional[bool] = None,
+        load_in_4bit: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -283,9 +290,10 @@ 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 = "int8" if load_in_8bit else "fp32"
+            compression_option = "fp32"
 
         main_export(
             model_name_or_path=model_id,
@@ -302,7 +310,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_base_seq2seq.py

@@ -253,7 +253,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,
             output=save_dir_path,
@@ -270,7 +270,7 @@ def _from_transformers(
 
         config.save_pretrained(save_dir_path)
         return cls._from_pretrained(
-            model_id=save_dir_path, config=config, use_cache=use_cache, load_in_8bit=False, **kwargs
+            model_id=save_dir_path, config=config, use_cache=use_cache, load_in_8bit=load_in_8bit, **kwargs
         )
 
     def _reshape(self, model: openvino.runtime.Model, batch_size: int, sequence_length: int, is_decoder=True):

optimum/intel/openvino/modeling_decoder.py

@@ -35,6 +35,7 @@
 from ..utils.modeling_utils import MULTI_QUERY_ATTN_MODELS
 from .modeling import _TOKENIZER_FOR_DOC, INPUTS_DOCSTRING, MODEL_START_DOCSTRING, OVModel
 from .utils import ONNX_WEIGHTS_NAME, OV_XML_FILE_NAME, STR_TO_OV_TYPE
+from .weight_quantization import OVWeightQuantizationConfig, compress_decoder_weights
 
 
 if is_transformers_version("<", "4.25.0"):
@@ -243,6 +244,8 @@ def _from_transformers(
         use_cache: bool = True,
         trust_remote_code: bool = False,
         load_in_8bit: Optional[bool] = None,
+        load_in_4bit: Optional[bool] = None,
+        quantization_config: Optional[Union[OVWeightQuantizationConfig, Dict]] = None,
         **kwargs,
     ):
         if config.model_type.replace("_", "-") not in _SUPPORTED_ARCHITECTURES:
@@ -259,9 +262,10 @@ 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:
-            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,
@@ -282,7 +286,14 @@ def _from_transformers(
         config.is_encoder_decoder = False
         config.save_pretrained(save_dir_path)
         return cls._from_pretrained(
-            model_id=save_dir_path, config=config, use_cache=use_cache, load_in_8bit=False, stateful=None, **kwargs
+            model_id=save_dir_path,
+            config=config,
+            use_cache=use_cache,
+            load_in_8bit=load_in_8bit,
+            stateful=None,
+            load_in_4bit=load_in_4bit,
+            quantization_config=quantization_config,
+            **kwargs,
         )
 
     def _reshape(
@@ -356,15 +367,14 @@ class OVModelForCausalLM(OVBaseDecoderModel, GenerationMixin):
             checkpoint="gpt2",
         )
     )
-    def forward(
+    def prepare_inputs(
         self,
         input_ids: torch.LongTensor,
         attention_mask: Optional[torch.LongTensor] = None,
         past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
         position_ids: Optional[torch.LongTensor] = None,
         **kwargs,
-    ) -> CausalLMOutputWithPast:
-        self.compile()
+    ) -> Dict:
         if self.use_cache and past_key_values is not None:
             input_ids = input_ids[:, -1:]
 
@@ -449,6 +459,26 @@ def forward(
                 self.next_beam_idx if self.next_beam_idx is not None else np.arange(batch_size, dtype=int)
             )
 
+        return inputs
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        self.compile()
+
+        inputs = self.prepare_inputs(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+            **kwargs,
+        )
+
         # Run inference
         self.request.start_async(inputs, share_inputs=True)
         self.request.wait()
@@ -532,6 +562,8 @@ def _from_pretrained(
         from_onnx: bool = False,
         local_files_only: bool = False,
         load_in_8bit: bool = False,
+        load_in_4bit: bool = False,
+        quantization_config: Union[OVWeightQuantizationConfig, Dict] = None,
         **kwargs,
     ):
         model_path = Path(model_id)
@@ -549,7 +581,9 @@ def _from_pretrained(
             local_files_only=local_files_only,
         )
 
-        model = cls.load_model(model_cache_path, load_in_8bit=load_in_8bit)
+        if load_in_8bit and load_in_4bit:
+            raise ValueError("Either load_in_8bit or load_in_4bit should be set to True.")
+        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":
@@ -563,7 +597,11 @@ def _from_pretrained(
         else:
             init_cls = cls
 
-        return init_cls(model=model, config=config, model_save_dir=model_cache_path.parent, **kwargs)
+        causal_model = init_cls(model=model, config=config, model_save_dir=model_cache_path.parent, **kwargs)
+
+        if load_in_4bit:
+            compress_decoder_weights(causal_model, quantization_config)
+        return causal_model
 
 
 class OVBloomForCausalLM(OVModelForCausalLM):