T5ForConditionalGeneration model runs on cuda:0 but not on cuda:1
Author: Kushdesh
Pull request: no

### System Info

transformers : 4.33.2
torch:  2.1.0a0+29c30b1
Python: 3.10.12
GPUs: Two identical NVIDIA GeForce RTX 4090 GPUs 
OS: Ubuntu 22.04.3 LTS


### Who can help?

_No response_

### 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

```python
import torch
from torch import cuda, tensor
from transformers import T5ForConditionalGeneration
device = 'cuda:1' if cuda.is_available() else 'cpu'

model = T5ForConditionalGeneration.from_pretrained("t5-small")
model = model.to(device)

ids = tensor([[ 363, 19,    8,  792,  381,   13, 7634, 7,   58, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).to(device, dtype=torch.long)
mask = tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).to(device, dtype=torch.long)

res = model.generate(input_ids=ids, max_length=50, attention_mask=mask)

```
gives following error:
`RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.`

The code run perfectly if in above code we change 'cuda:1' to 'cuda:0'.

I tried to find where the problem occurring and found that the error happens when program reaches and run following line of file https://github.com/NVIDIA/apex/blob/741bdf50825a97664db08574981962d66436d16a/apex/normalization/fused_layer_norm.py#L69
```python
  output, invvar = fused_layer_norm_cuda.rms_forward_affine(
        input_, ctx.normalized_shape, weight_, ctx.eps)
```


### Expected behavior

The code should run with 'cuda:1'

Comments:
Hey, you should try using accelerate's device map, or change the `"CUDA_VISIBLE_DEVICE"` os env variable! You probably have `apex` and are using the fused layer norm. Doesn't seem to work out of the box. 
cc @SunMarc 

@ArthurZucker  Thanks for your suggestions
@SunMarc 

I tried CUDA_VISIBLE_DEVICES by running 'CUDA_VISIBLE_DEVICES=1 python test.py '
and also tried by adding following lines at the beginning of code
```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] ="1"
```
But it gives same error.

I couldn't understand how to use device map to use 'cuda:1' for entire model. Do I have to map each layer of the model to cuda:1.

I also tried to run the code by adding  line `torch.cuda.set_device('cuda:1')` in the beginning. In this case the code run correctly. But it doesn't look correct way to handle this. 




Hi @Kushdesh , when you use 
```py
import os
os.environ["CUDA_VISIBLE_DEVICES"] ="1"
```
you also need to change the device to `device = 'cuda:0 if cuda.is_available() else 'cpu'` because the gpu 0 is not visible anymore.

To use device_map, you just need to add the device_map arg like that:
`model = T5ForConditionalGeneration.from_pretrained("t5-small", device_map={"":1})`

Let me know if it works now. 

Hi @SunMarc 
Thanks for your response.
changing device to `device = 'cuda:0 if cuda.is_available() else 'cpu'` does works and it uses 2nd GPU.

In following code I setting evironment variable CUDA_VISIBLE_DEVICES to "0,1" and then use cuda:1 and it gives same error. Look like the error is not because of CUDA_VISIBLE_DEVICES setting but use of cuda:1

```Python
import os
os.environ["CUDA_VISIBLE_DEVICES"] ="0,1"

import torch
from torch import cuda, tensor
from transformers import T5ForConditionalGeneration
device = 'cuda:1' if cuda.is_available() else 'cpu'

model = T5ForConditionalGeneration.from_pretrained("t5-small")
model = model.to(device)

ids = tensor([[ 363, 19,    8,  792,  381,   13, 7634, 7,   58, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).to(device, dtype=torch.long)
mask = tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).to(device, dtype=torch.long)

res = model.generate(input_ids=ids, max_length=50, attention_mask=mask)
```

`model = T5ForConditionalGeneration.from_pretrained("t5-small", device_map={"":1})` does the mapping, but it gives same error. I think @ArthurZucker mentioned **accelerate's device map** . How can I use accelerate in this case to use specific devices.


By passing `device_map={"":1}`, you are already using accelerate under the hood. On my setup, it works as expected. Can you check if you have the same behavior with other models ? 

