From 35daafdd6e3b233508a569835feb40a2acc0cb61 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 27 Jan 2024 19:53:57 +1100 Subject: [PATCH] Update fast_lora.py --- unsloth/kernels/fast_lora.py | 44 ++++++++++++++++++++++-------------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 71c5408677..441c01be52 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 = ( @@ -204,15 +205,10 @@ class LoRA_MLP_New(torch.autograd.Function): downW, downW_quant, downS, ) ctx.save_for_backward(gateA, gateB, upA, upB, downA, downB, - X, e, g, f, h, i) + X, e, g) 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 @@ -220,7 +216,7 @@ class LoRA_MLP_New(torch.autograd.Function): gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, = \ ctx.custom_saved_tensors gateA, gateB, upA, upB, downA, downB, \ - X, e, g, f, h, i = ctx.saved_tensors + X, e, g = ctx.saved_tensors gateA, gateB, upA, upB, downA, downB = \ gateA.t(), gateB.t(), upA.t(), upB.t(), downA.t(), downB.t() @@ -230,18 +226,19 @@ class LoRA_MLP_New(torch.autograd.Function): X = X .view(-1, X .shape[-1]) e = e .view(-1, e .shape[-1]) g = g .view(-1, g .shape[-1]) - f = f .view(-1, f .shape[-1]) - h = h .view(-1, h .shape[-1]) - i = i .view(-1, i .shape[-1]) dtype = X.dtype DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) - df = DW * f # 88us - dg = DW * g # 88us - de = cls._silu_backward(dg, e) # 90us - dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) - dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS) + # e = e.float() + # se = 1.0 / (1.0 + torch.exp(-e)) + # f = (se * e).to(dtype) + # h = f * g + # df = DW * f + # dg = DW * g + # de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype) + DW, e, g = swiglu_DWf_DW_dfg_kernel(DW, e, g) + h, df, de = DW, e, g # Down projection LoRA weights d_downA = h.t() @ (dY @ downB.t()) @@ -261,6 +258,19 @@ class LoRA_MLP_New(torch.autograd.Function): d_gateA *= gateS d_gateB *= gateS + # dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) + # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS) + + upW = fast_dequantize(upW.t(), upW_quant) + dX = torch.matmul(df, upW.t(), out = X) + del upW + dX += df @ upB.to(dtype).t() @ (upS * upA.to(dtype).t()) + + gateW = fast_dequantize(gateW.t(), gateW_quant) + dX += de @ gateW.t() + del gateW + dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t()) + # gateW, gateW_quant, gateA, gateB, gateS, # upW, upW_quant, upA, upB, upS, # downW, downW_quant, downA, downB, downS,