Cuda error when using T5 with the new Apex FusedRMSNorm 
Author: dblakely
Pull request: no

Hi there, thanks for the work integrating the new Apex kernel to accelerate T5! Unfortunately, however, I'm having some trouble getting it to run. I installed the Apex library following the instructions on their repo and am trying to run a small test script on an RTX A6000 with cuda 11.1.

## Environment info

- `transformers` version: 4.17.0.dev0
- Platform: Linux-4.15.0-135-generic-x86_64-with-glibc2.27
- Python version: 3.8.12
- PyTorch version (GPU?): 1.9.0+cu111 (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 (RTX A6000)
- Using distributed or parallel set-up in script?: no


### Who can help
@stas00 


## Information

Model I am using (Bert, XLNet ...): T5ForConditionalGeneration

## To reproduce

I get the error on my environment by running this script (named `debug.py`):

```python
import transformers
from transformers import T5ForConditionalGeneration, T5Tokenizer

transformers.logging.set_verbosity_info()

device = "cuda"
model_name = "t5-small"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
input_ids = tokenizer(
    "translate English to French: All grown-ups were once children—although few of them remember it.",
    return_tensors="pt",
).input_ids.to(device)

outputs = model.generate(input_ids, max_length=100)
outputs = tokenizer.decode(outputs[0])

print(f"{outputs = }")
```

and I get this stack trace when I run: `CUDA_LAUNCH_BLOCKING=1 python debug.py`:

```
Traceback (most recent call last):
  File "temp.py", line 17, in <module>
    outputs = model.generate(input_ids, max_length=100)
  File "/home/derrick/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/home/derrick/projects/transformers/src/transformers/generation_utils.py", line 1102, in generate
    model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
  File "/home/derrick/projects/transformers/src/transformers/generation_utils.py", line 510, in _prepare_encoder_decoder_kwargs_for_generation
    model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)
  File "/home/derrick/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/derrick/projects/transformers/src/transformers/models/t5/modeling_t5.py", line 1030, in forward
    layer_outputs = layer_module(
  File "/home/derrick/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/derrick/projects/transformers/src/transformers/models/t5/modeling_t5.py", line 665, in forward
    self_attention_outputs = self.layer[0](
  File "/home/derrick/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/derrick/projects/transformers/src/transformers/models/t5/modeling_t5.py", line 572, in forward
    attention_output = self.SelfAttention(
  File "/home/derrick/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/derrick/projects/transformers/src/transformers/models/t5/modeling_t5.py", line 498, in forward
    query_states = shape(self.q(hidden_states))  # (batch_size, n_heads, seq_length, dim_per_head)
  File "/home/derrick/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/derrick/env/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 96, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/derrick/env/lib/python3.8/site-packages/torch/nn/functional.py", line 1847, in linear
    return torch._C._nn.linear(input, weight, bias)
RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasCreate(handle)`
```

Any ideas about what the issue might be?

Comments:
thank you for the report, @dblakely 

I tested your code and it works on my rtx-3090 / pt-1.10

I wonder if you need to set device, with the older pytorch? Does this work?

```
import transformers
import os
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer

transformers.logging.set_verbosity_info()

local_rank = int(os.environ.get("LOCAL_RANK", 0))
torch.cuda.set_device(local_rank)

device = "cuda"
model_name = "t5-small"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
input_ids = tokenizer(
    "translate English to French: All grown-ups were once children—although few of them remember it.",
    return_tensors="pt",
).input_ids.to(device)

outputs = model.generate(input_ids, max_length=100)
outputs = tokenizer.decode(outputs[0])

print(f"{outputs = }")

```

Alternatively, could you please test with pt-1.10? Perhaps it only works with the newer version

cc: @eqy

Thanks! I did upgrade torch to 1.10.2+cu111 and ran your script and it worked. However, now I'm realizing that how I specify the device seems to be determining whether or not I get the error I mentioned above. I'm using a machine with 4 GPUs and if I just say `device = "cuda"`, it works. But if I say `device = "cuda:3"`, it doesn't work.

On the other hand, if I put in my script `device="cuda"` and then run it using `CUDA_VISIBLE_DEVICES=3 python debug.py`, it works (indicating it's not an issue with that GPU per se). So at the end of the day, I am able to get the Apex kernel to run. However, it's being kind of finicky with respect to device ordinals in a way that T5 normally is not. Any ideas about why that would be the case? If not, it's fine because I believe the issue is solved.

does the following make any difference?

```
local_rank = int(os.environ.get("LOCAL_RANK", 0))
device = torch.device(f'cuda:{local_rank}')
torch.cuda.set_device(device)
```

so did your original script w/o my changes work with pt-1.10? i.e. w/o setting the device explicitly - trying to understand if the initial issue lays with pt-1.9 - there could be multiple issues as well.

and you never shared how you launch this - sorry edit - you did share that info - not distributed. oh, I get it, but `CUDA_VISIBLE_DEVICES=3` still is `cuda:0`

so perhaps tell me what you're trying to do. running on one gpu out of 4, or running in parallel with DDP, or else?

- `python debug.py` will always run on gpu 0 if the code does `device = "cuda"`
- `CUDA_VISIBLE_DEVICES=3 python debug.py` is still running on gpu 0, software-wise for `device = "cuda"`

`CUDA_VISIBLE_DEVICES` simply narrows down or reorders the gpus visible to the program, so unless you use DDP `device("cuda") == device("cuda:0")` or just `tensor.cuda()`

see: https://huggingface.co/docs/transformers/master/main_classes/trainer#specific-gpus-selection

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.