Support weight-only quantization with quantized operators in intel-extension-for-transformers.

examples/neural_compressor/language-modeling/README.md

@@ -97,4 +97,4 @@ respectively `dynamic`, `static`, `weight_only` or `aware_training`.
 
 The flag `--verify_loading` can be passed along to verify that the resulting quantized model can be loaded correctly.
 
-> **_Note:_** `weight_only` quantization_approach requires neural-compressor >= 2.3
+> **_Note:_** `weight_only` quantization_approach requires neural-compressor >= 2.3 and intel-extension-for-transformers >= 1.3.

examples/neural_compressor/language-modeling/requirements.txt

@@ -3,3 +3,5 @@ torch >= 1.9
 datasets >= 1.8.0
 sentencepiece != 0.1.92
 protobuf
+intel-extension-for-transformers >= 1.3
+peft

examples/neural_compressor/language-modeling/run_clm.py

@@ -57,8 +57,12 @@
 from transformers.utils.versions import require_version
 
 from optimum.intel.neural_compressor import INCModelForCausalLM, INCQuantizer, INCTrainer
+from optimum.intel.utils.import_utils import is_intel_extension_for_transformers_available
 
 
+if is_intel_extension_for_transformers_available():
+    from intel_extension_for_transformers.transformers.utils.config import WeightOnlyQuantConfig
+
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
 # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
@@ -196,9 +200,9 @@ class OptimizationArguments:
         default=False,
         metadata={"help": "Whether or not to verify the loading of the quantized model."},
     )
