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

typing

Browse source
This commit is contained in:
Daniel Han 2024-11-03 21:25:32 -08:00
commit 53d98c8560
2 changed files with 11 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/kernels/cross_entropy_loss.py
View file

@ -190,7 +190,7 @@ def _cross_entropy_backward(
logits_ptr ,
logits_row_stride : tl.constexpr(tl.int64),
dloss_ptr ,
dloss_row_stride : tl.constexpr(tl.int32),
dloss_row_stride ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE : tl.constexpr(tl.int32),

14
unsloth/kernels/rms_layernorm.py
View file

@ -53,7 +53,7 @@ def _rms_layernorm_forward(
pass
@triton.heuristics({"GEMMA": lambda args: args["GEMMA"],})
@triton.heuristics({"GEMMA": lambda args: bool(args["GEMMA"]),})
@triton.jit
def _rms_layernorm_backward(
dY, dY_row_stride,
@ -130,11 +130,15 @@ pass
class Fast_RMS_Layernorm(torch.autograd.Function):
@staticmethod
def forward(ctx, X, W, eps, gemma = False):
def forward(ctx, X, W, eps :float, gemma : bool = False):
shape = X.shape
dim = shape[-1]
dim : int = shape[-1]
X = X.view(-1, dim)
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)
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0")
@ -161,9 +165,11 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
@staticmethod
def backward(ctx, dY):
shape = dY.shape
dim = shape[-1]
dim : int = shape[-1]
dY = dY.view(-1, dim)
X, W, r = ctx.saved_tensors
n_rows : int
n_cols : int
n_rows, n_cols = dY.shape
dW = X