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-Chen 2024-02-25 01:39:25 +11:00
commit 0c2d7e503b
1 changed files with 19 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

38
unsloth/kernels/cross_entropy_loss.py
View file

@ -246,23 +246,23 @@ def fast_cross_entropy_loss(logits, labels):
# We now support any vocab size due to Gemma!
# Prelim support Qwen, Deepseek other large vocab sizes > 2^16
if d > MAX_FUSED_SIZE:
logger.warning_once(
f"Unsloth: Vocab size of {d} exceeds the max CUDA blocksize of {MAX_FUSED_SIZE}.\n"\
"For now, Unsloth will use Pytorch's CrossEntropyLoss, which will entail a\n"\
"25% increase in memory usage and be slower. Make an issue on \n"\
"Unsloth's Github page if you want a faster and more memory efficient kernel!"
)
loss = slow_cross_entropy_loss(
logits.float().view(batch*seq_len, d), # Must cast to float32 for numerical stability
labels.view(-1),
)
return loss
else:
loss = Fast_CrossEntropyLoss.apply(
logits.view(batch*seq_len, d),
labels.view(-1),
)
n_items = torch.count_nonzero(labels != -100)
return loss.sum() / n_items
# if d > MAX_FUSED_SIZE:
# logger.warning_once(
# f"Unsloth: Vocab size of {d} exceeds the max CUDA blocksize of {MAX_FUSED_SIZE}.\n"\
# "For now, Unsloth will use Pytorch's CrossEntropyLoss, which will entail a\n"\
# "25% increase in memory usage and be slower. Make an issue on \n"\
# "Unsloth's Github page if you want a faster and more memory efficient kernel!"
# )
# loss = slow_cross_entropy_loss(
# logits.float().view(batch*seq_len, d), # Must cast to float32 for numerical stability
# labels.view(-1),
# )
# return loss
# else:
loss = Fast_CrossEntropyLoss.apply(
logits.view(batch*seq_len, d),
labels.view(-1),
)
n_items = torch.count_nonzero(labels != -100)
return loss.sum() / n_items
pass