From 168ded977ecae0bf22308f6ee6ec79b6c63076c8 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 2 Feb 2024 18:28:05 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6907432459..2f426f25a2 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -232,6 +232,8 @@ def LlamaAttention_fast_forward_inference( n_kv_heads = self.num_key_value_heads head_dim = self.head_dim # assert(n_kv_heads * n_groups == n_heads) + seq_len = K1.shape[-2] + kv_seq_len = seq_len + 1 if not hasattr(self, "paged_attention"): self.paged_attention = torch.zeros((2048, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda") @@ -251,8 +253,6 @@ def LlamaAttention_fast_forward_inference( Kn = Kn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2) Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2) - seq_len = K1.shape[-2] - kv_seq_len = seq_len + 1 # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) cos = self.rotary_emb.cos_cached[seq_len]