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 2024-11-04 23:34:11 -08:00
commit bf97546a16
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

@ -157,7 +157,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
r, r.stride(0),
n_cols = int(n_cols),
eps = float(eps),
BLOCK_SIZE = BLOCK_SIZE,
BLOCK_SIZE = triton.next_power_of_2(n_cols),
num_warps = 16,
)
else:
@ -167,7 +167,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
W, W.stride(0),
r, r.stride(0),
n_cols, eps,
BLOCK_SIZE = BLOCK_SIZE,
BLOCK_SIZE = triton.next_power_of_2(n_cols),
num_warps = 16,
)
ctx.eps = eps