Update fast_lora.py

This commit is contained in:
Daniel Han-Chen 2024-01-26 16:16:32 +11:00
commit 317bbc57fe

View file

@ -122,7 +122,7 @@ class LoRA_MLP(torch.autograd.Function):
e = e .view(-1, e .shape[-1])
g = g .view(-1, g .shape[-1])
dtype = X.dtype
DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS)
se = torch.nn.functional.sigmoid(e)
f = e * se
@ -130,6 +130,7 @@ class LoRA_MLP(torch.autograd.Function):
df = se * (1 - f) + f
DW_f = DW * f
DW_dfg = DW * df * g
h, DW_f, DW_dfg = DW, e, g
# DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS)
# DW, e, g = swiglu_DWf_DW_dfg_kernel(DW, e, g)
# h, DW_f, DW_dfg = DW, e, g