From 5da05558a036334a893e06e6c1c7e0a9e7e66ce6 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Wed, 31 Jan 2024 03:50:47 +1100 Subject: [PATCH] past_key_value --- unsloth/models/llama.py | 6 +++--- unsloth/models/mistral.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 87aca586b1..6029005265 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -310,7 +310,7 @@ def LlamaAttention_fast_forward( bsz, q_len, _ = hidden_states.size() # Check for inference - if past_key_value is not None and q_len == 1: + if past_key_value is not None: A, past_key_value = LlamaAttention_fast_forward_inference( self, hidden_states, @@ -426,7 +426,7 @@ def LlamaDecoderLayer_fast_forward( past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states """ bsz, q_len, hd = hidden_states.size() - if (past_key_value is not None and q_len == 1): + if past_key_value is not None: # Self Attention residual = hidden_states hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states) @@ -656,7 +656,7 @@ def LlamaModel_fast_forward( pass bsz, q_len, hd = hidden_states.size() - if (past_key_values is not None and q_len == 1): + if past_key_values is not None: hidden_states = fast_rms_layernorm_inference(self.norm, hidden_states) else: hidden_states = fast_rms_layernorm(self.norm, hidden_states) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 9944367d21..42f26b921a 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -49,7 +49,7 @@ def MistralAttention_fast_forward( bsz, q_len, _ = hidden_states.size() # Check for inference - if past_key_value is not None and q_len == 1: + if past_key_value is not None: A, past_key_value = LlamaAttention_fast_forward_inference( self, hidden_states,