From 84772de40fec66aa99726f11a8f3bf4e992c02cb Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 26 Jan 2024 13:59:57 +1100 Subject: [PATCH] Update fast_lora.py --- unsloth/kernels/fast_lora.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index b487ff95e2..4f91254e6c 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -90,7 +90,9 @@ class LoRA_MLP(torch.autograd.Function): e = matmul_lora(X, gateW, gateW_quant, gateA, gateB, gateS) g = matmul_lora(X, upW, upW_quant, upA, upB, upS) - h = swiglu_fg_kernel(e, g) + h = torch.nn.functional.silu(e, inplace = True) + h *= g + # h = swiglu_fg_kernel(e, g) i = matmul_lora(h, downW, downW_quant, downA, downB, downS) ctx.custom_saved_tensors = ( @@ -121,11 +123,17 @@ class LoRA_MLP(torch.autograd.Function): g = g .view(-1, g .shape[-1]) dtype = X.dtype - # DW_f = (D @ W.T * f) - # DW_dfg = (D @ W.T * df * g) - DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) - DW, e, g = swiglu_DWf_DW_dfg_kernel(DW, e, g) - h, DW_f, DW_dfg = DW, e, g + DW_f = (D @ W.T * f) + DW_dfg = (D @ W.T * df * g) + se = torch.nn.functional.sigmoid(e) + f = e * se + h = f * g + df = se * (1 - f) + f + DW_f = DW * f + DW_dfg = DW * df * g + # DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) + # DW, e, g = swiglu_DWf_DW_dfg_kernel(DW, e, g) + # h, DW_f, DW_dfg = DW, e, g # Down projection LoRA weights d_downA = h.t() @ (dY @ downB.t())