RuntimeError(s) when attempting multi-GPU fine-tuning of IDEFICS with naive model parallelism
Author: willemsenbram
Pull request: no

### System Info

- `transformers` version: 4.34.0
- Platform: Linux-5.4.0-148-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.24.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help?

@ArthurZucker @younesbelkada @amyeroberts

### Information

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

### Tasks

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

### Reproduction

Related issue: [#22561](https://github.com/huggingface/transformers/issues/22561)

The problems can be reproduced with [the IDEFICS PEFT demo](https://github.com/huggingface/notebooks/blob/main/examples/idefics/finetune_image_captioning_peft.ipynb).

Using `device_map="auto"` may result in `RuntimeError: Expected all tensors to be on the same device, but found at least two devices`. 
On my end, the error disappears when setting `CUDA_VISIBLE_DEVICES` to be only those that have no other processes already running, but it seems like having other processes use the same GPU is not always an issue: is it possible that the auto mapping can cause an improper split of the model when having to work around limited resource availability on devices with running processes?

With a working device mapping, the next problem is `RuntimeError: indices should be either on cpu or on the same device as the indexed tensor`. 
Relevant part of the traceback:
```
File "/home/user/miniconda3/envs/idefics/lib/python3.10/site-packages/transformers/models/idefics/modeling_idefics.py", line 1516, in forward
  shift_logits = logits[..., :-1, :][shift_attention_mask != 0].contiguous()
```
`logits`, `labels`, and `attention_mask` are expected to be on the same device.
A possible workaround is moving them all to the same device (in [modeling_idefics.py](https://github.com/huggingface/transformers/blob/ce315081340fdf6846f16c321eb53878b6272d53/src/transformers/models/idefics/modeling_idefics.py#L1512-L1515)), like so:
```diff
        if labels is not None:
+           labels = labels.to(logits.device)
            # Shift so that tokens < n predict n
            if attention_mask is not None:
-               shift_attention_mask = attention_mask[..., 1:]
+               shift_attention_mask = attention_mask[..., 1:].to(logits.device) 
```
Can someone confirm that this is proper? If so, I can create a PR.

### Expected behavior

Proper spread of the model over available GPUs when using `device_map="auto"` and have `logits`, `labels`, and `attention_mask` be on the same device to enable training with naive model parallelism.

Comments:
Hi @willemsenbram - thanks for raising this issue! 

The proposed solution sounds good to me - it's consistent with other model implementations. If you open a PR we'll be happy to review :) 

Hi @amyeroberts, thanks for checking! I've opened the [PR](https://github.com/huggingface/transformers/pull/27746)