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-24 13:33:05 +11:00
commit 9ab87f3648
1 changed files with 20 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

40
unsloth/kernels/cross_entropy_loss.py
View file

@ -231,7 +231,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
pass
# slow_cross_entropy_loss = torch.nn.functional.cross_entropy
slow_cross_entropy_loss = torch.nn.functional.cross_entropy
def fast_cross_entropy_loss(logits, labels):
"""
Arguments:
@ -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