This commit is contained in:
Daniel Han-Chen 2024-01-23 03:36:50 +11:00
commit 1bb1c3c2b9
2 changed files with 2 additions and 1 deletions

View file

@ -183,7 +183,7 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None):
if lora_A is not None:
dtype = X.dtype
temp_lora = torch.matmul(X, lora_A.to(dtype).t(), out = temp_lora)
out.addmv_(lora_B.to(dtype).t(), temp_lora, alpha = lora_S)
out.addmv_(lora_B.to(dtype), temp_lora, alpha = lora_S)
pass
return out
pass

View file

@ -180,6 +180,7 @@ def LlamaAttention_fast_forward_inference(
A = fast_linear_forward(self.o_proj, A)
A = A.reshape(1, 1, self.hidden_size)
# return A, (Kn, Vn)
return A, (Kn.unsqueeze(0), Vn.unsqueeze(0))
pass