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:20:53 +11:00
commit a0a409d5af
1 changed files with 10 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

20
unsloth/kernels/fast_lora.py
View file

@ -35,7 +35,7 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None):
if A is not None:
# LoRA is enabled
A, B = A.t(), B.t()
out += s * (X @ A.to(dtype)) @ (B.to(dtype))
out += (X @ A.to(dtype)) @ (s * B.to(dtype))
pass
return out.view(batch, seq_len, -1) if reshape else out
@ -134,20 +134,20 @@ class LoRA_MLP(torch.autograd.Function):
# h, DW_f, DW_dfg = 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() @ (DW_f @ upB.t())
d_upB = (upA.t() @ X.t()) @ DW_f
d_upA = X.t() @ DW_f @ upB.t()
d_upB = upA.t() @ X.t() @ DW_f
d_upA *= upS
d_upB *= upS
# Gate projection LoRA weights
d_gateA = X.t() @ (DW_dfg @ gateB.t())
d_gateB = (gateA.t() @ X.t()) @ DW_dfg
d_gateA = X.t() @ DW_dfg @ gateB.t()
d_gateB = gateA.t() @ X.t() @ DW_dfg
d_gateA *= gateS
d_gateB *= gateS
@ -156,15 +156,15 @@ class LoRA_MLP(torch.autograd.Function):
# (D @ W.T * f) @ U.T
upW = fast_dequantize(upW.t(), upW_quant)
# (D @ W.T * f) @ (U.T + B.T @ A.T)
dX = torch.matmul(DW_f, upW.t(), out = X)
dX = torch.matmul(DW_f, upW.t())
del upW
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_dX = DW_dfg @ gateW.t() + gateS * (DW_dfg @ gateB.to(dtype).t() @ (gateA.to(dtype).t()))
del gateW
dX += gateS * (DW_dfg @ gateB.to(dtype).t() @ (gateA.to(dtype).t()))
dX += new_dX
# gateW, gateW_quant, gateA, gateB, gateS,
# upW, upW_quant, upA, upB, upS,