Update fast_lora.py

This commit is contained in:
Daniel Han-Chen 2024-01-27 20:39:19 +11:00
commit e094af81a1

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())