Fix floating point precision issue for RoPE
Author: butsugiri
Pull request: yes

## What does this PR do?

This PR fixes the issue of floating point precision in `RotaryEmbedding`.
The purpose of this PR is to fix inconsistency between GPT-Neo-X and HF Transformers, which is causing a model performance degradation.

## Issue

In the current implementation of `RotaryEmbedding`, `inv_freq` is first initialized by float32. 
This value is then used for initializing `cos_cached` and `sin_cached` by float32.
As a result, `cos_cached` and `sin_cached` remain float32 even if the model (including inv_freq) uses float16; this is because these two variables are not the target of dtype conversion of `half()` method
Note that there is also a recomputation logic for these two variables, but it is very unlikely to occur https://github.com/huggingface/transformers/blob/f67dac97bdc63874f2288546b3fa87e69d2ea1c8/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L268

However, this implementation seems inconsistent to the one in the [EleutherAI/gpt-neox](https://github.com/EleutherAI/gpt-neox) library.
In their implementation, `cos_cached` and `sin_cached` are almost always recomputed in the [forward method](https://github.com/EleutherAI/gpt-neox/blob/23ad392fdfe0f0a22c986013013209a03b7c28a1/megatron/model/positional_embeddings.py#L51-L63).
Thus, dtype of `cos_cached` and `sin_cached` are always consistent to the dtype of inv_freq.

This inconsistency between two libraries (HF Transformers and GPT-Neo-X) causes the performance degradation of the model converted from gpt-neox.
For example, the perplexity score of the language model on Wikitext corpus is as follows:

- gpt-neo-x w/o conversion: 520.7840
- gpt-neo-x w/ conversion to HF format: 520.9911
- gpt-neo-x w/ conversion to HF format and this PR: 520.7840 

(Sorry that the perplexity value is really bad. I am reporting the performance of model trained on toy data for debugging purpose)

## Solution

I basically followed the previous PR https://github.com/huggingface/transformers/pull/22888 and made a similar fix.

## Possible Side Effect

In the original code, `cos_cashed` and `sin_cashed` are initialized in the model consturctor.
However, I had to move the initialization code to forward method.
Otherwise the library gave me the following error: "cos_vml_cpu" not implemented for 'Half'.
As a result, `torch.jit.trace` might be no longer available.
Since I am not sure what jit.trace is, I don't have any workaround for this.

## Similar Issues

- https://github.com/huggingface/transformers/pull/22888
- https://github.com/EleutherAI/gpt-neox/issues/873

## 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?
  - I would really appreciate it if the reviewers could point out the missing tests.


## Who can review?

@ArthurZucker and @younesbelkada

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

Thank you for the message.
While I appreciate that we have to keep the compatibility with existing models on the hub, my understanding is that all existing models converted from NeoX all have this precision issue.
I would like to explore alternative solutions to address this issue rather than simply closing the pull request. Is there any other approach we can consider to fix the problem?

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.

Hey! If you want to fix the problem without having to close the PR you should be aiming for a full backward compatibility, add tests to make sure that you are fixing the issue in place, and that previous behaviour is not broken.