diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 9b90a0d164..f542b8ef53 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -92,12 +92,13 @@ def GemmaAttention_fast_forward( K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2) V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2) - if True:#position_ids is None: + 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, position_ids, seq_len=None) + print(cos.shape, sin.shape) Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids) pass @@ -518,8 +519,8 @@ class FastGemmaModel(FastLlamaModel): # Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings. # https://github.com/huggingface/transformers/pull/27931 # https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py - import transformers.models.gemma.modeling_gemma - transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = LlamaRotaryEmbedding + # import transformers.models.gemma.modeling_gemma + # transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = LlamaRotaryEmbedding return pass