ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update llama.py

Browse source
This commit is contained in:
Daniel Han 2024-09-25 17:51:36 -07:00
commit 7fac7e288d
1 changed files with 2 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

7
unsloth/models/llama.py
View file

@ -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(