-    bits: int = field(
-        default=8,
-        metadata={"help": "Bits for weight only quantization, 1-8 bits."},
+    weight_dtype: str = field(
+        default="int8",
+        metadata={"help": "weight dtype for weight only quantization."},
     )
     group_size: int = field(
         default=-1,
@@ -625,26 +629,23 @@ def compute_metrics(eval_preds):
             else:
                 recipes = {}
             if optim_args.quantization_approach == "weight_only":
-                op_type_dict = {
-                    ".*": {
-                        "weight": {
-                            "bits": optim_args.bits,
-                            "group_size": optim_args.group_size,
-                            "scheme": optim_args.weight_only_scheme,
-                            "algorithm": optim_args.quantization_methodology,
-                        },
-                    },
-                }
-                if optim_args.quantization_methodology == "GPTQ":
-                    gptq_args = {
-                        "pad_max_length": block_size,
-                    }
-                    recipes.update({"gptq_args": gptq_args})
+                if not is_intel_extension_for_transformers_available():
+                    raise ImportError(
+                        "Didn't find out intel-etension-for-transformers package. "
+                        "Please install packages: pip install intel-etension-for-transformers and pip install peft."
+                    )
+                if optim_args.apply_pruning or optim_args.apply_distillation:
+                    raise ValueError("Weight only quantization and pruning or distillation cannot be combined.")
+                quantization_config = WeightOnlyQuantConfig(
+                    weight_dtype=optim_args.weight_dtype,
+                    group_size=optim_args.group_size,
+                    scheme=optim_args.weight_only_scheme,
+                    algorithm=optim_args.quantization_methodology,
+                )
             else:
-                op_type_dict = {}
-            quantization_config = PostTrainingQuantConfig(
-                approach=optim_args.quantization_approach, op_type_dict=op_type_dict, recipes=recipes
-            )
+                quantization_config = PostTrainingQuantConfig(
+                    approach=optim_args.quantization_approach, recipes=recipes
+                )
 
     if optim_args.apply_pruning:
         if optim_args.end_step is None:
@@ -735,12 +736,12 @@ def compute_metrics(eval_preds):
             calibration_dataset=train_dataset
             if optim_args.quantization_approach in ["static", "weight_only"]
             else None,
-            batch_size=1  # batch_size > 1 for GPTQ is WIP
-            if optim_args.quantization_approach == "weight_only" and optim_args.quantization_methodology == "GPTQ"
+            batch_size=1
+            if optim_args.quantization_approach == "weight_only"
             else training_args.per_device_train_batch_size,
-            weight_only=True if optim_args.quantization_approach == "weight_only" else False,
         )
         trainer.model = quantizer._quantized_model
+
     if optim_args.apply_quantization and optim_args.verify_loading:
         loaded_model = INCModelForCausalLM.from_pretrained(training_args.output_dir)
         tokens = tokenizer("This is a sample input", return_tensors="pt")

optimum/intel/neural_compressor/configuration.py

@@ -35,7 +35,7 @@ class INCConfig(BaseConfig):
 
     def __init__(
         self,
-        quantization: Optional[Union[Dict, _BaseQuantizationConfig]] = None,
+        quantization: Optional[Union[Dict, _BaseQuantizationConfig, "WeightOnlyQuantConfig"]] = None,
         pruning: Optional[Union[Dict, _BaseQuantizationConfig]] = None,
         distillation: Optional[Union[Dict, _BaseQuantizationConfig]] = None,
         save_onnx_model: bool = False,
@@ -50,7 +50,7 @@ def __init__(
         self.save_onnx_model = save_onnx_model
 
     @staticmethod
-    def _create_quantization_config(config: Union[Dict, _BaseQuantizationConfig]):
+    def _create_quantization_config(config):
         # TODO : add activations_dtype and weights_dtype
         if isinstance(config, _BaseQuantizationConfig):
             approach = _quantization_model[config.approach]

optimum/intel/neural_compressor/quantization.py

@@ -58,6 +58,7 @@
 from ..utils.import_utils import (
     _ipex_version,
     _neural_compressor_version,
+    is_intel_extension_for_transformers_available,
     is_ipex_version,
     is_neural_compressor_version,
 )
@@ -76,6 +77,14 @@
 from .utils import INCDataLoader, _cfgs_to_fx_cfgs
 
 
+if is_intel_extension_for_transformers_available():
+    from intel_extension_for_transformers.llm.quantization.utils import convert_to_quantized_model
+    from intel_extension_for_transformers.transformers.utils.config import WeightOnlyQuantConfig
+
+    Config = Union[PostTrainingQuantConfig, WeightOnlyQuantConfig]
+else:
+    Config = PostTrainingQuantConfig
+
 logger = logging.getLogger(__name__)
 
 NEURAL_COMPRESSOR_MINIMUM_VERSION = "2.1.0"
@@ -143,8 +152,8 @@ def from_pretrained(cls, model: PreTrainedModel, **kwargs):
 
     def quantize(
         self,
-        quantization_config: "PostTrainingQuantConfig",
         save_directory: Union[str, Path],
+        quantization_config: Config = None,
         calibration_dataset: Dataset = None,
         batch_size: int = 8,
         data_collator: Optional[DataCollator] = None,
@@ -157,7 +166,7 @@ def quantize(
         Quantize a model given the optimization specifications defined in `quantization_config`.
 
         Args:
-            quantization_config (`PostTrainingQuantConfig`):
+            quantization_config (`Union[PostTrainingQuantConfig, WeightOnlyQuantConfig]`):
                 The configuration containing the parameters related to quantization.
             save_directory (`Union[str, Path]`):
                 The directory where the quantized model should be saved.
@@ -177,30 +186,36 @@ def quantize(
         save_directory.mkdir(parents=True, exist_ok=True)
         save_onnx_model = kwargs.pop("save_onnx_model", False)
 
-        if save_onnx_model and isinstance(self._original_model, ORTModel):
+        if save_onnx_model and (isinstance(self._original_model, ORTModel) or weight_only):
             save_onnx_model = False
             logger.warning("Model provided is an ONNX model, `save_onnx_model` is set to False")
 
         default_name = WEIGHTS_NAME if not isinstance(self._original_model, ORTModel) else ONNX_WEIGHTS_NAME
         calibration_dataloader = None
         self._set_task()
 
-        if weight_only:
+        if weight_only or not isinstance(quantization_config, PostTrainingQuantConfig):
             # check neural-compressor version
             if is_neural_compressor_version("<", NEURAL_COMPRESSOR_WEIGHT_ONLY_MINIMUM_VERSION):
                 raise ImportError(
                     f"Found an incompatible version of neural-compressor. Found version {_neural_compressor_version}, "
                     f"but only version {NEURAL_COMPRESSOR_WEIGHT_ONLY_MINIMUM_VERSION} or higher supports weight-only quantization."
                 )
+            if not is_intel_extension_for_transformers_available():
+                raise ImportError(
+                    "Didn't find out intel-etension-for-transformers package. "
+                    "Please install packages: pip install intel-etension-for-transformers and pip install peft."
+                )
 
-            # If op_type_dict of quantization_config is not defined, it will use default values for weight-only quantization:
-            # {"bits": 4, "group_size": 32, "scheme": "sym", "algorithm": "RTN"}
-            if isinstance(quantization_config.op_type_dict, dict) and len(quantization_config.op_type_dict) > 0:
-                algo = []
-                for _, val in quantization_config.op_type_dict.items():
-                    algo += val.get("weight", {}).get("algorithm", ["RTN"])
-            else:
+            if quantization_config is None:
+                quantization_config = WeightOnlyQuantConfig()
                 algo = ["RTN"]
+            elif isinstance(quantization_config, WeightOnlyQuantConfig):
+                algo = quantization_config.algorithm
+            else:
+                raise TypeError(
+                    f"For weight-only quantization, `quantization_config` should be an instance of `WeightOnlyQuantConfig`, but got: {type(quantization_config)} instead."
+                )
 
             if calibration_dataset is None and ("GPTQ" in algo or "AWQ" in algo):
                 raise ValueError(
@@ -217,6 +232,9 @@ def quantize(
                     data_collator=data_collator,
                     use_label=False if "GPTQ" in algo else True,
                 )
+            quantization_config.calib_dataloader = calibration_dataloader
+
+            save_onnx_model = False
 
         elif INCQuantizationMode(quantization_config.approach) == INCQuantizationMode.STATIC:
             # Since PyTorch fx trace does not really require an example_inputs, only need calibration_dataset or calibration_fn here.
@@ -249,7 +267,8 @@ def quantize(
                 save_onnx_model = False
 
         if (
-            quantization_config.backend == "ipex"
+            isinstance(quantization_config, PostTrainingQuantConfig)
+            and quantization_config.backend == "ipex"
             and is_ipex_version("<", IPEX_MINIMUM_VERSION)
             and "generation" in self.task
         ):
@@ -258,76 +277,83 @@ def quantize(
                 f"but only version {IPEX_MINIMUM_VERSION} or higher is supported."
             )
 
-        if isinstance(self._original_model.config, PretrainedConfig):
-            self._original_model.config.backend = quantization_config.backend
-
-        if isinstance(self._original_model, ORTModel):
-            # TODO : enable seq2seq models
-            if isinstance(self._original_model, ORTModelForConditionalGeneration):
-                raise RuntimeError("ORTModelForConditionalGeneration not supported for quantization")
-
-            if isinstance(self._original_model, ORTModelForCausalLM):
-                model_or_path = self._original_model.onnx_paths
-                if len(model_or_path) > 1:
-                    raise RuntimeError(
-                        f"Too many ONNX model files were found in {self._original_model.onnx_paths}, only `use_cache=False` is supported"
-                    )
-                model_or_path = str(model_or_path[0])
-                default_name = ONNX_DECODER_NAME
-            else:
-                model_or_path = str(self._original_model.model_path)
+        if not isinstance(quantization_config, PostTrainingQuantConfig):
+            self._quantized_model = convert_to_quantized_model(self._original_model, quantization_config)
+            # Save the quantized model
+            output_path = save_directory.joinpath(file_name or default_name)
+            self._quantized_model.save_pretrained(output_path)
         else:
-            model_or_path = self._original_model
-
-        compressed_model = fit(
-            model_or_path,
-            conf=quantization_config,
-            calib_dataloader=calibration_dataloader,
-            eval_func=self.eval_fn,
-            calib_func=self.calibration_fn,
-        )
-
-        if not hasattr(compressed_model, "_model") or compressed_model._model is None:
-            raise RuntimeError(
-                "The maximum number of trials specified has been reached and no quantized model meeting the specified"
-                " accuracy tolerance has been found. Either the tolerance or the number of trials need to be increased."
+            if isinstance(self._original_model.config, PretrainedConfig):
+                self._original_model.config.backend = quantization_config.backend
+
+            if isinstance(self._original_model, ORTModel):
+                # TODO : enable seq2seq models
+                if isinstance(self._original_model, ORTModelForConditionalGeneration):
+                    raise RuntimeError("ORTModelForConditionalGeneration not supported for quantization")
+
+                if isinstance(self._original_model, ORTModelForCausalLM):
+                    model_or_path = self._original_model.onnx_paths
+                    if len(model_or_path) > 1:
+                        raise RuntimeError(
+                            f"Too many ONNX model files were found in {self._original_model.onnx_paths}, only `use_cache=False` is supported"
+                        )
+                    model_or_path = str(model_or_path[0])
+                    default_name = ONNX_DECODER_NAME
+                else:
+                    model_or_path = str(self._original_model.model_path)
+            else:
+                model_or_path = self._original_model
+
+            compressed_model = fit(
+                model_or_path,
+                conf=quantization_config,
+                calib_dataloader=calibration_dataloader,
+                eval_func=self.eval_fn,
+                calib_func=self.calibration_fn,
             )
 
-        if isinstance(self._original_model.config, PretrainedConfig):
-            # If backend is IPEX, then the quantized model is JIT model which will drop the config attribute,
-            # so need set config from original_model.
-            model_config = copy.deepcopy(self._original_model.config)
-            model_config.torch_dtype = "int8"
-            if isinstance(compressed_model, IPEXModel):
-                model_config.torchscript = True
-                model_config.backend = "ipex"
-            elif not isinstance(compressed_model, ONNXModel):
-                compressed_model._model.config = model_config
-            model_config.save_pretrained(save_directory)
-
-        self._quantized_model = compressed_model._model
-
-        if save_onnx_model:
-            model_type = self._original_model.config.model_type.replace("_", "-")
-            model_name = getattr(self._original_model, "name", None)
-            onnx_config_class = TasksManager.get_exporter_config_constructor(
-                exporter="onnx",
-                model=self._original_model,
-                task=self.task,
-                model_type=model_type,
-                model_name=model_name,
-            )
-            onnx_config = onnx_config_class(self._original_model.config)
-            compressed_model.eval()
-            output_onnx_path = save_directory.joinpath(ONNX_WEIGHTS_NAME)
-            # Export the compressed model to the ONNX format
-            self._onnx_export(compressed_model, onnx_config, output_onnx_path)
-
-        output_path = save_directory.joinpath(file_name or default_name)
-        # Save the quantized model
-        self._save_pretrained(compressed_model, output_path)
-        quantization_config = INCConfig(quantization=quantization_config, save_onnx_model=save_onnx_model)
-        quantization_config.save_pretrained(save_directory)
+            if not hasattr(compressed_model, "_model") or compressed_model._model is None:
+                raise RuntimeError(
+                    "The maximum number of trials specified has been reached and no quantized model meeting the specified"
+                    " accuracy tolerance has been found. Either the tolerance or the number of trials need to be increased."
+                )
+
+            if isinstance(self._original_model.config, PretrainedConfig):
+                # If backend is IPEX, then the quantized model is JIT model which will drop the config attribute,
+                # so need set config from original_model.
+                model_config = copy.deepcopy(self._original_model.config)
+                model_config.torch_dtype = "int8"
+                if isinstance(compressed_model, IPEXModel):
+                    model_config.torchscript = True
+                    model_config.backend = "ipex"
+                elif not isinstance(compressed_model, ONNXModel):
+                    compressed_model._model.config = model_config
+                model_config.save_pretrained(save_directory)
+
+            self._quantized_model = compressed_model._model
+
+            if save_onnx_model:
+                model_type = self._original_model.config.model_type.replace("_", "-")
+                model_name = getattr(self._original_model, "name", None)
+                onnx_config_class = TasksManager.get_exporter_config_constructor(
+                    exporter="onnx",
+                    model=self._original_model,
+                    task=self.task,
+                    model_type=model_type,
+                    model_name=model_name,
+                )
+                onnx_config = onnx_config_class(self._original_model.config)
+                compressed_model.eval()
+                output_onnx_path = save_directory.joinpath(ONNX_WEIGHTS_NAME)
+                # Export the compressed model to the ONNX format
+                self._onnx_export(compressed_model, onnx_config, output_onnx_path)
+
+            output_path = save_directory.joinpath(file_name or default_name)
+            # Save the quantized model
+            self._save_pretrained(compressed_model, output_path)
+            quantization_config = INCConfig(quantization=quantization_config, save_onnx_model=save_onnx_model)
+            quantization_config.save_pretrained(save_directory)
+        return self._quantized_model
 
     @staticmethod
     def _save_pretrained(model: Union[PyTorchModel, IPEXModel], output_path: str):