From 04fb97aee2bc921d48b302600dfe4928a9b82741 Mon Sep 17 00:00:00 2001
From: anonymous dev <44989609+devchilll@users.noreply.github.com>
Date: Sat, 14 Feb 2026 22:13:07 -0800
Subject: [PATCH] [FIX] Move labels to logits device in cross-entropy loss for
 multi-GPU support (#4041) (#4059)

When using device_map='balanced' with multiple GPUs, the labels tensor
may reside on a different device than the logits/losses tensors. This
causes a RuntimeError at the masked_fill_ call in the chunked
cross-entropy forward path.

Fix: explicitly move labels to the same device as logits at the start
of Fast_CrossEntropyLoss.forward(). This is a no-op on single-GPU
setups.

Fixes #4041
---
 unsloth/kernels/cross_entropy_loss.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index fbb14013ff..cdd5b1cf12 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -301,6 +301,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
         vocab_size: int
         n_rows, vocab_size = logits.shape
         device = logits.device
+        labels = labels.to(device)
 
         div, mod = divmod(vocab_size, MAX_FUSED_SIZE)
         n_chunks: int = div + (mod != 0)