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-26 19:41:48 +11:00
commit cb06ce849b
1 changed files with 5 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/kernels/fast_lora.py
View file

@ -158,13 +158,15 @@ class LoRA_MLP(torch.autograd.Function):
# (D @ W.T * f) @ (U.T + B.T @ A.T)
dX = torch.matmul(DW_f, upW.t(), out = X)
del upW
dX += upS * (DW_f @ upB.to(dtype).t() @ (upA.to(dtype).t()))
new_dX = upS * (DW_f @ upB.to(dtype).t() @ (upA.to(dtype).t()))
# And add the derivative for the gate projection
gateW = fast_dequantize(gateW.t(), gateW_quant)
dX += DW_dfg @ gateW.t()
new_dX2 = DW_dfg @ gateW.t()
# dX += DW_dfg @ gateW.t()
del gateW
dX += gateS * (DW_dfg @ gateB.to(dtype).t() @ (gateA.to(dtype).t()))
new_dX2 += gateS * (DW_dfg @ gateB.to(dtype).t() @ (gateA.to(dtype).t()))
dX += (new_dX + new_dX2)
# gateW, gateW_quant, gateA, gateB, gateS,
# upW, upW_quant, upA, upB, upS,