Add Mistral GPT-2 Stability Tweaks
Author: siddk
Pull request: yes

Following the release of [Mistral-v1](https://github.com/stanford-crfm/mistral), we are pushing a *Draft PR* with the stability fixes we made for training GPT-2 models directly to the base GPT-2 class definition (ensuring backwards compatibility).

This is in line with the following issues:
- https://github.com/stanford-crfm/mistral/issues/86: Enabling sharing Mistral Checkpoints via HF Hub (@osanseviero)
- https://github.com/huggingface/transformers/issues/13463: Upcasting Scaled Dot-Product Attention + Layerwise Scaling for stability (@lvwerra) 

Concretely we implement:
- Weight initialization from the original GPT-2 paper (by default, shouldn't affect folks unless they are training GPT-2 models from scratch)
- Layer-wise scaling in scaled dot-product attention (optional flag; necessary for running/loading Mistral GPT-2 Models)
- Scaled Dot-Product Attention Reordering (scale before dot-product) & FP32 Upcasting when training in Mixed Precision (optional flag; only necessary for training new Mistral/other GPT-2 models).

This is a Draft PR to aid in @lvwerra and @thomwolf's work training GPT-2 models stably; we plan on implementing tests (please let us know potential pain points), adding documentation, and will act on any other feedback you have 🙂. 

CC Mistral team: @lorr1, @J38, @santhnm2
CC others at HF + GPT-2 Model reviewers: @stas00, @patrickvonplaten, @LysandreJik

Resolves #13463 

Comments:
@lvwerra did us the courtesy of implementing many of the changes - most of the comments have been resolved!

Question for @LysandreJik @patrickvonplaten; what tests should we add for these changes, and what documentation should we add? I can take an initial pass later today.

@lvwerra - this should be good to start your GPT-2 training runs.

Separately; do we need to implement similar fixes for TF and Flax? I'm not as familiar with either of those libraries/those GPT implementations...

Thank you for working on this @siddk. After seeing the code, I agree that this can indeed be merged into GPT-2 rather than by creating a new class as I was suggesting a few weeks ago.

Regarding the documentation, I would put a comment about this in the "Tips" of GPT-2, which can be found [here](https://github.com/huggingface/transformers/blame/master/docs/source/model_doc/gpt2.rst#L32).

Regarding tests, I would aim to test added features:

> - Weight initialization from the original GPT-2 paper (by default, shouldn't affect folks unless they are training GPT-2 models from scratch)

I would check that a newly initialized GPT-2 with a specific initializer range results in a change of weights. If this is too hard to implement, I don't think that's too much of an issue to not test it.

> Layer-wise scaling in scaled dot-product attention (optional flag; necessary for running/loading Mistral GPT-2 Models)

I would check that the scaling is indeed applied; should be able to leverage the `output_hidden_states` flag for that.

Let us know if we can help with anything or if you had other tests in mind.

Thanks so much @LysandreJik - we'll implement the tests you suggest (CC @j38 @santhnm2). I noticed that we're getting this weird failed test with building documentation; any idea how to get to the bottom of it, the Sphinx error message isn't too helpful...

@lvwerra - hopefully this is enough to get started!

Hi everyone, I just pushed a first pass version of tests that just ensure passing the `scale_attn_by_layer_idx` and `reorder_and_upcast_attn` flags don't result in a crash. However, we were wondering it was worth doing more robust tests (for example actually verifying the attention weights changed as a result of upcasting), and if so, how we would go about doing this considering the new functions are pretty much black boxes from the point of view of the tests. Thanks for your help!

Thanks @stas00! @santhnm2 - can you make the suggested changes (and fix the "docs" test that @LysandreJik suggested a fix for!). 

I'll add some notes to the "Tips" section, then we can make this an official PR - thanks all for the help.

Hi all - @santhnm2 has resolved the remaining tests and docs, and all checks currently pass! 

Excited to get this merged!

I added a test for the new weight initializations that will check the empirical std and mean are within a reasonable range of the expected value. Please let me know if you'd like any modifications.

I just added an assert check as well that the upcasting is actually happening.

Thanks @patrickvonplaten @LysandreJik @sgugger - Merge conflicts should be addressed now!

Using a model trained with `reorder_and_upcast_attn=True` to generate text with `generate` function results in the following error:
```python
from transformers import GPT2LMHeadModel, AutoTokenizer, AutoConfig, pipeline
config = AutoConfig.from_pretrained('gpt2', reorder_and_upcast_attn=True)
model = GPT2LMHeadModel(config)
tokenizer = AutoTokenizer.from_pretrained('gpt2')
inputs = tokenizer('test', return_tensors='pt')
model.generate(**inputs)
```
```python
~/miniconda3/envs/codeparrot/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py in _upcast_and_reordered_attn(self, query, key, value, attention_mask, head_mask)
    244             with autocast(enabled=False):
    245                 q, k = query.reshape(-1, seq_len, dk), key.transpose(-1, -2).reshape(-1, dk, seq_len)
--> 246                 attn_weights = torch.baddbmm(attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor)
    247                 attn_weights = attn_weights.reshape(bsz, num_heads, seq_len, seq_len)
    248         else:

RuntimeError: Expected batch2_sizes[0] == bs && batch2_sizes[1] == contraction_size to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)
```
The reason is that when generating the previous keys and values are stacked to the current ones while the queries are not extended. 
https://github.com/huggingface/transformers/blob/babb7f40fb9087ad31cd38bb12117427a614a173/src/transformers/models/gpt2/modeling_gpt2.py#L322-L325
Thus the `q` tensor has a different shape than `k` which causes `torch.baddbmm` to throw an error. How should this be handled?

One way would be to add an extra condition to the if-statement:
```python
if self.reorder_and_upcast_attn and (query.shape==key.shape):
    attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
```
Or should the user be explicitly be warned? Another option is the repeat the query tensor along the sequence dimension:
```python
if q.shape != k.shape:
    q = q.repeat(1, 1, k.shape[2], 1)
```
But I have to think a bit harder and test that this is actually equivalent for the last option.

If you’re just using a model (for inference/generation) you should just be able to turn off that flag (as long as scaling by layer idx is still True). They should be equivalent.

This does conflict with typical use of `from_pretrained` though — any suggestions @sgugger @lvwerra?

Just wanted to follow-up - is there a plan to merge this PR?

Hi @siddk, the plan is definitely to merge this PR. The issue @lvwerra mentions seems to be real, however: if the model cannot be used for anything else than training with `reorder_and_upcast_attn=True`, it should be mentioned in the documentation and should fail with an explicit error message.

Would the fix offered by @lvwerra regarding the mismatched `q` and `k` tensor sizes work for you?

Sorry @LysandreJik @lvwerra I totally missed the last response.

That issue makes sense, I think the best fix is the "if" condition @lvwerra added (only upcast/reorder if q and k are the same shape). I don't think repeating the query actually keeps the same semantics - might actually give different results.

@lvwerra would it be easier if I made that change?

I believe Leandro is off for a few days, so if you have time to make that change it would be very welcome @siddk!

@LysandreJik - actually found a better, more general fix. Verified that the above snippet from @lvwerra works as expected - should be good to go.

@LysandreJik Just added tests with the `generate` method, let me know if these look alright to you.

These looks good, thank you @santhnm2. Merging!