From d87ef86991ed82c77bbcffd82d0313e64e555be2 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Tue, 23 Jan 2024 02:22:37 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 8f4a97d058..aadae6ac00 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -168,11 +168,12 @@ def LlamaAttention_fast_forward_inference( Knn, Vnn = Kn, Vn # Attention - A = torch.matmul(Qn, Knn.transpose(2, 3)) + # A = torch.matmul(Qn, Knn.transpose(2, 3)) + A = torch.matmul(Qn, Knn.transpose(1, 2)) A *= 1.0 / math_sqrt(self.head_dim) A[:] = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype) A = torch.matmul(A, Vnn, out = Qn) - A = A.transpose(1, 2) + # A = A.transpose(1, 2) A = A.view(self.hidden_size) # A = self.o_proj(A)