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

Author: PenghuiCheng

examples/neural_compressor/language-modeling/run_clm.py

@@ -33,7 +33,7 @@
 import torch
 import transformers
 from datasets import load_dataset
-from intel_extension_for_transformers.transformers.utils.quantization_config import WeightOnlyQuantConfig
+from intel_extension_for_transformers.transformers.utils.config import WeightOnlyQuantConfig
 from neural_compressor import (
     DistillationConfig,
     PostTrainingQuantConfig,
@@ -735,12 +735,7 @@ def compute_metrics(eval_preds):
         )
         trainer.model = quantizer._quantized_model
 
-    # TODO: Weight only quantization didn't support save/load function now. Will implement it soon.
-    if (
-        optim_args.apply_quantization
-        and optim_args.verify_loading
-        and optim_args.quantization_approach != "weight_only"
-    ):
+    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")
         with torch.no_grad():

optimum/intel/neural_compressor/__init__.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from ..utils.import_utils import is_diffusers_available
-from .configuration import INCConfig
+from .configuration import INCConfig, WeightOnlyQuantConfig
 from .modeling_base import (
     INCModel,
     INCModelForMaskedLM,

optimum/intel/neural_compressor/configuration.py

@@ -14,6 +14,7 @@
 
 from typing import Dict, Optional, Union
 
+from intel_extension_for_transformers.transformers.utils import WeightOnlyQuantConfig
 from neural_compressor.config import DistillationConfig, WeightPruningConfig, _BaseQuantizationConfig
 
 from optimum.configuration_utils import BaseConfig
@@ -35,7 +36,7 @@ class INCConfig(BaseConfig):
 
     def __init__(
         self,
-        quantization=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 +51,7 @@ def __init__(
         self.save_onnx_model = save_onnx_model
 
     @staticmethod
-    def _create_quantization_config(config):
+    def _create_quantization_config(config: Union[Dict, _BaseQuantizationConfig, WeightOnlyQuantConfig]):
         # TODO : add activations_dtype and weights_dtype
         if isinstance(config, _BaseQuantizationConfig):
             approach = _quantization_model[config.approach]

optimum/intel/neural_compressor/quantization.py

@@ -59,7 +59,6 @@
     _ipex_version,
     _neural_compressor_version,
     is_intel_extension_for_transformers_available,
-    is_intel_extension_for_transformers_version,
     is_ipex_version,
     is_neural_compressor_version,
 )
@@ -80,10 +79,7 @@
 
 if is_intel_extension_for_transformers_available():
     from intel_extension_for_transformers.llm.quantization.utils import convert_to_quantized_model
-    if is_intel_extension_for_transformers_version("<=", "1.2.2"):
-        from intel_extension_for_transformers.transformers.utils.quantization_config import WeightOnlyQuantConfig
-    else:
-        from intel_extension_for_transformers.transformers.utils.config import WeightOnlyQuantConfig
+    from intel_extension_for_transformers.transformers.utils.config import WeightOnlyQuantConfig
 
 logger = logging.getLogger(__name__)
 
@@ -186,7 +182,7 @@ 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")
 
@@ -278,6 +274,9 @@ def quantize(
 
         if isinstance(quantization_config, WeightOnlyQuantConfig):
             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:
             if isinstance(self._original_model.config, PretrainedConfig):
                 self._original_model.config.backend = quantization_config.backend

tests/neural_compressor/test_optimization.py

@@ -56,7 +56,7 @@
     INCSeq2SeqTrainer,
     INCStableDiffusionPipeline,
 )
-from intel_extension_for_transformers.transformers.utils.quantization_config import WeightOnlyQuantConfig
+from intel_extension_for_transformers.transformers.utils.config import WeightOnlyQuantConfig
 from optimum.intel.utils.constant import DIFFUSION_WEIGHTS_NAME
 from optimum.onnxruntime import ORTModelForCausalLM, ORTModelForSequenceClassification
 from optimum.pipelines import ORT_SUPPORTED_TASKS