MultiScaleDeformableAttentionFunction different results on different devices
Author: DonggeunYu
Pull request: no

### System Info

- `transformers` version: 4.39.0
- Platform: Linux-5.4.0-81-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.2
- Accelerate version: 0.27.2
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.0+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 GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@amyeroberts 

### 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

Command: CUDA_VISIBLE_DEVIECS=0 TORCH_USE_CUDA_DSA=1 python3 run.py
~~~
import torch
from transformers.models.deformable_detr.modeling_deformable_detr import load_cuda_kernels, MultiScaleDeformableAttentionFunction

load_cuda_kernels()
value = torch.ones([1, 21760, 8, 32])
spatial_shapes = torch.tensor([[128, 128],
    [64, 64], [32, 32], [16, 16]])
level_start_index = torch.tensor([0, 16384, 20480, 21504])
sampling_locations = torch.ones([1, 300, 8, 4, 4, 2])
attention_weights = torch.full([1, 300, 8, 4, 4], 0.0625)
im2col_step = 64

device = "cuda:0"
result = MultiScaleDeformableAttentionFunction.apply(
                    value.to(device),
                    spatial_shapes.to(device),
                    level_start_index.to(device),
                    sampling_locations.to(device),
                    attention_weights.to(device),
                    im2col_step
                )
torch.cuda.empty_cache()
print(result)
~~~
Output:
~~~
tensor([[[0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         ...,
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500]]],
       device='cuda:0')
~~~

---

Command: CUDA_VISIBLE_DEVIECS=1 TORCH_USE_CUDA_DSA=1 python3 run.py
~~~
import torch
from transformers.models.deformable_detr.modeling_deformable_detr import load_cuda_kernels, MultiScaleDeformableAttentionFunction

load_cuda_kernels()
value = torch.ones([1, 21760, 8, 32])
spatial_shapes = torch.tensor([[128, 128],
    [64, 64], [32, 32], [16, 16]])
level_start_index = torch.tensor([0, 16384, 20480, 21504])
sampling_locations = torch.ones([1, 300, 8, 4, 4, 2])
attention_weights = torch.full([1, 300, 8, 4, 4], 0.0625)
im2col_step = 64

device = "cuda:1"
result = MultiScaleDeformableAttentionFunction.apply(
                    value.to(device),
                    spatial_shapes.to(device),
                    level_start_index.to(device),
                    sampling_locations.to(device),
                    attention_weights.to(device),
                    im2col_step
                )
torch.cuda.empty_cache()
print(result)
~~~
Output:
~~~
tensor([[[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]]], device='cuda:1')
~~~

### Expected behavior

Results are different for GPU devices 0 and 1.
This issue affects our private model.

Comments:
cc @qubvel If you have time to dig into this 

Hi @DonggeunYu thanks for reporting the issue!
Unfortunately, I was not able to reproduce it with my envs. I tried:
 - latest torch (2.3.0+cu121) + latest transformers (4.41.2)
 - specified torch (2.1.0+cu121) + latest transformers (4.41.2)
 - latest torch (2.3.0+cu121) + specified transformers (4.39.0)
 - **specified torch (2.1.0+cu121) + specified transformers (4.39.0)**

My setup is 4 GPUs Tesla T4, I tried to launch on each of them, results were always identical
```
tensor([[[0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         ...,
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500]]],
       device='cuda:0')
```

Env:
```
- `transformers` version: 4.39.0
- Platform: Linux-6.5.0-1020-aws-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.1.0+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 GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
```

Do you have any ideas why that might happen in your env?

> Hi @DonggeunYu thanks for reporting the issue! Unfortunately, I was not able to reproduce it with my envs. I tried:
> 
> * latest torch (2.3.0+cu121) + latest transformers (4.41.2)
> * specified torch (2.1.0+cu121) + latest transformers (4.41.2)
> * latest torch (2.3.0+cu121) + specified transformers (4.39.0)
> * **specified torch (2.1.0+cu121) + specified transformers (4.39.0)**
> 
> My setup is 4 GPUs Tesla T4, I tried to launch on each of them, results were always identical
> 
> ```
> tensor([[[0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          ...,
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500]]],
>        device='cuda:0')
> ```
> 
> Env:
> 
> ```
> - `transformers` version: 4.39.0
> - Platform: Linux-6.5.0-1020-aws-x86_64-with-glibc2.35
> - Python version: 3.10.12
> - Huggingface_hub version: 0.23.4
> - Safetensors version: 0.4.3
> - Accelerate version: not installed
> - Accelerate config: not found
> - PyTorch version (GPU?): 2.1.0+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 GPU in script?: <fill in>
> - Using distributed or parallel set-up in script?: <fill in>
> ```
> 
> Do you have any ideas why that might happen in your env?

I don't know the cause.
I will also test it in various environments.

> Hi @DonggeunYu thanks for reporting the issue! Unfortunately, I was not able to reproduce it with my envs. I tried:
> 
> * latest torch (2.3.0+cu121) + latest transformers (4.41.2)
> * specified torch (2.1.0+cu121) + latest transformers (4.41.2)
> * latest torch (2.3.0+cu121) + specified transformers (4.39.0)
> * **specified torch (2.1.0+cu121) + specified transformers (4.39.0)**
> 
> My setup is 4 GPUs Tesla T4, I tried to launch on each of them, results were always identical
> 
> ```
> tensor([[[0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          ...,
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500],
>          [0.2500, 0.2500, 0.2500,  ..., 0.2500, 0.2500, 0.2500]]],
>        device='cuda:0')
> ```
> 
> Env:
> 
> ```
> - `transformers` version: 4.39.0
> - Platform: Linux-6.5.0-1020-aws-x86_64-with-glibc2.35
> - Python version: 3.10.12
> - Huggingface_hub version: 0.23.4
> - Safetensors version: 0.4.3
> - Accelerate version: not installed
> - Accelerate config: not found
> - PyTorch version (GPU?): 2.1.0+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 GPU in script?: <fill in>
> - Using distributed or parallel set-up in script?: <fill in>
> ```
> 
> Do you have any ideas why that might happen in your env?

@qubvel 
If the container image you used is public, can you share it?

@DonggeunYu I was using an Amazon EC2 instance `g4dn.12xlarge` with Ubuntu 22.04

I tested it on V100 and A100, but on devices other than 0, it comes out as 0.

used image: nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
command: `apt update && apt install python3 python3-pip -y && pip3 install torch ninja transformers && TORCH_CUDA_ARCH_LIST='5.0 5.2 6.0 6.1 7.0 7.5 8.0 8.6 8.9+PTX 9.0+PTX' python3 z.py`

@DonggeunYu thanks, I will try to reproduce with your environment!

If you can't reproduce, please let me know right away.

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.