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:11:12 +11:00
commit 440c29273f
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

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