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:04:29 +11:00
commit c31b27b7dc
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/kernels/rms_layernorm.py
View file

@ -43,8 +43,8 @@ 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.sum(X_row * X_row, axis = 0) / n_cols
inv_var = 1.0 / tl.math.sqrt(row_var + eps)
row_var = tl.mean(X_row * X_row, axis = 0)
inv_var = tl.math.rsqrt(row_var + eps)
tl.store(r, inv_var)
normed = X_row * inv_var
normed = normed.to(W_row.dtype) # Exact copy from HF