Update fast_lora.py

This commit is contained in:
Daniel Han-Chen 2024-01-27 03:12:53 +11:00
commit 83ceb11cc9

View file

@ -161,16 +161,16 @@ 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 += (DW_f @ upB.to(dtype).t() @ (upS * upA.to(dtype).t()))
# And add the derivative for the gate projection
gateW = fast_dequantize(gateW.t(), gateW_quant)
# new_dX2 = DW_dfg @ gateW.t()
dX += DW_dfg @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t())
dX += DW_dfg @ gateW.t()
del gateW
dX += DW_dfg @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t())
# gateW, gateW_quant, gateA, gateB, gateS,
# upW, upW_quant, upA, upB, upS,