[RuntimeError: Expected all tensors to be on the same device, but found at least two devices] when fine-tuning with peft and device_map=auto
Author: karoaper
Pull request: no
Labels: name: bug

### System Info

  transformers version: 4.52.4
  Platform: linux
  Python version: 3.10.16
  Accelerate version: 1.7.0
  PyTorch version: 2.6.0+cu124
  peft version: 0.15.2
  trl version: 0.18.1

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction


```
from datasets import load_dataset
import copy
import torch

from peft import AutoPeftModelForCausalLM, LoraConfig, get_peft_model, prepare_model_for_kbit_training

from transformers import TrainingArguments
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

from trl import SFTTrainer, SFTConfig

instruct_tune_dataset = load_dataset("mosaicml/instruct-v3")
instruct_tune_dataset = instruct_tune_dataset.filter(lambda x: x["source"] == "dolly_hhrlhf").rename_column('response','completion')

instruct_tune_dataset["train"] = instruct_tune_dataset["train"].select(range(5_000))
instruct_tune_dataset["test"] = instruct_tune_dataset["test"].select(range(200))

nf4_config = BitsAndBytesConfig(
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.1",
    device_map='auto',
    
    quantization_config=nf4_config,
    use_cache=False
)

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64,
    bias="none",
    task_type="CAUSAL_LM",
)

model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, peft_config)


args = SFTConfig(
  output_dir = "mistral_instruct_generation",
  max_steps = 100, # comment out this line if you want to train in epochs
  per_device_train_batch_size = 4,
  warmup_ratio = 0.03,
  logging_steps=10,
  save_strategy="epoch",
  eval_strategy="steps",
  eval_steps=20, # comment out this line if you want to evaluate at the end of each epoch
  learning_rate=2e-4,
  bf16=True,
  lr_scheduler_type='constant',
  max_seq_length=2048,
  packing=True,
  completion_only_loss = True,
  dataset_num_proc=1
)

trainer = MySFTTrainer(
  model=model,
  peft_config=peft_config,
  processing_class=tokenizer,
  args=args,
  train_dataset=instruct_tune_dataset["train"],
  eval_dataset=instruct_tune_dataset["test"]
)

trainer.train()
```

### Expected behavior

Should complete training

Comments:
Encountered the error when I tried to use device_map="auto" with peft fine-tuning of mistral 7b.  I took the code from this tutorial: https://www.digitalocean.com/community/tutorials/mistral-7b-fine-tuning

I received the error: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:7 and cuda:0!

The error is here: 
File ~/anaconda3/envs/myenv/lib/python3.10/site-packages/transformers/loss/loss_utils.py:38, in fixed_cross_entropy(source, target, num_items_in_batch, ignore_index, **kwargs)
     36 loss = nn.functional.cross_entropy(source, target, ignore_index=ignore_index, reduction=reduction)
     37 if reduction == "sum":
---> 38     loss = loss / num_items_in_batch
     39 return loss

Looking at the device map, I see that the lm head is on cuda:7, but num_items_in_batch is on device cuda:0.  The problem seems to originate in function ForCausalLMLoss, where num_items_in_batch (cuda:0) is not switched to the same device as the logits (cuda:2).


The fix is simple.  Add the following 2 lines in `transformers/loss/loss_utils.py` in function ForCausalLMLoss before `loss = fixed_cross_entropy(logits, shift_labels, num_items_in_batch, ignore_index, **kwargs)`

```
if num_items_in_batch is not None: 
     num_items_in_batch = num_items_in_batch.to(logits.device)
```



cc @sunmarc @mekkcyber because there's a lot of trainer and quantization stuff going on here, which I suspect caused the device mismatch. I'm not sure if it's an issue with the script or if we actually do need to update `loss_utils`

This should help understand how to use `num_items_in_batch`: https://github.com/huggingface/transformers/pull/38728
As for your specific issue, if you install main transformers, the issue should be fixed

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.