T5 int8 inference is not compatible with nvidia/apex
Author: lukaemon
Pull request: no

### System Info

- `transformers` version: 4.26.1
- Platform: Linux-5.15.0-60-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.12.0
- PyTorch version (GPU?): 1.14.0a0+44dac51 (True)
- Tensorflow version (GPU?): 2.11.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.6.4 (gpu)
- Jax version: 0.4.3
- JaxLib version: 0.4.3
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>


### Who can help?

@ArthurZucker and @younesbelkada

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Copy paste int8 inference code from flan-t5-xxl page:
```python
# pip install bitsandbytes accelerate
from transformers import T5Tokenizer, T5ForConditionalGeneration

tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-xxl")
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-xxl", device_map="auto", load_in_8bit=True)

input_text = "translate English to German: How old are you?"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")

outputs = model.generate(input_ids)
print(tokenizer.decode(outputs[0]))
```

Error message:
```
│ /usr/local/lib/python3.8/dist-packages/apex/normalization/fused_layer_norm.py:69 in forward      │
│                                                                                                  │
│    66 │   │   ctx.eps = eps                                                                      │
│    67 │   │   input_ = input.contiguous()                                                        │
│    68 │   │   weight_ = weight.contiguous()                                                      │
│ ❱  69 │   │   output, invvar = fused_layer_norm_cuda.rms_forward_affine(                         │
│    70 │   │   │   input_, ctx.normalized_shape, weight_, ctx.eps)                                │
│    71 │   │   ctx.save_for_backward(input_, weight_, invvar)                                     │
│    72 │   │   return output                                                                      │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: expected scalar type Float but found Half
```

If I understand correctly, apex is triggered automatically, ref T5 doc:
> If you’d like a faster training and inference performance, install [apex](https://github.com/NVIDIA/apex#quick-start) and then the model will automatically use apex.normalization.FusedRMSNorm instead of T5LayerNorm. The former uses an optimized fused kernel which is several times faster than the latter.

Any work around?
Should I turn off apex during `load_in_8bit` to use default layernorm? How? 

### Expected behavior

Want to work with int8 inference. Is it a bug or what did I miss? Thanks.

Ideally still want to keep apex during training. And turn it off during int8 inference.
Remove it definitely works, but -10% tflops at training from a quick benchmark. 

Comments:
Hi @lukaemon 
Yes this is a known issue, currently int8 and apex are not supported together, the fix I can propose now is to disable apex by uninstalling it until we found a proper fix!
Thanks a lot

Will do. Thanks.