Update code with comments

Signed-off-by: Cheng, Penghui <penghui.cheng@intel.com>

Author: PenghuiCheng

examples/neural_compressor/language-modeling/requirements.txt

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

examples/neural_compressor/language-modeling/run_clm.py

@@ -33,7 +33,6 @@
 import torch
 import transformers
 from datasets import load_dataset
-from intel_extension_for_transformers.transformers.utils.config import WeightOnlyQuantConfig
 from neural_compressor import (
     DistillationConfig,
     PostTrainingQuantConfig,
@@ -58,7 +57,10 @@
 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"] = ""
 
@@ -626,6 +628,11 @@ def compute_metrics(eval_preds):
             else:
                 recipes = {}
             if optim_args.quantization_approach == "weight_only":
+                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(

optimum/intel/neural_compressor/quantization.py

@@ -19,7 +19,7 @@
 from enum import Enum
 from itertools import chain
 from pathlib import Path
-from typing import Callable, Dict, Optional, Union
+from typing import Callable, Dict, Optional, TypeAlias, Union
 
 import torch
 from datasets import Dataset, load_dataset
@@ -80,6 +80,9 @@
 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: TypeAlias = Union[PostTrainingQuantConfig, WeightOnlyQuantConfig]
+else:
+    Config: TypeAlias = PostTrainingQuantConfig
 
 logger = logging.getLogger(__name__)
 
@@ -149,7 +152,7 @@ def from_pretrained(cls, model: PreTrainedModel, **kwargs):
     def quantize(
         self,
         save_directory: Union[str, Path],
-        quantization_config=None,
+        quantization_config: Config = None,
         calibration_dataset: Dataset = None,
         batch_size: int = 8,
         data_collator: Optional[DataCollator] = None,
@@ -162,7 +165,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.
@@ -261,8 +264,7 @@ def quantize(
                 save_onnx_model = False
 
         if (
-            not weight_only
-            and not isinstance(quantization_config, WeightOnlyQuantConfig)
+            isinstance(quantization_config, PostTrainingQuantConfig)
             and quantization_config.backend == "ipex"
             and is_ipex_version("<", IPEX_MINIMUM_VERSION)
             and "generation" in self.task
@@ -272,7 +274,7 @@ def quantize(
                 f"but only version {IPEX_MINIMUM_VERSION} or higher is supported."
             )
 
-        if isinstance(quantization_config, WeightOnlyQuantConfig):
+        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)

optimum/intel/utils/import_utils.py

@@ -350,7 +350,7 @@ def is_timm_version(operation: str, version: str):
 
 INTEL_EXTENSION_FOR_TRANSFORMERS_IMPORT_ERROR = """
 {0} requires the intel-extension-for-transformers library but it was not found in your environment. You can install it with pip:
-`pip install neural-compressor`. Please note that you may need to restart your runtime after installation.
+`pip install intel-extension-for-transformers`. Please note that you may need to restart your runtime after installation.
 """
 
 DATASETS_IMPORT_ERROR = """

tests/neural_compressor/test_optimization.py

@@ -202,14 +202,43 @@ def test_ipex_static_quantization_with_smoothquant(self, task, model_name, expec
 
     def test_weight_only_quantization(self):
         model_name = "hf-internal-testing/tiny-random-GPTNeoForCausalLM"
-        quantization_config = WeightOnlyQuantConfig(weight_dtype="int8")
         model = AutoModelForCausalLM.from_pretrained(model_name)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         tokenizer.add_special_tokens({"pad_token": "[PAD]"})
         quantizer = INCQuantizer.from_pretrained(copy.deepcopy(model), task="text-generation")
         calibration_dataset = _generate_dataset(quantizer, tokenizer, num_samples=2)
 
         with tempfile.TemporaryDirectory() as tmp_dir:
+            quantization_config = WeightOnlyQuantConfig(weight_dtype="int8")
+            q_model = quantizer.quantize(
+                quantization_config=quantization_config,
+                save_directory=tmp_dir,
+            )
+            inp = torch.tensor([calibration_dataset[0]["input_ids"]])
+            out = model(inp)[0]
+            q_out = q_model(inp)[0]
+            self.assertTrue(torch.all(torch.isclose(out, q_out, atol=5e-1)))
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            quantization_config = WeightOnlyQuantConfig(
+                algorithm="GPTQ",
+                weight_dtype="int4_clip",
+            )
+            q_model = quantizer.quantize(
+                quantization_config=quantization_config,
+                calibration_dataset=calibration_dataset,
+                save_directory=tmp_dir,
+            )
+            inp = torch.tensor([calibration_dataset[0]["input_ids"]])
+            out = model(inp)[0]
+            q_out = q_model(inp)[0]
+            self.assertTrue(torch.all(torch.isclose(out, q_out, atol=5e-1)))
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            quantization_config = WeightOnlyQuantConfig(
+                algorithm="AWQ",
+                weight_dtype="int4_clip",
+            )
             q_model = quantizer.quantize(
                 quantization_config=quantization_config,
                 calibration_dataset=calibration_dataset,
@@ -220,6 +249,16 @@ def test_weight_only_quantization(self):
             q_out = q_model(inp)[0]
             self.assertTrue(torch.all(torch.isclose(out, q_out, atol=5e-1)))
 
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            q_model = quantizer.quantize(
+                weight_only=True,  # use RTN quantization method and NF4 weight data type is default.
+                save_directory=tmp_dir,
+            )
+            inp = torch.tensor([calibration_dataset[0]["input_ids"]])
+            out = model(inp)[0]
+            q_out = q_model(inp)[0]
+            self.assertTrue(torch.all(torch.isclose(out, q_out, atol=5e-1)))
+
     def test_dynamic_accuracy_strategy_quantization(self):
         model_name = "distilbert-base-cased-distilled-squad"
         model = AutoModelForQuestionAnswering.from_pretrained(model_name)