From 765a7a9330b9745c6a859d59dfa8a24a22d70bb0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jul 2024 14:31:35 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index ee261b0f70..ba45bbbfbb 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -158,6 +158,14 @@ def LlamaAttention_fast_forward_inference( self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0") self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda:0") self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0") + + # Mistral Nemo 12b has weird dimensions + if attention_size != self.hidden_size: + self.temp_O = torch.empty((1, bsz, self.hidden_size), dtype = dtype, device = "cuda:0") + else: + self.temp_O = self.temp_QA[1][:,:,:self.hidden_size] + pass + self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda:0") self.scalar = 1.0 / math_sqrt(self.head_dim) self.half_head_dim = head_dim // 2 @@ -239,7 +247,7 @@ def LlamaAttention_fast_forward_inference( pass A = A.transpose(1, 2) A = A.reshape(bsz, 1, attention_size) - A = fast_linear_forward(self.o_proj, A, out = self.temp_QA[1][:,:,:self.hidden_size]) + A = fast_linear_forward(self.o_proj, A, out = self.temp_O) return A, (Kn, Vn) pass @@ -1152,7 +1160,7 @@ class FastLlamaModel: f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"\ f' "-____-" Free Apache license: http://github.com/unslothai/unsloth' print(statistics) - + # Warn about fast transfers if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", "0") == "1": logger.warning_once("Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!")