From 7fac7e288d5ff85b4d2427f8882a7e4deaeacdff Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Sep 2024 17:51:36 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index b7dc68586c..07bb5505ef 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -843,9 +843,6 @@ def LlamaModel_fast_forward( pass -global past_key_values_all -past_key_values_all = None - # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825 def LlamaModel_fast_forward_inference( self, @@ -858,8 +855,8 @@ def LlamaModel_fast_forward_inference( hidden_states = self.model.embed_tokens(input_ids) hidden_states = hidden_states.to(self.config.torch_dtype) bsz, q_len, hd = hidden_states.shape - global past_key_values_all - past_key_values_all = past_key_values + import os + os.environ["past_key_values_all"] = past_key_values seq_len = past_key_values[0][0].shape[-2] if bsz != 1: attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(