REsolved Conflicts

Signed-off-by: Amit Raj <[email protected]>

Author: quic-amitraj

QEfficient/transformers/models/gemma/modeling_gemma.py

@@ -86,8 +86,8 @@ def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     Returns:
         `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
     """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
 
     # Apply rotation
     q_embed = (q * cos) + (rotate_half(q) * sin)

QEfficient/transformers/models/gemma2/modeling_gemma2.py

@@ -40,11 +40,7 @@ class QEffGemma2RotaryEmbedding(Gemma2RotaryEmbedding):
     """
 
     def __init__(self, config: Gemma2Config, device=None):
-<<<<<<< HEAD
         super().__init__(config=config)
-=======
-        Gemma2RotaryEmbedding.__init__(self, config=config)
->>>>>>> 6ba4c76 (Code cleaning and formating)
 
         # Build here to make `torch.jit.trace` work.
         self._set_cos_sin_cache(
@@ -136,16 +132,6 @@ class QEffGemma2Attention(Gemma2Attention):
     - add new args cache idx for the kv retention
     """
 
-<<<<<<< HEAD
-=======
-    def __init__(self, config: Gemma2Config, layer_idx: Optional[int] = None):
-        super().__init__(config, layer_idx)
-        # Define the general __qeff_init__() for any changes in the init calls
-        # Set the init in the module mapping pytorch transforms
-        self.config = config
-        self.__qeff_init__()
-
->>>>>>> 6ba4c76 (Code cleaning and formating)
     def __qeff_init__(self):
         self.rotary_emb = QEffGemma2RotaryEmbedding(config=self.config)
 
@@ -355,10 +341,6 @@ def forward(
                 output_attentions=output_attentions,
                 use_cache=use_cache,
                 cache_position=cache_position,
-<<<<<<< HEAD
-=======
-                position_embeddings=position_embeddings,
->>>>>>> 6ba4c76 (Code cleaning and formating)
                 **kwargs,
             )
 
@@ -384,26 +366,6 @@ def forward(
         )
         return output if return_dict else output.to_tuple()
 
-<<<<<<< HEAD
-=======
-    def _update_causal_mask(
-        self,
-        attention_mask: torch.Tensor,
-        input_tensor: torch.Tensor,
-        cache_position: torch.Tensor,
-        position_ids: torch.Tensor,
-        past_key_values: Cache,
-        output_attentions: bool,
-    ):
-        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-        target_length = attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else past_seen_tokens
-        causal_mask = _create_causal_mask(
-            position_ids=position_ids, target_length=target_length, sliding_window=self.config.sliding_window
-        )
-
-        return causal_mask
-
->>>>>>> 6ba4c76 (Code cleaning and formating)
 
 class QEffGemma2ForCausalLM(Gemma2ForCausalLM, GenerationMixin):
     """

QEfficient/transformers/models/llama/modeling_llama.py

@@ -85,26 +85,15 @@ def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     Returns:
         `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
     """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
 
     # Apply rotation
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
     # Cast back to original dtype
     return q_embed.to(q.dtype), k_embed.to(k.dtype)
 
-def eager_attention_forward(
-    module: nn.Module,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: Optional[torch.Tensor],
-    scaling: float,
-    **kwargs,
-):
-    key_states = repeat_kv(key, module.num_key_value_groups)
-    value_states = repeat_kv(value, module.num_key_value_groups)
 
 def eager_attention_forward(
     module: nn.Module,
@@ -142,6 +131,8 @@ def forward(
         position_ids: Optional[torch.LongTensor] = None,
         past_key_value: Optional[Cache] = None,
         batch_index: Optional[torch.LongTensor] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
         cache_position: Optional[torch.LongTensor] = None,
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:

QEfficient/transformers/models/mistral/modeling_mistral.py

@@ -140,6 +140,8 @@ def forward(
         position_ids: Optional[torch.LongTensor] = None,
         past_key_value: Optional[Cache] = None,
         batch_index: Optional[torch.LongTensor] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
         cache_position: Optional[torch.LongTensor] = None,
         **kwargs,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
@@ -199,7 +201,6 @@ def forward(
         output_attentions: Optional[bool] = False,
         use_cache: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
         **kwargs,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         """
@@ -230,7 +231,6 @@ def forward(
             output_attentions=output_attentions,
             use_cache=use_cache,
             cache_position=cache_position,
-            position_embeddings=position_embeddings,
             **kwargs,
         )
         hidden_states = residual + hidden_states

QEfficient/transformers/models/qwen2/modeling_qwen2.py

@@ -68,10 +68,6 @@ def forward(self, x, seq_len=None):
         )
 
 
-<<<<<<< HEAD
-<<<<<<< HEAD
-=======
->>>>>>> e4503c5 (Minor fixes)
 def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/).
 
@@ -105,24 +101,15 @@ def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
         `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
     """
 
-<<<<<<< HEAD
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-=======
     cos = cos[position_ids].unsqueeze(unsqueeze_dim)
     sin = sin[position_ids].unsqueeze(unsqueeze_dim)
->>>>>>> e4503c5 (Minor fixes)
 
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
 
     return q_embed.to(q.dtype), k_embed.to(k.dtype)
 
 
-<<<<<<< HEAD
-=======
->>>>>>> d0f7ffd (Ruff check and format)
-=======
 def eager_attention_forward(
     module: nn.Module,
     query: torch.Tensor,
@@ -146,7 +133,6 @@ def eager_attention_forward(
     return attn_output, attn_weights
 
 
->>>>>>> e4503c5 (Minor fixes)
 class QEffQwen2Attention(Qwen2Attention):
     """
     Copied from Qwen2Attention: https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2/modeling_qwen2.py
@@ -177,11 +163,7 @@ def forward(
         kv_seq_len = key_states.shape[-2]
 
         kv_seq_len = past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
-<<<<<<< HEAD
-        cos, sin = position_embeddings
-=======
         cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
->>>>>>> e4503c5 (Minor fixes)
         query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
         if past_key_value is not None: