ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Swiglu

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-27 19:03:04 +11:00
commit 85e87d9ba2
2 changed files with 39 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

19
unsloth/kernels/fast_lora.py
View file

@ -229,14 +229,15 @@ class LoRA_MLP_New(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)
# 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())
@ -256,10 +257,8 @@ class LoRA_MLP_New(torch.autograd.Function):
d_gateA *= gateS
d_gateB *= gateS
# dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS)
# dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS)
upW = fast_dequantize(upW.t(), upW_quant)
dX = torch.matmul(df, upW.t(), out = X)
del upW

56
unsloth/kernels/swiglu.py
View file

@ -49,40 +49,44 @@ pass
@triton.jit
def _DWf_DW_dfg_kernel(DW, gate, up, n_elements, BLOCK_SIZE : tl.constexpr,):
def _DWf_DW_dfg_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,):
"""
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)
"""
block_idx = tl.program_id(0)
offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
DW_row = tl.load(DW + offsets, mask = mask, other = 0)#.to(tl.float32)
gate_row = tl.load(gate + offsets, mask = mask, other = 0)#.to(tl.float32)
up_row = tl.load(up + offsets, mask = mask, other = 0)#.to(tl.float32)
DW_row = tl.load(DW + offsets, mask = mask, other = 0)#.to(tl.float32)
e_row = tl.load(e + offsets, mask = mask, other = 0).to(tl.float32)
g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32)
# f = gate * sigmoid(gate)
se_row = 1 / (1 + tl.exp(-e_row.to(tl.float32)))
se_row = se_row.to(e_row.dtype) # Exact copy from HF
# f = e * se
f_row = e_row * se_row
# e = e.float()
# se = 1.0 / (1.0 + torch.exp(-e))
se_row = 1.0 / (1.0 + tl.exp(-e_row))
# f = (se * e).to(dtype)
f_row = se_row * e_roe
f_row = f_row.to(DW_row.dtype)
# h = f * g
h_row = f_row * g_row
# DW_f = DW * f
DWf_row = DW_row * f_row
# [TODO] Weirdly the below actually loses precision???
# DW_dfg = DW * (se*(g - h) + h)
# DW_dfg_row = DW_row * (se_row*(g_row - h_row) + h_row)
# dh/dgate = sigmoid(gate)*up + gate*up*sigmoid'(gate)
# dh/dgate = sigmoid(gate)*up + gate*up*[sigmoid(gate) * (1 - sigmoid(gate))]
# dh/dgate = sigmoid(gate)*up * [1 + gate*(1 - sigmoid(gate))]
DW_dfg_row = DW_row * se_row * g_row * (1.0 + e_row*(1.0 - se_row)) # 5 FMAs / mults
# DW_dfg_row = DW_row * (se_row * g_row + h_row*(1.0 - se_row)) # 4 FMAs / mults
# DW_dfg_row = DW_row * (se_row*(g_row - h_row) + h_row) BREAKS bad accuracy
h_row = f_row * g_row
# df = DW * f
df_row = DW_row * f_row
# dg = DW * g
dg_row = DW_row * g_row
# de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
de_row = dg_row.to(tl.float32) * se_row * (1.0 + e_row * (1.0 - se_row))
de_row = de_row.to(DW_row.dtype)
# Store derivatives in buffers
tl.store(DW + offsets, h_row, mask = mask)
tl.store(e + offsets, DWf_row, mask = mask)
tl.store(g + offsets, DW_dfg_row, mask = mask)
tl.store(DW + offsets, h_row, mask = mask) # h = f * g
tl.store(e + offsets, df_row, mask = mask) # df = DW * f
tl.store(g + offsets, de_row, mask = mask) # de
pass