update prepare_inputs_for_generation

Author: eaidova

optimum/exporters/openvino/model_patcher.py

@@ -1608,25 +1608,28 @@ def _phi3_self_attn_sdpa_forward(
     return attn_output, None, past_key_value
 
 
-@torch.jit.script
-def select_ext_factor(seq_len: torch.Tensor, max_pos_embeddings: torch.Tensor, short_factor: torch.Tensor, long_factor: torch.Tensor):
-    if seq_len > max_pos_embeddings:
-        return long_factor
-    return short_factor
+# @torch.jit.script
+def select_ext_factor(
+    seq_len: torch.Tensor, max_pos_embeddings: torch.Tensor, short_factor: torch.Tensor, long_factor: torch.Tensor
+):
+    return torch.where(
+        seq_len < max_pos_embeddings, short_factor, long_factor
+    )  # short_factor * (seq_len <= max_pos_embeddings) + long_factor * (seq_len > max_pos_embeddings)
+
 
 def long_rope(self, x, position_ids, seq_len=None):
     seq_len = torch.max(position_ids) + 1
     original_max_position_embeddings = (
         self.original_max_position_embeddings
-        if hasattr(self, "original_max_positional_embeddings") else self.config.original_max_position_embeddings
+        if hasattr(self, "original_max_positional_embeddings")
+        else self.config.original_max_position_embeddings
     )
-    max_position_embeddings = self.max_position_embeddings if hasattr(self, "max_position_embeddings") else self.config.max_position_embeddings
-    inv_freq = select_ext_factor(
-        seq_len,
-        torch.tensor(original_max_position_embeddings),
-        self.inv_freq,
-        self.long_inv_freq
+    max_position_embeddings = (
+        self.max_position_embeddings
+        if hasattr(self, "max_position_embeddings")
+        else self.config.max_position_embeddings
     )
+    inv_freq = select_ext_factor(seq_len, original_max_position_embeddings, self.inv_freq, self.long_inv_freq)
 
     inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
     position_ids_expanded = position_ids[:, None, :].float()
@@ -1679,9 +1682,16 @@ def __enter__(self):
                 layer.self_attn.rotary_emb.inv_freq = 1.0 / (
                     rotary_emb.base ** (torch.arange(0, rotary_emb.dim, 2, dtype=torch.int64).float() / rotary_emb.dim)
                 )
-        
-        if hasattr(self._model.model, "rotary_emb") and getattr(self._model.model.rotary_emb, "rope_type", "default") == "longrope":
-            long_inv_freq, _ = self._model.model.rotary_emb.rope_init_fn(self._model.config, torch.device("cpu"), seq_len=self._model.config.original_max_position_embeddings + 1)
+
+        if (
+            hasattr(self._model.model, "rotary_emb")
+            and getattr(self._model.model.rotary_emb, "rope_type", "default") == "longrope"
+        ):
+            long_inv_freq, _ = self._model.model.rotary_emb.rope_init_fn(
+                self._model.config,
+                torch.device("cpu"),
+                seq_len=self._model.config.original_max_position_embeddings + 1,
+            )
             self._model.model.rotary_emb.long_inv_freq = long_inv_freq
             self._model.model.rotary_emb._orig_forward = self._model.model.rotary_emb.forward
             self._model.model.rotary_emb.forward = types.MethodType(long_rope, self._model.model.rotary_emb)
@@ -1690,7 +1700,6 @@ def __enter__(self):
         ):
             self._model.config.max_position_embeddings = self._model.config.original_max_position_embeddings
 
-
     def __exit__(self, exc_type, exc_value, traceback):
         super().__exit__(exc_type, exc_value, traceback)
         if hasattr(self._model.model, "_orig_forward"):

optimum/intel/openvino/modeling_decoder.py

@@ -846,6 +846,8 @@ def _from_pretrained(
             init_cls = OVBloomForCausalLM
         elif model_type == "gpt-bigcode":
             init_cls = OVGPTBigCodeForCausalLM
+        elif model_type == "phi3":
+            init_cls = OVPhi3ForCausalLM
         else:
             init_cls = cls
 
@@ -915,6 +917,47 @@ def _from_pretrained(
         return causal_model
 
 
+class OVPhi3ForCausalLM(OVModelForCausalLM):
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        logits_to_keep=None,
+        **kwargs,
+    ):
+        # Overwritten -- this model may need to switch between short and long rope, invalidating the cache in the
+        # process
+
+        # When the first time input length reached long and short factor switching point, enforce re-compute cache
+        # It will cause downside of slower at this single token position, however, better than current failure.
+        if (
+            past_key_values
+            and self.config.rope_scaling
+            and input_ids.shape[1] >= self.config.original_max_position_embeddings + 1
+        ):
+            past_length = cache_position[0]
+            if past_length <= self.config.original_max_position_embeddings:
+                past_key_values = None
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            use_cache=use_cache,
+            logits_to_keep=logits_to_keep,
+            **kwargs,
+        )
+        return model_inputs
+
+
 class OVBloomForCausalLM(OVModelForCausalLM):
     # Adapted from transformers.models.bloom.modeling_bloom.BloomForCausalLM.prepare_inputs_for_generation
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):