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-Chen 2024-01-22 04:33:44 +11:00
commit 278de9f375
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/models/llama.py
View file

@ -108,7 +108,7 @@ def LlamaAttention_fast_forward_inference(
head_dim = self.head_dim
# assert(n_kv_heads * n_groups == n_heads)
Xn = hidden_states.view(1, self.hidden_size)
Xn = hidden_states.view(self.hidden_size)
K1, V1 = past_key_value
# LoRA or general matrix multiplication
@ -167,7 +167,7 @@ def LlamaAttention_fast_forward_inference(
A[:] = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype)
A = torch.matmul(A, Vnn, out = Qn)
A = A.transpose(1, 2)
A = A.reshape(1, self.hidden_size)
A = A.reshape(self.hidden_size)
# A = self.o_proj(A)
o_proj = self.o_proj
@ -186,7 +186,7 @@ pass
torch_silu = torch.nn.functional.silu
def fast_mlp_inference(self, X):
X = X.view(1, self.hidden_size)
X = X.view(self.hidden_size)
dtype = X.dtype
gate_proj = self.gate_proj
up_proj = self.up_proj