@SunMarc 
I tried a different model and it works on cuda:1. Following is the code
```Python
import torch
from torch import cuda
device = 'cuda:1' if cuda.is_available() else 'cpu'

from transformers import EncoderDecoderModel, AutoTokenizer
sentence_fuser = EncoderDecoderModel.from_pretrained("google/roberta2roberta_L-24_discofuse", device_map={"":1})

tokenizer = AutoTokenizer.from_pretrained("google/roberta2roberta_L-24_discofuse")

input_ids = tokenizer(

    "This is the first sentence. This is the second sentence.", add_special_tokens=False, return_tensors="pt"

).input_ids.to(device, dtype=torch.long)

outputs = sentence_fuser.generate(input_ids)

print(tokenizer.decode(outputs[0]))
```

I tried to find at which line of the code accessing data on GPU produce error while running T5ForConditionalGeneration. Below is output of pdb `where` command and the last line gives position of https://github.com/NVIDIA/apex/blob/741bdf50825a97664db08574981962d66436d16a/apex/normalization/fused_layer_norm.py#L69
Before running this line I can print variables input_, weight_ mentioned on the last line and is on GPU. But executing this line if I tried to print input_ then it gives the error I mentioned in the first post.

```
res = model.generate(input_ids=ids, max_length=50, attention_mask=mask)
  /usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py(115)decorate_context()
-> return func(*args, **kwargs)
  /home/qai/.local/lib/python3.10/site-packages/transformers/generation/utils.py(1490)generate()
-> model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
  /home/qai/.local/lib/python3.10/site-packages/transformers/generation/utils.py(660)_prepare_encoder_decoder_kwargs_for_generation()
-> model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1505)_wrapped_call_impl()
-> return self._call_impl(*args, **kwargs)
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1514)_call_impl()
-> return forward_call(*args, **kwargs)
  /home/qai/.local/lib/python3.10/site-packages/accelerate/hooks.py(165)new_forward()
-> output = old_forward(*args, **kwargs)
  /home/qai/.local/lib/python3.10/site-packages/transformers/models/t5/modeling_t5.py(1123)forward()
-> layer_outputs = layer_module(
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1505)_wrapped_call_impl()
-> return self._call_impl(*args, **kwargs)
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1514)_call_impl()
-> return forward_call(*args, **kwargs)
  /home/qai/.local/lib/python3.10/site-packages/transformers/models/t5/modeling_t5.py(695)forward()
-> self_attention_outputs = self.layer[0](
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1505)_wrapped_call_impl()
-> return self._call_impl(*args, **kwargs)
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1514)_call_impl()
-> return forward_call(*args, **kwargs)
  /home/qai/.local/lib/python3.10/site-packages/transformers/models/t5/modeling_t5.py(601)forward()
-> normed_hidden_states = self.layer_norm(hidden_states)
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1505)_wrapped_call_impl()
-> return self._call_impl(*args, **kwargs)
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py(1514)_call_impl()
-> return forward_call(*args, **kwargs)
  /usr/local/lib/python3.10/dist-packages/apex/normalization/fused_layer_norm.py(386)forward()
-> return fused_rms_norm_affine(input, self.weight, self.normalized_shape, self.eps)
  /usr/local/lib/python3.10/dist-packages/apex/normalization/fused_layer_norm.py(189)fused_rms_norm_affine()
-> return FusedRMSNormAffineFunction.apply(*args)
  /usr/local/lib/python3.10/dist-packages/torch/autograd/function.py(506)apply()
-> return super().apply(*args, **kwargs)  # type: ignore[misc]
> /usr/local/lib/python3.10/dist-packages/apex/normalization/fused_layer_norm.py(70)forward()
-> output, invvar = fused_layer_norm_cuda.rms_forward_affine(
(Pdb) 
```


@Kushdesh thanks for this in-depth investigation ! Looks like it is more of an issue on apex side. Unfortunately, I don't think that we will be able to solve it on our side. In the modeling script of t5t, we are just calling `self.layer_norm(hidden_states)` and since you have apex, it uses `fused_rms_norm_affine`. I suggest you to raise that issue on apex library. LMK if it makes sense. 

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.

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.