Improve WOQ algo autoround (#1330)

Signed-off-by: changwangss <chang1.wang@intel.com>

examples/huggingface/pytorch/code-generation/quantization/requirements.txt

@@ -12,7 +12,7 @@ neural-compressor
 intel_extension_for_pytorch==2.2.0
 optimum-intel
 git+https://github.com/bigcode-project/bigcode-evaluation-harness@00967d12093ef614de7bdad0772aed8e4118f1fd
-git+https://github.com/intel/auto-round.git@a868c805de4be271cfe7403309a64d9bf03a0ecf
+git+https://github.com/intel/auto-round.git@b65830f3f6cb32d92a5c8ba5f80ace12d517357b
 
 
 

examples/huggingface/pytorch/code-generation/quantization/run_generation.py

@@ -136,7 +136,7 @@
 # ============AUTOROUND configs==============
 parser.add_argument(
     "--autoround_nsamples",
-    type=int, default=128,
+    type=int, default=512,
     help="Number of calibration data samples.",
 )
 parser.add_argument(
@@ -301,6 +301,7 @@
             "iters": args.calib_iters,
             "scale_dtype": "fp32",
             "device": "cpu",
+            "export_args": {"format": "itrex", "inplace": False}
         }
         quantization_config = WeightOnlyQuantConfig(
             compute_dtype=args.woq_compute_dtype,

examples/huggingface/pytorch/text-generation/quantization/requirements.txt

@@ -14,4 +14,4 @@ tiktoken  #qwen
 einops  #qwen
 git+https://github.com/intel/neural-compressor.git
 git+https://github.com/EleutherAI/lm-evaluation-harness.git@cc9778fbe4fa1a709be2abed9deb6180fd40e7e2
-git+https://github.com/intel/auto-round.git@a868c805de4be271cfe7403309a64d9bf03a0ecf
+git+https://github.com/intel/auto-round.git@b65830f3f6cb32d92a5c8ba5f80ace12d517357b

examples/huggingface/pytorch/text-generation/quantization/run_generation.py

@@ -162,7 +162,7 @@
 # ============AUTOROUND configs==============
 parser.add_argument(
     "--autoround_nsamples",
-    type=int, default=128,
+    type=int, default=512,
     help="Number of calibration data samples.",
 )
 parser.add_argument(
@@ -312,6 +312,7 @@
             "iters": args.calib_iters,
             "scale_dtype": "fp32",
             "device": "cpu",
+            "export_args": {"format": "itrex", "inplace": False}
         }
         quantization_config = WeightOnlyQuantConfig(
             compute_dtype=args.woq_compute_dtype,

intel_extension_for_transformers/llm/quantization/utils.py

@@ -24,6 +24,7 @@
 from datasets import load_dataset
 from neural_compressor import quantization
 from neural_compressor.adaptor.torch_utils.model_wrapper import WeightOnlyLinear
+from auto_round.export.export_to_itrex.model_wrapper import WeightOnlyLinear as auto_round_woqlinear
 from neural_compressor.utils.utility import LazyImport
 from neural_compressor.config import PostTrainingQuantConfig
 from ...utils.utils import is_ipex_available
@@ -105,8 +106,9 @@ def _replace_linear(
         is_removed = False
 
         if (isinstance(module, torch.nn.Linear) or isinstance(module, WeightOnlyLinear)
-            or (is_ipex_available() and isinstance(module, ipex.nn.utils._weight_prepack._IPEXLinear))) \
-           and (name not in modules_to_not_convert):
+            or isinstance(module, auto_round_woqlinear) or (is_ipex_available()
+            and isinstance(module, ipex.nn.utils._weight_prepack._IPEXLinear))) \
+            and (name not in modules_to_not_convert):
             # Check if the current key is not in the `modules_to_not_convert`
             if not any(
                 key in ".".join(current_key_name) for key in modules_to_not_convert
@@ -186,7 +188,7 @@ def _replace_linear(
                                 int_weight,
                                 gptq_scales,
                                 gptq_zeros,
-                                module.g_idx,
+                                module.g_idx if hasattr(module, "g_idx") else None,
                                 quantization_config,
                                 bias=None if module.bias is None else module.bias.data,
                             )
@@ -279,12 +281,35 @@ def collate_batch(batch):
                 input_ids_padded.append(input_ids)
             return torch.vstack(input_ids_padded)
 
-        calib_dataloader = DataLoader(
-            tokenized_dataset,
-            batch_size=1,
-            shuffle=False,
-            collate_fn=collate_batch,
-        )
+        def collate_batch_for_autoround(batch):
+            input_ids_padded = []
+            for text in batch:
+                input_ids = text["input_ids"]
+                if input_ids.shape[0] < config.algorithm_args["seq_len"]:
+                    continue
+                input_ids = input_ids[:config.algorithm_args["seq_len"]]
+                input_ids_list = input_ids.tolist()
+                if input_ids_list.count(input_ids_list[-1]) > config.algorithm_args["seq_len"] // 2:
+                    continue
+                input_ids_padded.append(input_ids)
+            if len(input_ids_padded) == 0:
+                return None
+
+            return torch.vstack(input_ids_padded)
+        if config.algorithm == "AUTOROUND":
+            calib_dataloader = DataLoader(
+                tokenized_dataset,
+                batch_size=1,
+                shuffle=False,
+                collate_fn=collate_batch_for_autoround,
+            )
+        else:
+            calib_dataloader = DataLoader(
+                tokenized_dataset,
+                batch_size=1,
+                shuffle=False,
+                collate_fn=collate_batch,
+            )
     if calib_func is None and config.algorithm in ["AWQ"]:
 
         def default_calib_func(model):
@@ -390,7 +415,6 @@ def default_calib_func(model):
                 setattr(config, "gptq_quantize_config", quantize_config)
                 q_model = replace_linear(inc_model, None, None, config, device=device)
             elif config.algorithm == "AUTOROUND":
-                inc_model = inc_model.export_compressed_model(use_optimum_format=True)
                 inc_model.eval()
                 quantize_config = {
                     "bits": bits,
@@ -403,7 +427,7 @@ def default_calib_func(model):
                 }
 
                 setattr(config, "gptq_quantize_config", quantize_config)
-                q_model = replace_linear(inc_model, None, None, config, device=device)
+                q_model = replace_linear(inc_model._model, None, None, config, device=device)
             else:
                 q_model = replace_linear(inc_model.model, None, None, config, device=device)
         if orig_dtype != torch.float32:

intel_extension_for_transformers/transformers/modeling/modeling_auto.py

@@ -408,7 +408,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
                 import intel_extension_for_pytorch as ipex
             except ImportError:
                 logger.warning("Please install Intel Extension for PyTorch to accelerate the model inference.")
-            assert (ipex.__version__ >= "2.1.0+cpu"), "Please use Intel Extension for PyTorch >=2.1.0+cpu."
+            assert (ipex.__version__ >= "2.2.0+cpu"), "Please use Intel Extension for PyTorch >=2.2.0+cpu."
             model = cls.ORIG_MODEL.from_pretrained(
                 pretrained_model_name_or_path,
                 low_cpu_mem_usage=True,
@@ -424,8 +424,6 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
                 model = model.float()
             model.eval()
             model_type = model.config.model_type.replace("_", "-")
-            if "falcon" in model_type:
-                logger.warning("Please use transformers 4.33.3 if you would like to apply smoothquant to Falcon.")
             if "llama" in model_type and transformers.__version__ >= "4.36.0":
                 quantization_config.ipex_opt_llm = False
             logger.info("Applying SmoothQuant.")
@@ -434,7 +432,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
                 if model_type in IPEX_OPT_LLM_SUPPORTED:
                     quantization_config.ipex_opt_llm = True
                     logger.info("quantization_config.ipex_opt_llm set to True and ipex.optimize_transformers is used.")
-                    logger.warning("The suggested transformers version is 4.31.0.")
+                    logger.warning("The suggested transformers version is 4.35.2.")
                 else:
                     quantization_config.ipex_opt_llm = False
             if quantization_config.ipex_opt_llm:
@@ -487,12 +485,12 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
                     calib_dataset = calib_dataset.shuffle(seed=42)
 
                 def tokenize_function(examples):
-                    if "prompt" in examples:
+                    if "code" in examples:
+                        example = tokenizer(examples["code"])
+                    elif "prompt" in examples:
                         example = tokenizer(examples["prompt"])
                     elif "text" in examples:
                         example = tokenizer(examples["text"])
-                    elif "code" in examples:
-                        example = tokenizer(examples["code"])
                     else:
                         logger.error("Please check dataset prompt identifier," +
                                      " NeelNanda/pile-10k is default used calibration dataset.")

requirements.txt

@@ -2,6 +2,7 @@
 accelerate
 cmake>=3.16
 gguf
+git+https://github.com/intel/auto-round.git@b65830f3f6cb32d92a5c8ba5f80ace12d517357b
 ninja
 optimum-intel
 py-cpuinfo

tests/CI/test_quantization.py

@@ -465,6 +465,7 @@ def test_quantization_for_llm(self):
             "iters": 5,
             "scale_dtype": "fp32",
             "device": "cpu",
+            "export_args": {"format": "itrex", "inplace": False}
         }
         woq_config = WeightOnlyQuantConfig(weight_dtype="int4_clip",
                                         algorithm_args=algorithm_args,
@@ -476,7 +477,7 @@ def test_quantization_for_llm(self):
                                                 )
         woq_model.eval()
         output = woq_model(dummy_input)
-        self.assertTrue(isclose(float(output[0][0][0][0]), 0.18015708029270172, rel_tol=1e-04))
+        self.assertTrue(isclose(float(output[0][0][0][0]), 0.173023983836174, rel_tol=1e-04))
 
     def test_export(self):
         # test model with model_id

tests/requirements.txt

@@ -5,7 +5,7 @@ datasets==2.14.7
 einops
 evaluate
 gguf
-git+https://github.com/intel/auto-round.git@a868c805de4be271cfe7403309a64d9bf03a0ecf
+git+https://github.com/intel/auto-round.git@b65830f3f6cb32d92a5c8ba5f80ace12d517357b
 git+https://github.com/intel/neural-compressor.git
 intel-extension-for-pytorch==2.2.0
 intel-tensorflow==2.14.0