From 440c29273f55c604f1cc60fccea741d0cd6703cc Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 4 Mar 2024 16:11:12 +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 22f0545167..ccd9f89948 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -43,7 +43,7 @@ def _rms_layernorm_forward( X_row = tl.load(X + col_offsets, mask = mask, other = 0).to(tl.float32) W_row = tl.load(W + col_offsets, mask = mask, other = 0)#.to(tl.float32) - row_var = tl.mean(X_row * X_row, axis = 0) + row_var = tl.sum(X_row * X_row, axis = 0) / n_cols inv_var = tl.math.rsqrt(row_var + eps) tl.store(r, inv_var) normed = X_row * inv_var