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-10-30 16:37:37 -07:00
commit 6ef574b58a
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/kernels/cross_entropy_loss.py
View file

@ -223,7 +223,7 @@ def _cross_entropy_backward(
else:
dloss = 0.0
x = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf"))
x = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32)
# Do logit scaling for Cohere
if DO_LOGIT_SCALING:
@ -239,7 +239,7 @@ def _cross_entropy_backward(
pass
logsumexp = tl.load(logsumexp_ptr + row_idx)
y = tl.exp(x.to(tl.float32) - logsumexp)
y = tl.exp(x - logsumexp)
y = tl.where(
col_offsets == label_idx,
y - 1.0, # exp(x - logsumexp) - 1