APEX: use MixedFusedRMSNorm instead of FusedRMSNorm for numerical consistency
Author: fxmarty
Pull request: yes

As per title.

APEX `FusedRMSNorm` initialize the returned tensor to the `input` dtype, which raises an error in case the model is set on fp16, where we may have `fp32` input to the layer norm layer:
```
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/apex-0.1-py3.8-linux-x86_64.egg/apex/normalization/fused_layer_norm.py", line 189, in fused_rms_norm_affine
    return FusedRMSNormAffineFunction.apply(*args)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/apex-0.1-py3.8-linux-x86_64.egg/apex/normalization/fused_layer_norm.py", line 69, in forward
    output, invvar = fused_layer_norm_cuda.rms_forward_affine(
```

Comparing to [T5LayerNorm](https://github.com/huggingface/transformers/blob/820c46a707ddd033975bc3b0549eea200e64c7da/src/transformers/models/t5/modeling_t5.py#L247-L261) where the output is on the weight dtype. That is exactly what `MixedFusedRMSNorm` is for, see https://github.com/NVIDIA/apex/blob/52e18c894223800cb611682dce27d88050edf1de/apex/normalization/fused_layer_norm.py#L420 and https://github.com/NVIDIA/apex/blob/52e18c894223800cb611682dce27d88050edf1de/csrc/layer_norm_cuda.cpp#L205

For example, the test `pytest tests/test_pipeline_mixin.py::VisualQuestionAnsweringPipelineTests::test_small_model_pt_blip2 -s -vvvvv` fails when `apex` is available (and accelerate is installed). This error was never detected because the docker images used for testing do not have APEX installed (e.g. `nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04`)

This is an issue for the AMD CI as the image used `rocm/pytorch:rocm5.6_ubuntu20.04_py3.8_pytorch_2.0.1` has RoCm APEX installed by default.

---
Slightly out of topic: something I don't get is why [neither t5x](https://github.com/google-research/t5x/blob/ea66ec835a5b413ca9d211de96aa899900a84c13/t5x/examples/t5/layers.py#L445) nor transformers seem to recast to fp16 after the FFN. Due to the `keep_in_fp32` attribute, [this weight](https://github.com/huggingface/transformers/blob/820c46a707ddd033975bc3b0549eea200e64c7da/src/transformers/models/t5/modeling_t5.py#L284C14-L284C14) is always in fp32 and then fp32 is propagated in the model. t5x seem to do the same.

Related: https://github.com/huggingface/transformers/pull/26225

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26226). All of your documentation changes will be reflected on that endpoint.

Good for me, but let's wait for @mfuntowicz 's response.

But why we can't just uninstall apex on the AMD docker image? The goal here is not to test the functionality of APEX.

> why we can't just uninstall apex on the AMD docker image

IMO it makes sense to test both paths, the APEX path is broken as well on NVIDIA GPUs. APEX is installed by default in the PyTorch docker image provided by AMD, and it makes sense to expect Transformers to work with the image.

Those 3rd party libraries cause problems quite frequently especially when they change versions or torch get a new version.

There is `docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile` used by the job `run_all_tests_torch_cuda_extensions_gpu` - this has apex installed, but this is a separate job, in particular for DeepSpeed.

Let's not spend too much time but just make the **usual** model and common tests run on AMD CI for now.

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.