TorchDynamo graph needlessly fragmented for GPTNeoX due to baddbmm type mistake
### System Info

- `transformers` version: 4.32.0.dev0
- Platform: Linux-5.19.0-42-generic-x86_64-with-glibc2.27
- Python version: 3.10.10
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

```py
from transformers import AutoModelForCausalLM
import torch


def debug_backend(gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]):
    print("debug_backend() called with FX graph:")
    gm.graph.print_tabular()
    return gm.forward  # return a python callable

model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-160m")
jitted = torch.compile(model, backend=debug_backend)

jitted(**model.dummy_inputs)
```
The output is too long to fit in a comment, so you'll have to run the code yourself. It features `"debug_backend() called with FX graph:"` being printed several times, each time followed with a fragment of the whole computation graph. This is not expected since NeoX has no data-dependent control flow.

### Expected behavior

The `torch.compile` backend should only be called once, and therefore `"debug_backend() called with FX graph:"` should only appear once, because GPT NeoX does not actually require any data-dependent control flow.

I've already checked that this can be fixed by turning `GPTNeoXAttention.norm_factor` into a Python scalar instead of a tensor. This is actually what `torch.baddbmm` expects for its `alpha` parameter; it's [supposed to be](https://pytorch.org/docs/stable/generated/torch.baddbmm.html) a scalar. But it seems to silently convert tensors into scalars, so this doesn't cause a crash in normal use.
<img width="577" alt="Captura de pantalla 2023-07-19 a la(s) 5 27 42 p m" src="https://github.com/huggingface/transformers/assets/39116809/24274bdb-2599-4ab6-896b-dd77ff98461e">

The exact fix is, in `modeling_gpt_neox.py`, replace lines 103-107 with:
```py
self.norm_factor = self.head_size ** -0.5
```
and replace the `baddbmm` call inside `_attn` with:
```py
attn_scores = torch.baddbmm(
    attn_scores,
    query,
    key.transpose(1, 2),
    beta=1.0,
    alpha=self.norm_factor,
)
```
@fxmarty I think you are more familiar with this topic? If so, could you take a look, thanks!<|||||>Hi @norabelrose, would you like to submit a PR?<|||||>> Hi @norabelrose, would you like to submit a PR?

I already did! 😊 See #24941.