Fix SQ model restore loading (#1406)

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

Author: changwangss

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

@@ -401,6 +401,22 @@
         use_neural_speed=False,
     )
 
+# save model
+if args.output_dir is not None:
+    tokenizer.save_pretrained(args.output_dir)
+    if args.sq:
+        config.save_pretrained(args.output_dir)
+        user_model.save(args.output_dir)
+    elif args.mixed_precision or args.woq:
+        # user_model will be changed.
+        user_model.save_pretrained(args.output_dir)
+        # loading saved woq model
+        user_model = AutoModelForCausalLM.from_pretrained(
+            args.output_dir, 
+            trust_remote_code=args.trust_remote_code,
+            use_neural_speed=args.use_neural_speed
+            )
+
 if args.int8 or args.int8_bf16_mixed:
     # TorchScript model don't attribute generate method, the wrapper is provided.
     import intel_extension_for_pytorch as ipex
@@ -528,12 +544,3 @@
         args=args,
     )
     print(results)
-
-# save model
-if args.output_dir is not None:
-    tokenizer.save_pretrained(args.output_dir)
-    if args.sq:
-        config.save_pretrained(args.output_dir)
-        user_model.save(args.output_dir)
-    elif args.mixed_precision or args.woq:
-        user_model.save_pretrained(args.output_dir)

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

@@ -425,6 +425,21 @@
             _commit_hash=args._commit_hash,
         )
 
+# save model
+if args.output_dir is not None:
+    tokenizer.save_pretrained(args.output_dir)
+    if args.sq:
+        config.save_pretrained(args.output_dir)
+        user_model.save(args.output_dir)
+    elif args.mixed_precision or args.woq:
+        # user_model will be changed.
+        user_model.save_pretrained(args.output_dir)
+        # loading saved woq model
+        user_model = AutoModelForCausalLM.from_pretrained(
+            args.output_dir, 
+            trust_remote_code=args.trust_remote_code,
+            use_neural_speed=args.use_neural_speed
+            )
 
 
 # int8 model loading
@@ -552,11 +567,3 @@
                 % (task_name, results["results"][task_name]["acc"])
             )
 
-# save model
-if args.output_dir is not None:
-    tokenizer.save_pretrained(args.output_dir)
-    if args.sq:
-        config.save_pretrained(args.output_dir)
-        user_model.save(args.output_dir)
-    elif args.mixed_precision or args.woq:
-        user_model.save_pretrained(args.output_dir)

intel_extension_for_transformers/transformers/utils/utility.py

@@ -277,7 +277,8 @@ def get_example_inputs(model_config, batch_size=1, tokenizer=None, num_beams=4):
     prompt = "Welcome to use Intel Extension for Transformers."
     prompt = [prompt] * batch_size
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    if model_config.model_type in IPEX_OPT_LLM_SUPPORTED:
+    model_type = model_config.model_type.replace("_", "-")
+    if model_type in IPEX_OPT_LLM_SUPPORTED:
         past_key_values = generate_dummy_past_key_values_for_opt_llm(
                                                                     config=model_config,
                                                                     input_bs=batch_size,
@@ -287,10 +288,14 @@ def get_example_inputs(model_config, batch_size=1, tokenizer=None, num_beams=4):
         past_key_values = generate_dummy_past_key_values(config=model_config, input_bs=batch_size)
 
     input_ids = input_ids[:, :512]
-    attention_mask = torch.ones(input_ids.shape)
+    if model_type in ["bloom", "qwen"]:
+        attention_mask = torch.ones(input_ids.shape[0], input_ids.shape[1] + 1)
+        attention_mask[:,0] = 0
+    else:
+        attention_mask = torch.ones(input_ids.shape)
     position_ids = torch.arange(input_ids.shape[1]).repeat(batch_size, 1)
 
-    if model_config.model_type in MODEL_TYPES_REQUIRING_POSITION_IDS:
+    if model_type in MODEL_TYPES_REQUIRING_POSITION_IDS:
         example_inputs = {
                     "input_ids": input_ids,
                     "attention_mask": attention_mask,
@@ -342,5 +347,6 @@ def recover_model_from_json(user_model, json_file_path, trust_remote_code=False)
 
     # pylint: disable=E0611
     from neural_compressor.utils.pytorch import recover_model_from_json as inc_recover_model_from_json
+    user_model.config.torchscript = True
     user_model = inc_recover_model_from_json(user_model, json_file_path, example_inputs)
     return user_model