fix: GPTNeoX half inference error
Author: SeongBeomLEE
Pull request: yes

norm_factor is still torch.float32 after using model.half

So I changed it to register_buffer so I can change it to torch.float16 after using model.half

This error does not occur in all cases, but it does happen occasionally. 

Thanks!

Error Message:
```
File "/data2/sblee/anaconda3/lib/python3.9/site-packages/transformers/models/gpt_neox/modeling_gpt_neox.py", line 206, in _attn
attn_scores = torch.baddbmm(
RuntimeError: expected scalar type Half but found Float
```

Error Code:
```
attn_scores = torch.baddbmm(
            attn_scores,
            query,
            key.transpose(1, 2),
            beta=1.0,
            alpha=(torch.tensor(1.0, dtype=self.norm_factor.dtype, device=self.norm_factor.device) / self.norm_factor),
        )
```

before Code:
```
class GPTNeoXAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_attention_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.head_size = self.hidden_size // self.num_attention_heads
        self.rotary_ndims = int(self.head_size * config.rotary_pct)
        max_positions = config.max_position_embeddings
        self.register_buffer(
            "bias",
            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
                1, 1, max_positions, max_positions
            ),
        )
        self.register_buffer("masked_bias", torch.tensor(-1e9))
        self.rotary_emb = RotaryEmbedding(
            self.rotary_ndims, config.max_position_embeddings, base=config.rotary_emb_base
        )
        self.norm_factor = torch.sqrt(torch.tensor(self.head_size, dtype=torch.float32)).to(torch.get_default_dtype())
        self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size)
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
```

after Code:
```
class GPTNeoXAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_attention_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.head_size = self.hidden_size // self.num_attention_heads
        self.rotary_ndims = int(self.head_size * config.rotary_pct)
        max_positions = config.max_position_embeddings
        self.register_buffer(
            "bias",
            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
                1, 1, max_positions, max_positions
            ),
        )
        self.register_buffer("masked_bias", torch.tensor(-1e9))
        self.rotary_emb = RotaryEmbedding(
            self.rotary_ndims, config.max_position_embeddings, base=config.rotary_emb_base
        )
        self.register_buffer("norm_factor", torch.sqrt(torch.tensor(self.head_size, dtype=torch.float32)).to(torch.get_default_dtype()))
        self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size)
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)

```

before:
```
model = GPTNeoXForCausalLM.from_pretrained(F_MODEL_PATH, config=model_config)
model.half()
model.to("cuda")
model.gpt_neox.layers[0].attention.norm_factor.dtype

output: torch.float32
```

after:
```
model = GPTNeoXForCausalLM.from_pretrained(F_MODEL_PATH, config=model_config)
model.half()
model.to("cuda")
model.gpt_neox.layers[0].attention.norm_factor.dtype

output: torch.float16
```

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Hi! @younesbelkada 

It works fine when executed in the way you described.

However, what do you think about considering the case of using model.half?

Thanks!

before:
```
model = GPTNeoXForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float16)
model.gpt_neox.layers[0].attention.norm_factor.dtype

output: torch.float16
```
after:
```
model = GPTNeoXForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float16)
model.gpt_neox.layers[0].attention.norm_factor.dtype

output: torch.float16
````

The canonical way is indeed to use `torch_dtype=torch.float16` as it saves a lot of memory (otherwise you instantiate your model in float32 so take a lot of space, then convert it to half). But `model.half()` should work nethertheless. Using a buffer seems like a good solution, but make it non-persistent so it doesn't get inside the `state_dict` (with `persistent=False`).

Added a parameter "persistent=False"

Thanks!

Of course. Thanks!

Failure of the FLax test is unrelated and already fixed on main, so merging!