ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[FIX] Move loss and n_items to logits device in fast_cross_entropy_loss loss for multi-GPU support (#4063)

Browse source 
* bug fix for multi-GPU

* Apply suggestion from @gemini-code-assist[bot]

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
nole69 2026-02-15 01:09:40 -08:00 • committed by GitHub
commit e3c9482cfb
1 changed files with 3 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/kernels/cross_entropy_loss.py
View file

@ -438,6 +438,7 @@ def fast_cross_entropy_loss(
batch, seq_len, d = logits.shape
assert labels.shape == (batch, seq_len)
device = logits.device
loss = Fast_CrossEntropyLoss.apply(
logits.view(batch * seq_len, d),
labels.view(-1),
@ -446,6 +447,8 @@ def fast_cross_entropy_loss(
)
if n_items is None:
n_items = torch.count_nonzero(labels != -100)
if torch.is_tensor(n_items):
n_items = n_items.to(device)
return loss.sum() / n_items