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

Update cross_entropy_loss.py

Browse source
This commit is contained in:
Daniel Han 2024-11-03 20:58:31 -08:00
commit ad79b86e86
1 changed files with 23 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

43
unsloth/kernels/cross_entropy_loss.py
View file

@ -37,7 +37,7 @@ def _cross_entropy_forward(
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr ,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,
@ -69,7 +69,7 @@ def _cross_entropy_forward(
logsumexp_ptr += row_idx
labels_ptr += row_idx
col_offsets = tl.arange(0, BLOCK_SIZE)
col_offsets = tl.arange(0, BLOCK_SIZE : tl.constexpr)
mask = col_offsets < VOCAB_SIZE
label_idx = tl.load(labels_ptr).to(tl.int32)
@ -111,7 +111,7 @@ def _chunked_cross_entropy_forward(
labels_ptr ,
VOCAB_SIZE ,
N_CHUNKS ,
BLOCK_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr ,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,
@ -148,7 +148,7 @@ def _chunked_cross_entropy_forward(
logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx
labels_ptr += row_idx
col_offsets = chunk_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
col_offsets = chunk_idx*BLOCK_SIZE : tl.constexpr + tl.arange(0, BLOCK_SIZE : tl.constexpr)
mask = col_offsets < VOCAB_SIZE
label_idx = tl.load(labels_ptr).to(tl.int32)
@ -194,7 +194,7 @@ def _cross_entropy_backward(
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr ,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,
@ -220,7 +220,7 @@ def _cross_entropy_backward(
logits_ptr += row_idx * logits_row_stride
dloss_ptr += row_idx * dloss_row_stride
col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
col_offsets = block_idx*BLOCK_SIZE : tl.constexpr + tl.arange(0, BLOCK_SIZE : tl.constexpr)
mask = col_offsets < VOCAB_SIZE
label_idx = tl.load(labels_ptr + row_idx).to(tl.int32)
@ -279,9 +279,12 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
n_chunks = div + (mod != 0)
losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
DO_SOFTCAPPING : bool = bool(logit_softcapping != 0)
DO_LOGIT_SCALING : bool = bool(logit_scaling != 0)
if n_chunks == 1:
# For small vocabs <= 65336 like Llama, Mistral
BLOCK_SIZE, num_warps = calculate_settings(vocab_size)
BLOCK_SIZE : tl.constexpr, num_warps = calculate_settings(vocab_size)
logsumexp = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
_cross_entropy_forward[(n_rows,)](
@ -290,10 +293,10 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
logsumexp,
labels,
VOCAB_SIZE = vocab_size,
BLOCK_SIZE = BLOCK_SIZE,
DO_SOFTCAPPING = bool(logit_softcapping != 0),
BLOCK_SIZE : tl.constexpr = BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING = DO_SOFTCAPPING,
SOFTCAP = logit_softcapping,
DO_LOGIT_SCALING = bool(logit_scaling != 0),
DO_LOGIT_SCALING = DO_LOGIT_SCALING,
LOGIT_SCALE = logit_scaling,
num_warps = num_warps,
)
@ -308,10 +311,10 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
labels,
VOCAB_SIZE = vocab_size,
N_CHUNKS = n_chunks,
BLOCK_SIZE = MAX_FUSED_SIZE,
DO_SOFTCAPPING = bool(logit_softcapping != 0),
BLOCK_SIZE : tl.constexpr = MAX_FUSED_SIZE,
DO_SOFTCAPPING = DO_SOFTCAPPING,
SOFTCAP = logit_softcapping,
DO_LOGIT_SCALING = bool(logit_scaling != 0),
DO_LOGIT_SCALING = DO_LOGIT_SCALING,
LOGIT_SCALE = logit_scaling,
num_warps = 32,
)
@ -323,9 +326,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
pass
ctx.save_for_backward(logits, logsumexp, labels)
ctx.DO_SOFTCAPPING = bool(logit_softcapping != 0)
ctx.DO_SOFTCAPPING = DO_SOFTCAPPING
ctx.logit_softcapping = logit_softcapping
ctx.DO_LOGIT_SCALING = bool(logit_scaling != 0)
ctx.DO_LOGIT_SCALING = DO_LOGIT_SCALING
ctx.logit_scaling = logit_scaling
return losses
pass
@ -335,8 +338,8 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
logits, logsumexp, labels = ctx.saved_tensors
n_rows, vocab_size = logits.shape
BLOCK_SIZE = 4096
div, mod = divmod(vocab_size, BLOCK_SIZE)
BLOCK_SIZE : tl.constexpr = 4096
div, mod = divmod(vocab_size, BLOCK_SIZE : tl.constexpr)
n_blocks = div + (mod != 0)
_cross_entropy_backward[(n_rows, n_blocks,)](
@ -345,10 +348,10 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
logsumexp,
labels,
VOCAB_SIZE = vocab_size,
BLOCK_SIZE = BLOCK_SIZE,
DO_SOFTCAPPING = bool(ctx.DO_SOFTCAPPING),
BLOCK_SIZE : tl.constexpr = BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING = ctx.DO_SOFTCAPPING,
SOFTCAP = ctx.logit_softcapping,
DO_LOGIT_SCALING = bool(ctx.DO_LOGIT_SCALING),
DO_LOGIT_SCALING = ctx.DO_LOGIT_SCALING,
LOGIT_SCALE = ctx.logit_scaling,
num_warps = 8,
)