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 19:59:59 -08:00
commit 9a68c7fec9
1 changed files with 6 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

15
unsloth/kernels/cross_entropy_loss.py
View file

@ -279,9 +279,6 @@ 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)
@ -294,9 +291,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
labels,
VOCAB_SIZE = vocab_size,
BLOCK_SIZE = BLOCK_SIZE,
DO_SOFTCAPPING = DO_SOFTCAPPING,
DO_SOFTCAPPING = bool(logit_softcapping != 0),
SOFTCAP = logit_softcapping,
DO_LOGIT_SCALING = DO_LOGIT_SCALING,
DO_LOGIT_SCALING = bool(logit_scaling != 0),
LOGIT_SCALE = logit_scaling,
num_warps = num_warps,
)
@ -312,9 +309,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
VOCAB_SIZE = vocab_size,
N_CHUNKS = n_chunks,
BLOCK_SIZE = MAX_FUSED_SIZE,
DO_SOFTCAPPING = DO_SOFTCAPPING,
DO_SOFTCAPPING = bool(logit_softcapping != 0),
SOFTCAP = logit_softcapping,
DO_LOGIT_SCALING = DO_LOGIT_SCALING,
DO_LOGIT_SCALING = bool(logit_scaling != 0),
LOGIT_SCALE = logit_scaling,
num_warps = 32,
)
@ -326,9 +323,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
pass
ctx.save_for_backward(logits, logsumexp, labels)
ctx.DO_SOFTCAPPING = DO_SOFTCAPPING
ctx.DO_SOFTCAPPING = bool(logit_softcapping != 0)
ctx.logit_softcapping = logit_softcapping
ctx.DO_LOGIT_SCALING = DO_LOGIT_SCALING
ctx.DO_LOGIT_SCALING = bool(logit_scaling != 0)
ctx.logit_scaling = logit_scaling
return losses
pass