diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index b487ff95e2..550539d927 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -90,7 +90,8 @@ 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) * g + # h = swiglu_fg_kernel(e, g) i = matmul_lora(h, downW, downW_quant, downA, downB, downS) ctx.custom_saved_tensors = ( @@ -121,29 +122,39 @@ 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 + se = 1 / (1 + torch.exp(-e.float())) + f = torch.nn.functional.silu(e) + h = f * g + DW_f = (DW * f) + DW_dfg = (DW * g).float() * se * (1.0 + e.float() * (1.0 - se)) + DW_dfg = DW_dfg.to(dtype) + # 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()) - d_downB = (downA.t() @ h.t()) @ dY - d_downA *= downS - d_downB *= downS + d_downA = h.t() @ (dY @ (downS * downB).t()) + d_downB = ((downS * downA).t() @ h.t()) @ dY + # d_downA *= downS + # d_downB *= downS # Up projection LoRA weights - d_upA = X.t() @ (DW_f @ upB.t()) - d_upB = (upA.t() @ X.t()) @ DW_f - d_upA *= upS - d_upB *= upS + d_upA = X.t() @ (DW_f @ (upS * upB).t()) + d_upB = ((upS * upA).t() @ X.t()) @ DW_f + # d_upA *= upS + # d_upB *= upS # Gate projection LoRA weights - d_gateA = X.t() @ (DW_dfg @ gateB.t()) - d_gateB = (gateA.t() @ X.t()) @ DW_dfg - d_gateA *= gateS - d_gateB *= gateS + d_gateA = X.t() @ (DW_dfg @ (gateS * gateB).t()) + d_gateB = ((gateS * gateA).t() @ X.t()) @ DW_dfg + # d_gateA *= gateS + # d_gateB *= gateS # Final derivatives to backpropagate backwards. # See our blogpost for more details. @@ -152,13 +163,14 @@ class LoRA_MLP(torch.autograd.Function): # (D @ W.T * f) @ (U.T + B.T @ A.T) dX = torch.matmul(DW_f, upW.t(), out = X) del upW - dX += DW_f @ upB.to(dtype).t() @ (upS * upA.to(dtype).t()) + dX += (DW_f @ upB.to(dtype).t() @ (upS * upA.to(dtype).t())) # And add the derivative for the gate projection gateW = fast_dequantize(gateW.t(), gateW_quant) + # new_dX2 = DW_dfg @ gateW.t() + dX += DW_dfg @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t()) dX += DW_dfg @ gateW.t() del gateW - dX += DW_dfg @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t()) # gateW, gateW_quant, gateA, gateB, gateS, # upW, upW_quant, upA, upB, upS, @@ -171,11 +183,106 @@ class LoRA_MLP(torch.autograd.Function): pass +class LoRA_MLP_New(torch.autograd.Function): + @classmethod + @torch.cuda.amp.custom_fwd + def forward(cls, ctx, X : torch.Tensor, + gateW, gateW_quant, gateA, gateB, gateS, + upW, upW_quant, upA, upB, upS, + downW, downW_quant, downA, downB, downS): + dtype = X.dtype + + 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 + i = matmul_lora(h, downW, downW_quant, downA, downB, downS) + + ctx.custom_saved_tensors = ( + gateW, gateW_quant, gateS, + upW, upW_quant, upS, + downW, downW_quant, downS, + ) + ctx.save_for_backward(gateA, gateB, upA, upB, downA, downB, + X, e, g, f, h, i) + 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 + def backward(cls, ctx, dY : torch.Tensor): + gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, = \ + ctx.custom_saved_tensors + gateA, gateB, upA, upB, downA, downB, \ + X, e, g, f, h, i = ctx.saved_tensors + + gateA, gateB, upA, upB, downA, downB = \ + gateA.t(), gateB.t(), upA.t(), upB.t(), downA.t(), downB.t() + + batch, seq_len, hd = X.shape + dY = dY.view(-1, dY.shape[-1]) + X = X .view(-1, X .shape[-1]) + e = e .view(-1, e .shape[-1]) + g = g .view(-1, g .shape[-1]) + f = f .view(-1, f .shape[-1]) + h = h .view(-1, h .shape[-1]) + i = i .view(-1, i .shape[-1]) + dtype = X.dtype + + DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) + df = DW * f # 88us + dg = DW * g # 88us + de = cls._silu_backward(dg, e) # 90us + + dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) + dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS) + + # Down projection LoRA weights + d_downA = h.t() @ (dY @ downB.t()) + d_downB = (downA.t() @ h.t()) @ dY + d_downA *= downS + d_downB *= downS + + # Up projection LoRA weights + d_upA = X.t() @ (df @ upB.t()) + d_upB = (upA.t() @ X.t()) @ df + d_upA *= upS + d_upB *= upS + + # Gate projection LoRA weights + d_gateA = X.t() @ (dg @ gateB.t()) + d_gateB = (gateA.t() @ X.t()) @ dg + d_gateA *= gateS + d_gateB *= gateS + + # gateW, gateW_quant, gateA, gateB, gateS, + # upW, upW_quant, upA, upB, upS, + # downW, downW_quant, downA, downB, downS, + return dX.view(batch, seq_len, hd), \ + None, None, d_gateA.t(), d_gateB.t(), None, \ + None, None, d_upA.t(), d_upB.t(), None, \ + None, None, d_downA.t(), d_downB.t(), None, + pass +pass + +from transformers.models.llama.modeling_llama import logger def apply_lora_mlp(self, X): + logger.warning_once("Hello!2") + # gate = self.gate_proj(X) + # up = self. up_proj(X) + # h = torch.nn.functional.silu(gate) * up + # down = self.down_proj(h) + # return down gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj) upW, upW_quant, upA, upB, upS = get_lora_parameters(self. up_proj) downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj) - out = LoRA_MLP.apply(X, + out = LoRA_MLP_New.apply(X, gateW, gateW_quant, gateA, gateB, gateS, upW, upW_quant, upA, upB, upS, downW, downW_quant, downA, downB, downS) diff --git a/unsloth/kernels/swiglu.py b/unsloth/kernels/swiglu.py index 4e9b7ba2ae..6eec0b04cb 100644 --- a/unsloth/kernels/swiglu.py +++ b/unsloth/kernels/swiglu.py @@ -49,26 +49,35 @@ pass @triton.jit -def _DWf_DW_dfg_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,): +def _DWf_DW_dfg_kernel(DW, gate, up, n_elements, BLOCK_SIZE : tl.constexpr,): 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) - 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) + 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) - # f = e * sigmoid(e) + # 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 # 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) + # 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 # Store derivatives in buffers tl.store(DW + offsets, h_row, mask = mask)