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 13:59:57 +11:00
commit 84772de40f
1 changed files with 14 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

20
unsloth/kernels/fast_lora.py
View file

@ -90,7 +90,9 @@ class LoRA_MLP(torch.autograd.Function):
e = matmul_lora(X, gateW, gateW_quant, gateA, gateB, gateS)
g = matmul_lora(X, upW, upW_quant, upA, upB, upS)
h = swiglu_fg_kernel(e, g)
h = torch.nn.functional.silu(e, inplace = True)
h *= g
# h = swiglu_fg_kernel(e, g)
i = matmul_lora(h, downW, downW_quant, downA, downB, downS)
ctx.custom_saved_tensors = (
@ -121,11 +123,17 @@ class LoRA_MLP(torch.autograd.Function):
g = g .view(-1, g .shape[-1])
dtype = X.dtype
# DW_f = (D @ W.T * f)
# DW_dfg = (D @ W.T * df * 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
DW_f = (D @ W.T * f)
DW_dfg = (D @ W.T * df * g)
se = torch.nn.functional.sigmoid(e)
f = e * se
h = f * g
df = se * (1 - f) + f
DW_f = DW * f
DW_dfg = DW * df * 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
# Down projection LoRA weights
d_downA = h.t() @ (dY @ downB.t())