From 00f50876f7b308ea12e2f498e40d24bf2576d516 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Tue, 23 Jan 2024 02:03:25 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 9dd410a1e2..c30025768d 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -112,9 +112,8 @@ def LlamaAttention_fast_forward_inference( Xn = hidden_states.view(self.hidden_size) K1, V1 = past_key_value seq_len = K1.shape[1] - position_id = seq_len + 1 - K1 = K1.view(n_kv_heads, position_id, head_dim) - V1 = V1.view(n_kv_heads, position_id, head_dim) + K1 = K1.view(n_kv_heads, seq_len, head_dim) + V1 = V1.view(n_kv_heads, seq_len, head_dim) # LoRA or general matrix multiplication dtype = Xn.dtype @@ -135,6 +134,7 @@ def LlamaAttention_fast_forward_inference( # kv_seq_len = K1.shape[-2] + 1 # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) + position_id = seq_len + 1 cos = self.rotary_emb.cos_cached[position_id] sin = self.rotary_emb.sin_cached[position_id] h = head_dim // 2