ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update fast_lora.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-27 19:18:59 +11:00
commit 8e0e4ccdc6
1 changed files with 22 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

43
unsloth/kernels/fast_lora.py
View file

@ -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,