Fix premature downcast in LlamaRMSNorm
Author: Birch-san
Pull request: yes

# What does this PR do?

Fixes dtype mismatch encountered after LlamaRMSNorm during full-finetuning of Llama models:

```
Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/mahouko/anaconda3/envs/p311-qlora/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py", line 85, in _worker
    output = module(*input, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^
(this was just a WrappedException bubbling up to parallel apply. the root cause is below)
…

  File "/home/mahouko/anaconda3/envs/p311-qlora/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16          
```

I'll backtrack a bit and explain how I got here.

At the first layernorm in the Llama model:

- we start with bfloat16 hidden states
- we put the bfloat16 hidden states into a float32 LayerNorm
- the LayerNorm upcasts the hidden states to float32 (this is sensible)
- the LayerNorm **tries** to downcast the result back to original bfloat16, but **fails** due to a typo. it downcasts the hidden states operand, instead of the result.
- hence, bfloat16 hidden states went into the LayerNorm, but float32 hidden states came out

the float32 hidden states continue onward and enter our bfloat16 `q_proj`:  
![image](https://github.com/huggingface/transformers/assets/6141784/a599d04c-8228-4b06-92f8-d6d844174f69)

```
hidden_states.dtype
torch.float32
self.q_proj.weight.dtype
torch.bfloat16
```

_That's_ when the `float != c10::BFloat16` error is thrown.

Fixing this typo will **also** improve the precision of the LayerNorm calculation (we avoid downcasting the hidden states prematurely, so they remain full-precision for the Hadamard product).

I'm not the only person who encountered this problem. Other people are doing this too:  
https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/13/files

## Before submitting
- [x] 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?

@ArthurZucke, @younesbelkada

Comments:
Well, the hidden states have to be casted to the input hidden state's dtype following [the original code](https://github.com/facebookresearch/llama/blob/main/llama/model.py#L34).
I would think that if `self.weight` is in `bfloat16` then autocast would resolve it to `bfloat16`

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

Could you provide a small reproducer? I am having a hard time understanding why this is still discussed? 
If the model's weights have the proper type (the same as the input) then the output would be `bfloat16` no? Others have successfully trained with `bloat16` as well as `float16`. 

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.