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-27 18:15:34 +11:00
commit 3d3e7f5b42
1 changed files with 11 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

24
unsloth/kernels/fast_lora.py
View file

@ -194,8 +194,9 @@ class LoRA_MLP_New(torch.autograd.Function):
e = matmul_lora(X, gateW, gateW_quant, gateA, gateB, gateS)
g = matmul_lora(X, upW, upW_quant, upA, upB, upS)
f = torch.nn.functional.silu(e)
h = f * g
# f = torch.nn.functional.silu(e)
# h = f * g
h = swiglu_fg_kernel(e, g)
i = matmul_lora(h, downW, downW_quant, downA, downB, downS)
ctx.custom_saved_tensors = (
@ -208,11 +209,6 @@ class LoRA_MLP_New(torch.autograd.Function):
return i
pass
def _silu_backward(dy, X):
# https://github.com/pytorch/pytorch/blob/563b065f5a4b4055fa6b025c2514b566d5fd9439/aten/src/ATen/native/Activation.cpp#L483
sigm = 1 / (1 + torch.exp(-X.float()))
return (dy.float() * sigm * (1 + X.float() * (1 - sigm))).to(X.dtype)
pass
@classmethod
@torch.cuda.amp.custom_bwd
@ -232,13 +228,15 @@ class LoRA_MLP_New(torch.autograd.Function):
g = g .view(-1, g .shape[-1])
dtype = X.dtype
f = torch.nn.functional.silu(e)
h = f * g
DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS)
df = DW * f # 88us
dg = DW * g # 88us
sigm = 1.0 / (1.0 + torch.exp(-e.float()))
de = (dg.float() * sigm * (1.0 + e.float() * (1.0 - sigm))).to(dtype)
e = e.float()
se = 1.0 / (1.0 + torch.exp(-e))
f = se * e
h = f * g
df = DW * f
dg = DW * g
de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
# Down projection LoRA weights
d_downA = h.t() @ (dY @ downB.t())