diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 05829cbc94..e3d43b6596 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -194,8 +194,9 @@ class LoRA_MLP_New(torch.autograd.Function): e = matmul_lora(X, gateW, gateW_quant, gateA, gateB, gateS) g = matmul_lora(X, upW, upW_quant, upA, upB, upS) - f = torch.nn.functional.silu(e) - h = f * g + # f = torch.nn.functional.silu(e) + # h = f * g + h = swiglu_fg_kernel(e, g) i = matmul_lora(h, downW, downW_quant, downA, downB, downS) ctx.custom_saved_tensors = ( @@ -208,11 +209,6 @@ class LoRA_MLP_New(torch.autograd.Function): return i pass - def _silu_backward(dy, X): - # https://github.com/pytorch/pytorch/blob/563b065f5a4b4055fa6b025c2514b566d5fd9439/aten/src/ATen/native/Activation.cpp#L483 - sigm = 1 / (1 + torch.exp(-X.float())) - return (dy.float() * sigm * (1 + X.float() * (1 - sigm))).to(X.dtype) - pass @classmethod @torch.cuda.amp.custom_bwd @@ -232,13 +228,15 @@ class LoRA_MLP_New(torch.autograd.Function): g = g .view(-1, g .shape[-1]) dtype = X.dtype - f = torch.nn.functional.silu(e) - h = f * g DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) - df = DW * f # 88us - dg = DW * g # 88us - sigm = 1.0 / (1.0 + torch.exp(-e.float())) - de = (dg.float() * sigm * (1.0 + e.float() * (1.0 - sigm))).to(dtype) + + e = e.float() + se = 1.0 / (1.0 + torch.exp(-e)) + f = se * e + h = f * g + df = DW * f + dg = DW * g + de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype) # Down projection LoRA weights d_downA = h.t() @ (dY @ downB.t())