Add attention upcast for GPT-NeoX
Author: wilswu99
Pull request: yes

# What does this PR do?

This PR adds a flag to upcast the QK.T matmul step in the attention layer, for better numerical stability when using fp16. 
This is the same as the upcast added to GPT-2 in https://github.com/huggingface/transformers/pull/13573

Without this step, smaller models often overflow during the softmax. For example:
```
train = load_dataset('EleutherAI/pile-deduped-pythia-random-sampled', split='train')
train = train.cast_column('Tokens', datasets.Sequence(datasets.Value('int64')))
train = train.with_format('torch')
train_loader = DataLoader(train, batch_size=32)
x = next(iter(train_loader))['Tokens'].to('cuda')

model = AutoModelForCausalLM.from_pretrained(f'EleutherAI/pythia-14m').to('cuda')
with torch.autocast(enabled=True, device_type='cuda', dtype=torch.float16):
    out = model(input_ids=x, labels=x)
print(out.loss.item()) # nan

config = AutoConfig.from_pretrained('EleutherAI/pythia-14m')
config.upcast_attn = True
model = AutoModelForCausalLM.from_pretrained('EleutherAI/pythia-14m', config=config).to('cuda')
with torch.autocast(enabled=True, device_type='cuda', dtype=torch.float16):
    out = model(input_ids=x, labels=x)
print(out.loss.item()) # 4.708
```

## 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#create-a-pull-request),
      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?
@younesbelkada @zphang 

Comments:
Hi @wiwu2390, thanks for opening this PR! 

We try to avoid adding arguments like this that lead to more complex forward passes e.g. with if/else structures. I'm also not sure that adding this autocast logic is compatible with the rest of the library, as there's assumptions about what part of the library handles things like mixed precision and intialization which it might break. 

cc @ArthurZucker 

@amyeroberts Thanks for looking at this!

It seems that there shouldn't be any compatibility problems, because this change is just mimicking what is already being done in GPT-2's `_upcast_and_reordered_attn` method (first mentioned in https://github.com/huggingface/transformers/issues/13463 and added in https://github.com/huggingface/transformers/pull/13573). I don't think GPT-NeoX has any major differences in mixed-precision handling that would make it incompatible with doing this same thing, but please let me know if I'm mistaken.

The GPT-2 implementation has the if/else logic in `forward`, and switches between `_attn` or `_upcast_and_reorderd_attn` based on the `config.reorder_and_upcast_attn` flag (see below). If that's the preferred style, I can do the same thing here, though it would result in more duplicate code.

https://github.com/huggingface/transformers/blob/09f9f566de83eef1f13ee83b5a1bbeebde5c80c1/src/transformers/models/gpt2/modeling_gpt2.py#L328-L331

In any case, I think this would be a useful change to make -- without it, I don't know of a way to train or finetune GPT-NeoX models in fp16. (I tried a range of Pythia sizes from 14m to 1b, and they all fail due to loss overflow early in training.)

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.