From a0a409d5afa2290d4a96357b1bc7ae278dcec353 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 26 Jan 2024 19:20:53 +1100 Subject: [PATCH] Update fast_lora.py --- unsloth/kernels/fast_lora.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 6d1f65b4ee..d6dfd9f92f 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -35,7 +35,7 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None): if A is not None: # LoRA is enabled A, B = A.t(), B.t() - out += s * (X @ A.to(dtype)) @ (B.to(dtype)) + out += (X @ A.to(dtype)) @ (s * B.to(dtype)) pass return out.view(batch, seq_len, -1) if reshape else out @@ -134,20 +134,20 @@ class LoRA_MLP(torch.autograd.Function): # h, DW_f, DW_dfg = DW, e, g # Down projection LoRA weights - d_downA = h.t() @ (dY @ downB.t()) - d_downB = (downA.t() @ h.t()) @ dY + d_downA = h.t() @ dY @ downB.t() + d_downB = downA.t() @ h.t() @ dY d_downA *= downS d_downB *= downS # Up projection LoRA weights - d_upA = X.t() @ (DW_f @ upB.t()) - d_upB = (upA.t() @ X.t()) @ DW_f + d_upA = X.t() @ DW_f @ upB.t() + d_upB = upA.t() @ X.t() @ DW_f d_upA *= upS d_upB *= upS # Gate projection LoRA weights - d_gateA = X.t() @ (DW_dfg @ gateB.t()) - d_gateB = (gateA.t() @ X.t()) @ DW_dfg + d_gateA = X.t() @ DW_dfg @ gateB.t() + d_gateB = gateA.t() @ X.t() @ DW_dfg d_gateA *= gateS d_gateB *= gateS @@ -156,15 +156,15 @@ class LoRA_MLP(torch.autograd.Function): # (D @ W.T * f) @ U.T upW = fast_dequantize(upW.t(), upW_quant) # (D @ W.T * f) @ (U.T + B.T @ A.T) - dX = torch.matmul(DW_f, upW.t(), out = X) + dX = torch.matmul(DW_f, upW.t()) del upW dX += upS * (DW_f @ upB.to(dtype).t() @ (upA.to(dtype).t())) # And add the derivative for the gate projection gateW = fast_dequantize(gateW.t(), gateW_quant) - dX += DW_dfg @ gateW.t() + new_dX = DW_dfg @ gateW.t() + gateS * (DW_dfg @ gateB.to(dtype).t() @ (gateA.to(dtype).t())) del gateW - dX += gateS * (DW_dfg @ gateB.to(dtype).t() @ (gateA.to(dtype).t())) + dX += new_dX # gateW, gateW_quant, gateA, gateB, gateS, # upW, upW_quant, upA, upB, upS,