Prevent Dynamo graph fragmentation in GPTNeoX with torch.baddbmm fix
# What does this PR do?

Fixes #24940.


## 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?

@ArthurZucker and @younesbelkada

I'm struggling to see what's wrong with my code style since when I run `black` locally it says it's fine.<|||||>> Hey! Thanks for submitting a PR! You should be using `make style` to properly format the code.

This doesn't seem to work for me. `make style` makes a ton of changes to _other_ files that I didn't modify, but doesn't change my code at all.<|||||>> My only concern here is that the dtype is not set. A small test like this shows that this will have some effect for some head size

I suppose I can manually cast `norm_factor` to the appropriate dtype and then convert back to a Python scalar. On the other hand, inverse square root isn't going to have _any_ numerical error as long as the head size is a power of 2, which it almost always is (in fact, it's almost always 64). So in practice this shouldn't matter at all.

Also, presumably the ground truth should be the Eleuther `gpt-neox` implementation. That implementation seems to already encode `norm_factor` as a Python scalar and uses `math.sqrt`, see [here](https://github.com/EleutherAI/gpt-neox/blob/408e29d9c746a02d842917bb7447c5c4be0b42d4/megatron/model/transformer.py#L298). It passes a Python scalar to `torch.baddbmm` [here](https://github.com/EleutherAI/gpt-neox/blob/408e29d9c746a02d842917bb7447c5c4be0b42d4/megatron/model/transformer.py#L420).

So in fact, our current implementation is already "wrong" with respect to the original, and this PR would correct the discrepancy. Although, as I said, in almost all cases this will make no difference.<|||||>Ok! Got what you mean. Part of the code was added in #22888 in order to have proper float16 casting. Yes we should match the original results (which is currently the case) but we have also other functionalities on which a lot of users rely! 
Let's make sure that the value is casted to the correct type for the next computations. 

(make style is probably behaving wrongly because of the black / ruff versioning) <|||||>> Ok! Got what you mean. Part of the code was added in #22888 in order to have proper float16 casting. Yes we should match the original results (which is currently the case) but we have also other functionalities on which a lot of users rely! Let's make sure that the value is casted to the correct type for the next computations.
> 
> (make style is probably behaving wrongly because of the black / ruff versioning)

Hmm okay, it looks like prior to #22888 we were just using `float32` all the time which is clearly wrong. This PR shouldn't cause any regressions because the Python scalar will work with any parameter dtype.

Is there anything I can do to get the codestyle check to pass? That seems to be the only thing preventing merging right now.<|||||>The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_24941). All of your documentation changes will be reflected on that endpoint.<|||||>Sorry but no, in transformers we autocast the weights to the required dtype. You can try the following:
```python 
>>> from transformers import GPTNeoXModel
>>> import torch 

>>> model = GPTNeoXModel.from_pretrained("EleutherAI/gpt-neox-20b", torch_dtype = torch.float16)
>>> print(model.layers[1].attention.norm_factor)
tensor(9.7969, dtype=torch.float16)
```
So the dtype was properly handle. We initialized it with float32 but then it can be casted to any dtype. 

For the red ci run `make style`! 