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 2025-02-02 13:47:14 -08:00
commit 50d5250a57
1 changed files with 5 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

10
unsloth/models/llama.py
View file

@ -920,6 +920,11 @@ def LlamaModel_fast_forward_inference(
bsz, q_len = input_ids.shape
hd = self.config.hidden_size
mlp_size = self.config.intermediate_size
X = self.model.embed_tokens(input_ids)
X = X.to(self.config.torch_dtype)
bsz, q_len, hd = X.shape
assert(q_len == 1)
# Get saved buffers to reduce memory movement
residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0")
@ -929,11 +934,6 @@ def LlamaModel_fast_forward_inference(
temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
temp_gate, temp_up = temp_mlp[0], temp_mlp[1]
X = self.model.embed_tokens(input_ids)
X = X.to(self.config.torch_dtype)
bsz, q_len, hd = X.shape
assert(q_len == 1)
seq_len = past_key_values[0][0].shape[-2]
if bsz != 1:
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(