From 8180e9803c877f41de68c3e93063e315a13e5cff Mon Sep 17 00:00:00 2001 From: Mohamed Mekkouri <93391238+MekkCyber@users.noreply.github.com> Date: Sat, 1 Mar 2025 23:23:22 +0100 Subject: [PATCH] Fix Layernorm when num_cols not a power of 2 (#1867) * fix * Update layernorm.py --------- Co-authored-by: Daniel Han --- unsloth/kernels/layernorm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py index 11f82b8ff3..ffcc5cc13c 100644 --- a/unsloth/kernels/layernorm.py +++ b/unsloth/kernels/layernorm.py @@ -49,7 +49,8 @@ def layernorm_forward( b_row = tl.load(b + col_offsets, mask = mask, other = 0).to(tl.float32) mean_X = tl.sum(X_row, axis = 0) / n_cols - XX = X_row - mean_X + # (X[0] - mean) == -mean so we need to mask it out + XX = tl.where(mask, X_row - mean_X, 0) row_var = tl.sum(XX * XX, axis = 0) / n_cols inv_var = tl.math.rsqrt(row_var + eps) tl.store (r, inv_var)