Update fast_lora.py

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,