From 8e0e4ccdc69de918b79edec65ce62c40bba8ba58 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Sat, 27 Jan 2024 19:18:59 +1100
Subject: [PATCH] Update fast_lora.py

---
 unsloth/kernels/fast_lora.py | 43 ++++++++++++++++++------------------
 1 file changed, 22 insertions(+), 21 deletions(-)

diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py
index cc4b03fef8..c34ef90a05 100644
--- a/unsloth/kernels/fast_lora.py
+++ b/unsloth/kernels/fast_lora.py
@@ -229,15 +229,14 @@ class LoRA_MLP_New(torch.autograd.Function):
         dtype = X.dtype
 
         DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS)
-        # 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
+
+        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)
 
         # Down projection LoRA weights
         d_downA = h.t() @ (dY @ downB.t())
@@ -252,22 +251,24 @@ class LoRA_MLP_New(torch.autograd.Function):
         d_upB  *= upS
 
         # Gate projection LoRA weights
-        d_gateA = X.t() @ (dg @ gateB.t())
-        d_gateB = (gateA.t() @ X.t()) @ dg
+        d_gateA = X.t() @ (de @ gateB.t())
+        d_gateB = (gateA.t() @ X.t()) @ de
         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())
+        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,