ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update rms_layernorm.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-03-04 16:00:35 +11:00
commit 6fa081ae77
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/kernels/rms_layernorm.py
View file

@ -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