Repository: transformers
Issue: Tensor size mismatch when trying to run RT-DETR on multiple gpus
Labels: bug

### System Info

- `transformers` version: 4.44.2
- Platform: Linux-5.4.0-174-generic-x86_64-with-glibc2.31
- Python version: 3.11.6
- Huggingface_hub version: 0.24.6
- Safetensors version: 0.4.4
- Accelerate version: 0.33.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.2+cu121 (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: Tesla V100-DGXS-16GB

### Who can help?

@amyeroberts @muellerz @SunMarc 

### 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

Following the example on [the official pytorch example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/object-detection/run_object_detection.py) and [here](https://huggingface.co/docs/transformers/en/tasks/object_detection) it seems that I get the stack trace below after following these steps:

1. Set the initial model class to RT-DETR assuming other parts of the example have been followed 
```
IMAGE_SIZE = 1280
CHECKPOINT = "PekingU/rtdetr_r50vd_coco_o365"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModelForObjectDetection.from_pretrained(
    CHECKPOINT,
    id2label=id2label,
    label2id=label2id,
    anchor_image_size=None,
    ignore_mismatched_sizes=True
)
```

2. Set the batch size to 4 and set to 4 visible GPUs assuming other parts of the example have been followed
```
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

training_args = TrainingArguments(
    output_dir=output_path,
    num_train_epochs=20,
    max_grad_norm=0.1,
    learning_rate=5e-5,
    warmup_steps=300,
    per_device_train_batch_size=4,
    dataloader_num_workers=2,
    metric_for_best_model="eval_map",
    greater_is_better=True,
    load_best_model_at_end=True,
    eval_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    remove_unused_columns=False,
    eval_do_concat_batches=False,
)
```

3. Run training
```
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=pytorch_dataset_train,
    eval_dataset=pytorch_dataset_valid,
    tokenizer=processor,
    data_collator=collate_fn,
    compute_metrics=eval_compute_metrics_fn,
)

trainer.train()
```

Stack trace:
```
RuntimeError                              Traceback (most recent call last)
Cell In[13], [line 11](vscode-notebook-cell:?execution_count=13&line=11)
      [1](vscode-notebook-cell:?execution_count=13&line=1) trainer = Trainer(
      [2](vscode-notebook-cell:?execution_count=13&line=2)     model=model,
      [3](vscode-notebook-cell:?execution_count=13&line=3)     args=training_args,
   (...)
      [8](vscode-notebook-cell:?execution_count=13&line=8)     compute_metrics=eval_compute_metrics_fn,
      [9](vscode-notebook-cell:?execution_count=13&line=9) )
---> [11](vscode-notebook-cell:?execution_count=13&line=11) trainer.train()

File ~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1938, in Trainer.train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   [1936](https://vscode-remote+ssh-002dremote-002bbrain-002ds-002d1-002ehq-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1936)         hf_hub_utils.enable_progress_bars()
   [1937](https://vscode-remote+ssh-002dremote-002brain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1937) else:
-> [1938](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1938)     return inner_training_loop(
   [1939](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1939)         args=args,
   [1940](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1940)         resume_from_checkpoint=resume_from_checkpoint,
   [1941](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1941)         trial=trial,
   [1942](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1942)         ignore_keys_for_eval=ignore_keys_for_eval,
   [1943](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:1943)     )

File ~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:2279, in Trainer._inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   [2276](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:2276)     self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
   [2278](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:2278) with self.accelerator.accumulate(model):
-> [2279](https://vscode-remote+ssh-002dremote-002bbrain-002ecom.vscode-resource.vscode-cdn.net/home/jb/repos/ml/ll_ml/models/rt_detr/~/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/trainer.py:2279)     tr_loss_step = self.training_step(model, inputs)
...
  File "/home/jb/.cache/pypoetry/virtualenvs/ml-Mf12zaqr-py3.11/lib/python3.11/site-packages/transformers/models/rt_detr/modeling_rt_detr.py", line 1850, in forward
    reference_points_unact = torch.concat([denoising_bbox_unact, reference_points_unact], 1)
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 16 but got size 4 for tensor number 1 in the list.
```

It seems that the first tensor is not split amongst the GPUs? Is there something else I'm missing here? I have not set up `accelerate` for instance, but the example linked does not state that it is required, and also states that at least one GPU is needed, making no specification about setting up `accelerate`.

### Expected behavior

I would expect to see training taking place as it does when I use just one GPU:

![image](https://github.com/user-attachments/assets/9bc3c5b9-b486-4223-8b16-dec9edce95dd)