From 6fa081ae77eed48e051134852e9068889afd8e24 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 4 Mar 2024 16:00:35 +1100 Subject: [PATCH] Update rms_layernorm.py --- unsloth/kernels/rms_layernorm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index ccd9f89948..ca2c75a0c2 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -44,7 +44,7 @@ def _rms_layernorm_forward( W_row = tl.load(W + col_offsets, mask = mask, other = 0)#.to(tl.float32) row_var = tl.sum(X_row * X_row, axis = 0) / n_cols - inv_var = tl.math.rsqrt(row_var + eps) + inv_var = 1.0 / tl.math.sqrt(row_var + eps) tl.store(r, inv_var) normed = X_row * inv_var normed = normed.to(W_row.dtype) # Exact copy from HF