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-02-23 18:12:15 +11:00
commit 25339b71f7
1 changed files with 9 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

18
unsloth/kernels/fast_lora.py
View file

@ -128,31 +128,31 @@ class LoRA_MLP(torch.autograd.Function):
h, df, de = DW, e, g
# Down projection LoRA weights
d_downA = (h.t() @ dY) @ downB.t()
d_downB = downA.t() @ (h.t() @ dY)
d_downA = h.t() @ (dY @ downB.t())
d_downB = (downA.t() @ h.t()) @ dY
d_downA *= downS
d_downB *= downS
# Up projection LoRA weights
d_upA = (X.t() @ df) @ upB.t()
d_upB = upA.t() @ (X.t() @ df)
d_upA = X.t() @ (df @ upB.t())
d_upB = (upA.t() @ X.t()) @ df
d_upA *= upS
d_upB *= upS
# Gate projection LoRA weights
d_gateA = (X.t() @ de) @ gateB.t()
d_gateB = gateA.t() @ (X.t() @ de)
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 = 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)
gateW = fast_dequantize(gateW.t(), gateW_quant, out = upW.t())
dX += de @ gateW.t()
del gateW
dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t())