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 02:06:11 +11:00
commit 88bf684c61
1 changed files with 13 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

21
unsloth/kernels/cross_entropy_loss.py
View file

@ -99,16 +99,19 @@ def _large_cross_entropy_forward(
# Maximum stops overflow
lse = tl.log(tl.sum(tl.exp(logits - max_logits), 0)) + max_logits
tl.store(lse_ptr, lse)
if (label_idx != -100) and \
(label_idx >= (col_idx+0)*BLOCK_SIZE) and \
(label_idx < min((col_idx+1)*BLOCK_SIZE, n_cols)):
loss = tl.load(logits_ptr + label_idx).to(tl.float32)
lse = 0.0
loss = lse - logits_label # We add the final logsumexp after a reduction
else:
loss = 0.0
loss = 0.0
# chained boolean operators (A or B or C) are not supported; use parentheses to split the chain.
if (label_idx != -100):
if (label_idx >= (col_idx+0)*BLOCK_SIZE) and \
(label_idx < min((col_idx+1)*BLOCK_SIZE, n_cols)):
loss = tl.load(logits_ptr + label_idx).to(tl.float32)
lse = 0.0
loss = lse - logits_label # We add the final logsumexp after a reduction
pass
pass
tl.store(loss_ptr, loss)
pass