Update fast_lora.py

This commit is contained in:
Daniel Han-Chen 2024-01-28 00:45:55 +11:00
commit d01ba458df

View file

@ -125,15 +125,15 @@ class LoRA_MLP(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)
# DW, e, g = swiglu_DWf_DW_dfg_kernel(DW, e, g)
# h, df, de = DW, e, g
# Down projection LoRA weights
d_downA = h.t() @ (dY @ downB.t())
@ -148,8 +148,8 @@ class LoRA_MLP(torch.autograd.Function):
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() @ (dg @ gateB.t())
d_gateB = (gateA.t() @ X.t()) @ dg
d_gateA *= gateS
d_gateB *= gateS