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:24:33 -08:00
commit d966cbd5ea
1 changed files with 14 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

21
unsloth/kernels/rms_layernorm.py
View file

@ -16,6 +16,7 @@ import triton
import triton.language as tl
import torch
from .utils import calculate_settings
next_power_of_2 = triton.next_power_of_2
@triton.jit
@ -142,9 +143,15 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
n_rows : int
n_cols : int
n_rows, n_cols = X.shape
BLOCK_SIZE : int
num_warps : int
# BLOCK_SIZE, num_warps = calculate_settings(n_cols)
BLOCK_SIZE : int = next_power_of_2(n_cols)
MAX_FUSED_SIZE : int = 65536
if BLOCK_SIZE > MAX_FUSED_SIZE:
raise RuntimeError(f"Cannot launch Triton kernel since n = {n_cols} exceeds "\
f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.")
num_warps : int = 4
if BLOCK_SIZE >= 32768: num_warps = 32
elif BLOCK_SIZE >= 8192: num_warps = 16
elif BLOCK_SIZE >= 2048: num_warps = 8
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0")
r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
@ -157,8 +164,8 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
r, r.stride(0),
n_cols = int(n_cols),
eps = float(eps),
BLOCK_SIZE = 4096,
num_warps = 16,
BLOCK_SIZE = BLOCK_SIZE,
num_warps = num_warps,
)
else:
_gemma_rms_layernorm_forward[(n_rows,)](
@ -171,8 +178,8 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
num_warps = num_warps,
)
ctx.eps = eps
ctx.BLOCK_SIZE = 4096
ctx.num_warps = 16
ctx.BLOCK_SIZE = BLOCK_SIZE
ctx.num_warps = num_warps
ctx.GEMMA = gemma
ctx.save_for_backward(X, W, r)
return Y.view(*shape)