Repository: transformers
Issue: Qwen2-VL: Multi-GPU training
Labels: Distributed Training / Models,trainer,Feature request,bug,Vision,Multimodal

### System Info

- `transformers` version: 4.45.0.dev0
- Platform: Linux-4.18.0-477.10.1.el8_8.x86_64-x86_64-with-glibc2.28
- Python version: 3.11.5
- Huggingface_hub version: 0.24.0
- Safetensors version: 0.4.3
- Accelerate version: 0.34.2
- Accelerate config: 	- compute_environment: LOCAL_MACHINE
	- distributed_type: NO
	- mixed_precision: bf16
	- use_cpu: False
	- debug: False
	- num_processes: 1
	- machine_rank: 0
	- num_machines: 1
	- gpu_ids: all
	- rdzv_backend: static
	- same_network: True
	- main_training_function: main
	- enable_cpu_affinity: False
	- downcast_bf16: no
	- tpu_use_cluster: False
	- tpu_use_sudo: False
	- tpu_env: []
- PyTorch version (GPU?): 2.2.1+rocm5.7 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: AMD Instinct MI250X


### Who can help?

@muellerzr @ArthurZucker @gante 

Issue about both the Qwen-VL model and perhaps the trainer so not sure who is best suited to answer :) 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Replicating the setup is a bit tough, so this is more of a preliminary discussion issue to see if there is an obvious problem that surfaces.

1. Multi-GPU setup  + Huggingface trainer
2. Train Qwen2-VL model with dynamic image resolution
3. The processor creates BatchEncodings with pixel_values, input_ids, attention_mask and image_grid_thw.
4. Run a model forward pass with the model in data parallel mode of the trainer.

We observe that compared to mono-gpu setups, the rope values are disaligned with the hidden_states size.

Typically, in line 1109 (Qwen2VisionTransformerPretrainedModel forward pass):
```python
    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
        hidden_states = self.patch_embed(hidden_states)
        rotary_pos_emb = self.rot_pos_emb(grid_thw)
```
we can see rotary_pos_emb is hidden_states have a sligtly  different dimension 0. 
ex: torch.Size([7820, 40]) torch.Size([7736, 1280])

Upon further inspection, we see rotary_pos_emb has the same dimension as what we would get in mono-gpu runs (normal since it only depends on the grid_thw argument). However, hidden_states (that correspond to pixel values) have a different  size.

This makes training crash:

```bash
  File "/lus/home/CT10/cad15443/mfaysse/colpali/venv/lib/python3.11/site-packages/transformers/models/qwen2_vl/modeling_qwen2_vl.py", line 395, in forward
    q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lus/home/CT10/cad15443/mfaysse/colpali/venv/lib/python3.11/site-packages/transformers/models/qwen2_vl/modeling_qwen2_vl.py", line 254, in apply_rotary_pos_emb_vision
    output = (tensor * cos) + (rotate_half(tensor) * sin)
              ~~~~~~~^~~~~
RuntimeError: The size of tensor a (7736) must match the size of tensor b (7808) at non-singleton dimension 1
```

### Expected behavior

[edited] see below for more details being investigated

Thanks !