add position_ids in forward

optimum/intel/generation/modeling.py

@@ -70,6 +70,7 @@ def prepare_jit_inputs(model: PreTrainedModel, task: str, use_cache: bool = Fals
 
 def jit_trace(model: PreTrainedModel, task: str, use_cache: bool = False):
     model_inputs = prepare_jit_inputs(model, task, use_cache)
+    has_position_ids = True if "position_ids" in model_inputs else False
     # check if the model_inputs is correct.
     model(**model_inputs)
     torch._C._jit_set_texpr_fuser_enabled(False)
@@ -88,7 +89,7 @@ def jit_trace(model: PreTrainedModel, task: str, use_cache: bool = False):
     traced_model(**model_inputs)
     traced_model(**model_inputs)
 
-    return traced_model
+    return traced_model, has_position_ids
 
 
 class PreTrainedModel(OptimizedModel):
@@ -107,6 +108,7 @@ def __init__(
         config: PretrainedConfig = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         use_cache: bool = True,
+        has_position_ids: bool = False,
         **kwargs,
     ):
         super(BaseModelForCausalLM, self).__init__(model=model, config=config)
@@ -116,6 +118,7 @@ def __init__(
         self._device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         self.normalized_config = NormalizedConfigManager.get_normalized_config_class(config.model_type)(config)
         self.model_dtype = kwargs.get("model_dtype", None)
+        self.has_position_ids = has_position_ids
 
         if is_transformers_version("<=", "4.25.1"):
             self.generation_config = None
@@ -145,7 +148,7 @@ def _save_pretrained(self, save_directory: Union[str, Path], file_name: Optional
         torch.jit.save(self.model, os.path.join(save_directory, WEIGHTS_NAME))
 
     # Adapted from transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel.prepare_inputs_for_generation
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
         past_key_values = past_key_values or kwargs.get("past", None)
 
         if self.use_cache and past_key_values is not None:
@@ -156,11 +159,19 @@ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwarg
             if past_key_values[0][0].shape[0] == input_ids.shape[0]:
                 past_key_values = self._convert_to_bloom_cache(past_key_values)
 
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
         return {
             "input_ids": input_ids,
             "past_key_values": past_key_values,
             "use_cache": self.use_cache,
-            "position_ids": None,
+            "position_ids": position_ids,
             "attention_mask": kwargs.get("attention_mask", None),
             "token_type_ids": None,
         }
@@ -292,8 +303,8 @@ def forward(
                     pkv = tuple(empty_tensor for _ in range(nb_pkv))
                 else:
                     pkv = ()
-                    for nb_pkv in range(nb_pkv):
-                        if nb_pkv % 2 == 0:
+                    for i in range(nb_pkv):
+                        if i % 2 == 0:
                             new_shape = [input_ids.shape[0] * num_key_value_heads, d_k, 0]
                         else:
                             new_shape = [input_ids.shape[0] * num_key_value_heads, 0, d_k]
@@ -305,6 +316,23 @@ def forward(
                     past_key_values = tuple(tuple(pkv) for _ in range(num_layers))
 
             inputs["past_key_values"] = past_key_values
+
+        position_ids = kwargs.get("position_ids", None)
+        if self.has_position_ids and position_ids is not None:
+            inputs.update({"position_ids": position_ids})
+        elif self.has_position_ids and position_ids is None:
+            seq_length = input_ids.shape[-1]
+            if not self.use_cache:
+                past_key_values_length = 0
+            else:
+                past_key_values_length = past_key_values[0][1].shape[-2]
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=self._device
+            ).unsqueeze(0)
+            inputs.update({"position_ids": position_ids})
+        elif not self.has_position_ids and position_ids is not None:
+            logger.warning("You miss the position_ids in the inputs")
+
         outputs = self.model(**inputs)
 
         if isinstance(outputs, (list, tuple)):
@@ -313,6 +341,7 @@ def forward(
         else:
             logits = outputs["logits"]
             past_key_values = outputs["past_key_values"] if self.use_cache else None
+
         return CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values)
 
 
@@ -412,7 +441,7 @@ def _from_transformers(
         if model.config.model_type == "llama":
             model.model._prepare_decoder_attention_mask = _prepare_decoder_attention_mask
 
-        traced_model = jit_trace(model, task, use_cache)
+        traced_model, has_position_ids = jit_trace(model, task, use_cache)
         save_dir = TemporaryDirectory()
         save_dir_path = Path(save_dir.name)
         torch.jit.save(traced_model, save_dir_path / WEIGHTS_NAME)
@@ -427,5 +456,6 @@ def _from_transformers(
             force_download=force_download,
             cache_dir=cache_dir,
             local_files_only=local_files_only,
+            has_position_ids=has_position_ids,
             **kwargs,
         )

optimum/intel/ipex/inference.py

@@ -115,7 +115,7 @@ def __enter__(self):
                                     use_cache = False
                                     if hasattr(self._original.config, "use_cache") and self._original.config.use_cache:
                                         use_cache = True
-                                    model = jit_trace(
+                                    model, has_position_ids = jit_trace(
                                         model=model,
                                         task=self._model.task,
                                         use_cache=use_cache,
@@ -126,6 +126,7 @@ def __enter__(self):
                                             config=self._original.config,
                                             use_cache=use_cache,
                                             model_dtype=self._original.dtype,
+                                            has_position_ids=has_position_ids,
                                         )
                                 except Exception as e:
                                     logger.warning(f"failed to use PyTorch jit mode due to: {e}.")

optimum/intel/neural_compressor/modeling_base.py

@@ -235,7 +235,7 @@ def _from_transformers(
         if task == "text-generation":
             model = patch_decoder_attention_mask(model)
 
-        traced_model = jit_trace(model, task, use_cache)
+        traced_model, has_position_ids = jit_trace(model, task, use_cache)
         save_dir = TemporaryDirectory()
         save_dir_path = Path(save_dir.name)
         torch.jit.save(traced_model, save_dir_path / WEIGHTS_NAME)
@@ -250,6 +250,7 @@ def _from_transformers(
             force_download=force_download,
             cache_dir=cache_dir,
             local_files_only=local_files_only,
+            has_position_ids=has_position_ids,
             **kwargs,
         )