diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index aa0bcb55ee..cbbebeec5c 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -75,6 +75,7 @@ def CohereAttention_fast_forward( output_attentions: bool = False, use_cache: bool = False, padding_mask: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *args, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: @@ -112,12 +113,11 @@ def CohereAttention_fast_forward( if past_key_value is not None: kv_seq_len += past_key_value[0].shape[-2] + cos, sin = position_embeddings if position_ids is None: - cos = self.rotary_emb.cos_cached - sin = self.rotary_emb.sin_cached Q, K = fast_rope_embedding(Q, K, cos, sin) else: - cos, sin = self.rotary_emb(V, seq_len = kv_seq_len) + cos, sin = cos[position_ids], sin[position_ids] Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids) pass @@ -190,6 +190,7 @@ def CohereDecoderLayer_fast_forward( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, padding_mask: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *args, **kwargs, ): if use_cache and hasattr(self, "_flag_for_generation"): #past_key_value is not None: diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index cfeeab0800..03336c134a 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -337,6 +337,7 @@ def LlamaAttention_fast_forward( output_attentions: bool = False, use_cache: bool = False, padding_mask: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *args, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: @@ -368,20 +369,24 @@ def LlamaAttention_fast_forward( if past_key_value is not None: kv_seq_len += past_key_value[0].shape[-2] - # Extend RoPE dynamically to fit in VRAM - rotary_emb = self.rotary_emb - rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len) - - if position_ids is None: - # Useful for LongRoPE - cos, sin = rotary_emb.get_cached(kv_seq_len) - # cos = self.rotary_emb.cos_cached - # sin = self.rotary_emb.sin_cached - Q, K = fast_rope_embedding(Q, K, cos, sin) + if position_embeddings: + cos, sin = position_embeddings else: - cos, sin = rotary_emb(V, seq_len = kv_seq_len) - Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids) - pass + # Extend RoPE dynamically to fit in VRA + rotary_emb = self.rotary_emb + rotary_emb.extend_rope_embedding(V, seq_len=kv_seq_len) + + if position_ids is None: + # Useful for LongRoPE + cos, sin = rotary_emb.get_cached(kv_seq_len) + else: + cos, sin = rotary_emb(V, seq_len=kv_seq_len) + + Q, K = ( + fast_rope_embedding(Q, K, cos, sin) + if position_ids is None + else inplace_rope_embedding(Q, K, cos, sin, position_ids) + ) if past_key_value is not None: K = torch.cat([past_key_value[0], K], dim = 2) @@ -452,6 +457,7 @@ def LlamaDecoderLayer_fast_forward( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, padding_mask: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *args, **kwargs, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: """ @@ -479,6 +485,7 @@ def LlamaDecoderLayer_fast_forward( output_attentions=output_attentions, use_cache=use_cache, padding_mask=padding_mask, + position_embeddings = position_embeddings, ) hidden_states += residual @@ -499,6 +506,7 @@ def LlamaDecoderLayer_fast_forward( output_attentions=output_attentions, use_cache=use_cache, padding_mask=padding_mask, + position_embeddings = position_embeddings, ) hidden_states = residual + hidden_states @@ -777,8 +785,11 @@ def LlamaModel_fast_forward( pass - if IS_GRANITE: - position_embeddings = self.rotary_emb(hidden_states, position_ids, self.max_position_embeddings) + if transformers_version > "4.47.1" and hasattr(self,'rotary_emb'): + # Transformers main has made it mandatory to pass position_embeddings + # https://github.com/huggingface/transformers/pull/34858 + position_embeddings = self.rotary_emb(hidden_states, position_ids, self.config.max_position_embeddings) + print(f'position_embeddings: {position_embeddings}') else: position_embeddings = None diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 00dcc5cd1d..cdc73c7610 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -47,6 +47,7 @@ def MistralAttention_fast_forward( output_attentions: bool = False, use_cache: bool = False, padding_mask: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *args, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: