From 99eeebf72a85744ac3c5b35eab0418cfe12d697e Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 26 Jan 2024 02:41:15 +1100 Subject: [PATCH] Update swiglu.py --- unsloth/kernels/swiglu.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/kernels/swiglu.py b/unsloth/kernels/swiglu.py index 4e9b7ba2ae..3106656f32 100644 --- a/unsloth/kernels/swiglu.py +++ b/unsloth/kernels/swiglu.py @@ -24,11 +24,11 @@ def _fg_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,): offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < n_elements - e_row = tl.load(e + 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 = e * sigmoid(e) - f_row = e_row / (1 + tl.exp(-e_row)) + f_row = e_row / (1 + tl.exp(-e_row.to(tl.float32)).to(g_row.dtype)) f_row = f_row.to(g_row.dtype) # Exact copy from HF # h = f * g h_row = f_row * g_row