Fix ipex mistral export for bs > 1 (#1276)

* fix bug when bs > 1 and do not provide `position_ids` for input

Signed-off-by: Liu, Kaixuan <[email protected]>

* Update tests/ipex/test_modeling_causal_lm.py

Co-authored-by: Ella Charlaix <[email protected]>

* delete code for test case

Signed-off-by: Liu, Kaixuan <[email protected]>

---------

Signed-off-by: Liu, Kaixuan <[email protected]>
Co-authored-by: Ella Charlaix <[email protected]>

Author: kaixuanliu

optimum/exporters/ipex/modeling_utils.py

@@ -665,7 +665,10 @@ def _mistral_model_forward(
         )
 
     if position_ids is None:
-        position_ids = cache_position.unsqueeze(0)
+        position_ids = torch.arange(
+            past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+        )
+        position_ids = position_ids.unsqueeze(0).repeat_interleave(input_ids.shape[0], 0)
 
     causal_mask = self._update_causal_mask(
         attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
@@ -1021,7 +1024,7 @@ def __init__(self, module, device, config) -> None:
         self.module_device = device
 
         if not config.compile and getattr(config, "quantization_config", None) is None:
-            # LinearAllreduce and LinearLayer cannot use fused op LinearAdd
+            # LinearAllreduce cannot use fused op LinearAdd
             if module.down_proj.__class__.__name__ not in ["LinearAllreduce"]:
                 self.mlp_linear_add = LinearAdd(module.down_proj)
             if isinstance(self.act_fn, nn.SiLU):
@@ -1049,7 +1052,7 @@ def __init__(self, module, device, config) -> None:
         self.config = config
         self.module_device = device
         if not config.compile and getattr(config, "quantization_config", None) is None:
-            # LinearAllreduce and LinearLayer cannot use fused op LinearAdd
+            # LinearAllreduce cannot use fused op LinearAdd
             self.linear_gelu = LinearGelu(module.dense_h_to_4h)
 
             if module.dense_4h_to_h.__class__.__name__ not in ["LinearAllreduce"]:

tests/ipex/test_modeling_causal_lm.py

@@ -70,35 +70,39 @@ def test_compare_to_transformers(self, model_arch):
         ipex_model = IPEXModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=DEVICE)
         self.assertIsInstance(ipex_model.config, PretrainedConfig)
         tokenizer = AutoTokenizer.from_pretrained(model_id)
-        tokens = tokenizer(
-            "This is a sample",
-            return_tensors="pt",
-            return_token_type_ids=False if model_arch in ("llama2",) else None,
-        ).to(DEVICE)
-        inputs = ipex_model.prepare_inputs_for_generation(**tokens)
-        outputs = ipex_model(**inputs)
+        texts = ["This is a sample", ["This is the first input", "This is the second input"]]
+        for text in texts:
+            tokens = tokenizer(
+                text,
+                return_tensors="pt",
+                return_token_type_ids=False if model_arch in ("llama2",) else None,
+            ).to(DEVICE)
+            outputs = ipex_model(**tokens)
+            inputs = ipex_model.prepare_inputs_for_generation(**tokens)
+            outputs_2 = ipex_model(**inputs)
+            self.assertTrue(torch.allclose(outputs.logits, outputs_2.logits, atol=1e-3))
 
-        self.assertIsInstance(outputs.logits, torch.Tensor)
+            self.assertIsInstance(outputs.logits, torch.Tensor)
 
-        transformers_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=DEVICE)
-        with torch.no_grad():
-            transformers_outputs = transformers_model(**tokens)
+            transformers_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=DEVICE)
+            with torch.no_grad():
+                transformers_outputs = transformers_model(**tokens)
 
-        # Test re-load model
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            ipex_model.save_pretrained(tmpdirname)
-            loaded_model = self.IPEX_MODEL_CLASS.from_pretrained(tmpdirname, torch_dtype=dtype, device_map=DEVICE)
-            loaded_model_outputs = loaded_model(**inputs)
+            # Test re-load model
+            with tempfile.TemporaryDirectory() as tmpdirname:
+                ipex_model.save_pretrained(tmpdirname)
+                loaded_model = self.IPEX_MODEL_CLASS.from_pretrained(tmpdirname, torch_dtype=dtype, device_map=DEVICE)
+                loaded_model_outputs = loaded_model(**inputs)
 
-        # Test init method
-        init_model = self.IPEX_MODEL_CLASS(transformers_model)
-        init_model_outputs = init_model(**inputs)
+            # Test init method
+            init_model = self.IPEX_MODEL_CLASS(transformers_model)
+            init_model_outputs = init_model(**inputs)
 
-        # Compare tensor outputs
-        self.assertTrue(torch.allclose(outputs.logits, transformers_outputs.logits, atol=1e-3))
-        # To avoid float pointing error
-        self.assertTrue(torch.allclose(outputs.logits, loaded_model_outputs.logits, atol=1e-7))
-        self.assertTrue(torch.allclose(outputs.logits, init_model_outputs.logits, atol=1e-7))
+            # Compare tensor outputs
+            self.assertTrue(torch.allclose(outputs.logits, transformers_outputs.logits, atol=1e-3))
+            # To avoid float pointing error
+            self.assertTrue(torch.allclose(outputs.logits, loaded_model_outputs.logits, atol=1e-7))
+            self.assertTrue(torch.allclose(outputs.logits, init_model_outputs.logits, atol=1e-7))
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
     @unittest.skip(reason="Paged attention do not support assisted decoding for now")