diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 8dbac0d28b..6ea5345714 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -158,13 +158,15 @@ class LoRA_MLP(torch.autograd.Function): # (D @ W.T * f) @ (U.T + B.T @ A.T) dX = torch.matmul(DW_f, upW.t(), out = X) del upW - dX += upS * (DW_f @ upB.to(dtype).t() @ (upA.to(dtype).t())) + new_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_dX2 = DW_dfg @ gateW.t() + # dX += DW_dfg @ gateW.t() del gateW - dX += gateS * (DW_dfg @ gateB.to(dtype).t() @ (gateA.to(dtype).t())) + new_dX2 += gateS * (DW_dfg @ gateB.to(dtype).t() @ (gateA.to(dtype).t())) + dX += (new_dX + new_dX2) # gateW, gateW_quant, gateA, gateB, gateS, # upW, upW_quant, upA, upB, upS,