Bugfix: LLaMA layer norm incorrectly changes input type and consumers lots of memory
# What does this PR do?

It is a common setup to run LLaMA in `bfloat16` or `float16` while the `RMSNorm` layers are in `float32`. The current implementation converts an input into `float32` but never converts it back to the original input type. This means in the first layer the input type if `bfloat16` but for all remaining its `float32` this increases the overall memory-footprint of finetuning considerably.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@younesbelkada @sgugger 

Please also see: https://github.com/huggingface/transformers/pull/23479


_The documentation is not available anymore as the PR was closed or merged._<|||||>> but for all remaining its `float32` this increases the overall memory-footprint of finetuning considerably.

...how bad was this? like 30% worse?<|||||>Please merge ^_^<|||||>I am assuming this does not impact 8bit training as I noticed no change in memory.<|||||>@official-elinas did an experiment with transformers==4.28.1 vs the code added here, and was unable to reproduce any speed/memory gains

https://wandb.ai/officialelinas/tests?workspace=user-officialelinas

![image](https://github.com/huggingface/transformers/assets/54623771/1cd5af7d-96ba-4217-a0e4-12d5cea4ef5c)
![image](https://github.com/huggingface/transformers/assets/54623771/4aaa93e4-5cc5-45e1-9cb5-7c3e724734a1)

@TimDettmers what do you think we did wrong? are there any experiments you can share regarding this?
<|||||>I think unless you've manually set RMSNorm layer's parameters to fp32 while the rest of your model is in fp16/bf16 you won't see any change.

I'm curious if as a design we should instead allow to specify `dtype` in the layer arguments? `F.softmax` has a dtype argument that allows one to control such a thing more finely? https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html