Prevent Dynamo graph fragmentation in GPTNeoX with torch.baddbmm fix
Author: norabelrose
Pull request: yes

# 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


Comments:
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`! 

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.

> Sorry but no, in transformers we autocast the weights to the required dtype.

But this isn't a _weight_, though. It's a constant value deterministically computed from the config.

> For the red ci run `make style`!

I'm sorry, but this simply does not work! It changes over 200 files that I haven't changed, but doesn't do anything to the files I did change. Please help me here, can you just clone my fork and somehow fix it or figure out what's wrong?

Just pushed the changes! 
You are right it is not a weight, but it was 😓 
Which mean that this fix would be breaking previous outputs. Since the model is pretty old, would like to just cast the norm factor to the correct dtype (since I am not sure we can be 100% sure all previous outputs will be the sameWDYT?) 

Thanks a lot for that!

I'm not sure what you mean by "cast the norm factor to the correct dtype." At the very least, we must pass a Python scalar to `torch.baddbmm` in order to fix the Dynamo graph fragmentation issue. Python scalars will work with any model dtype.

If you're saying we cast it to the model's dtype and then call `.item()` or something— I _suppose_ we could do that, but I would rather not. Right now there is _technically_ a discrepancy between our implementation and the original NeoX implementation. If we merge this PR as-is, the discrepancy would be fixed.

That said, this discrepancy will virtually never matter since `head_size` is always a power of 2 in all the pretrained checkpoints. They will behave identically.

Thanks @norabelrose for bearing with me and merging this!