[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

@@ -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)