diff --git a/unsloth/kernels/geglu.py b/unsloth/kernels/geglu.py index 06ef2078c5..7001b8ff0a 100644 --- a/unsloth/kernels/geglu.py +++ b/unsloth/kernels/geglu.py @@ -27,10 +27,10 @@ def _forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,): # f = 1/2 * e * (1 + erf(1/sqrt(2) * e)) # h = f * up 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) + g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32) f_row = 0.5 * e_row * (tl.math.erf(tl.math.rsqrt(2.0) * e_row) + 1.0) - # f_row = f_row.to(g_row.dtype) # Exact copy from HF + f_row = f_row.to(g_row.dtype) # Exact copy from HF h_row = f_row * g_row # Store h @@ -64,16 +64,16 @@ def _backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,): 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) + 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) + g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32) # Break e_row away for re-use # f = 1/2 * e * (1 + erf(1/sqrt(2) * e)) f_partial_row = 0.5 * (tl.math.erf(tl.math.rsqrt(2.0) * e_row) + 1.0) f_row = f_partial_row * e_row - # f_row = f_row.to(DW_row.dtype) + f_row = f_row.to(DW_row.dtype) # h = f * g h_row = f_row * g_row # df = DW * f diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 5242a1e23b..81af8a1cee 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1400,7 +1400,7 @@ class FastLlamaModel: (down_proj.base_layer if hasattr(down_proj, "base_layer") else down_proj).bias is None: # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module - layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp) + # layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp) n_mlp += 1 else: logger.warning_once(