FlaxGPTNeoForCausalLM not working properly with fp16 when using left padding.
Author: T-Almeida
Pull request: no

### System Info

WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/transformers/commands/env.py:52: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.list_physical_devices('GPU')` instead.
WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/transformers/commands/env.py:52: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.list_physical_devices('GPU')` instead.
2023-01-18 15:47:59.442290: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:42] Overriding orig_value setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.25.1
- Platform: Linux-5.10.147+-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.13.1+cu116 (True)
- Tensorflow version (GPU?): 2.9.2 (True)
- Flax version (CPU?/GPU?/TPU?): 0.6.3 (gpu)
- Jax version: 0.3.25
- JaxLib version: 0.3.25
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@gante 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Hello there, I am having a bit of trouble to successfully generate text (beam search) with the FlaxGPTNeoForCausalLM model when using fp16.

I provide two colab notebooks to replicate this issue:
 - torch version, which works fine both on fp32 and fp16: https://colab.research.google.com/drive/15Fy3VmTfUVGGGC1NAGP8p_DqZajDzxZk?usp=sharing
 - flax version, which fails on fp16: https://colab.research.google.com/drive/1t588H8_1SGSj6g1yVXgkeRiIvsxQiOKA?usp=sharing

Very briefly in torch I am converting the model to fp16 by doing this: `GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M", torch_dtype=torch.float16)`, while in Flax I am doing the following:
```python
jax_model = FlaxGPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M", dtype=jax.numpy.float16)
jax_model.params = jax_model.to_fp16(jax_model.params)
```
For both cases, I am using the following sentences as input **with** left padding:
```
texts = ["My name is Julien and I like to", "Why float16 is giving such a strange outputs?"]
```
Output of the torch version:
```python
['





My name is Julien and I like to call you Julien. I like to call you Julien. I like to call you Julien. I like to call you Julien. I like to call you Julien. I like to call you Julien. I like to call you',
 '




Why float16 is giving such a strange outputs?\n\nA:\n\nfloat16 is giving such a strange outputs?\n\nYes, it does.\n\nA:\n\nYes, it does.\n\nA:\n\nYes, it does.\n\nA:\n\n']
```

Output of the flax version:
```python
['





My name is Julien and I like to!
















































',
 '




Why float16 is giving such a strange outputs?!
















































']
```

As you can see, in the case of flax I will always get the `!` token (which corresponds to id:0) and then I will get the `
` token (id: 50256). Strangely if I dont do the left padding and process each sentence individually I will get the same output as the torch version.

### Expected behavior

Basically, I want to have the equivalent of `GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M", torch_dtype=torch.float16)` but in Flax. So, from the [docs](https://huggingface.co/docs/transformers/v4.25.1/en/main_classes/model#transformers.FlaxPreTrainedModel.from_pretrained.dtype), I get that using `to_fp16()` will convert the model params to fp16 and changing the dtype to `jnp.float16` will force the computation to be in fp16. However, when I set `dtype=jnp.float16` and using left padding, the generation does not work properly. If instead, I just use the `to_fp16()` to convert the params and leave `dtype=jnp.float32`, the code works properly, but it is two times slower than the pytorch version, which means that is not truly fp16.

I also want to add that this issue only seems to appear when I add left padding to the inputs in Flax.

Any idea why is this happening?

P.S. I am also not sure if what I am doing is correct, but I couldn't find anything similar to this issue.

**UPDATE**
I also noticed that in the pytorch version, if I use the default padding behaviour (right padding) I get the following warning, which **does not appear** in flax. 
```
A decoder-only architecture is being used, but right-padding was detected! For correct generation results, please set `padding_side='left'` when initializing the tokenizer.
```
So I tried using right padding in the case of Flax and for my surprise, it worked! It gave me the same outputs as the left padded version in torch.

I do not understand if this behaviour is intended or not, but I find it to be a bit confusing, since I believe that the left padding would make more sense.

Comments:
Hey @T-Almeida 👋 To be candid with you, I've never played with Flax + generate + fp16, so I can't confirm whether it is a model, generate, or flax issue without a deep dive :) In any case, I can tell you that we've stopped development on Flax, and that is why you won't see newer features there (such as the left padding warning). 

With decoder-only models, you must use left padding. Otherwise, the first generated token will take `<PAD>` as the previous token and will return different results (the results you get with right padding are slightly different if you pay attention).

______________________________________

@sanchit-gandhi nvm, I've located the bug. Numerical masking strikes again!

(OLD: TL;DR before going into debug mode, a quick check with you :D Flax + generate + fp16 on GPTNeo returns gibberish, where fp32 works fine. Have you seen anything similar before? There is also a chance that the example makes incorrect use of fp16.)


Hi @gante ty for the feedback, you are right with right padding the outputs are indeed different (my bad). Sadly, I also notice that the Flax generate API is, indeed, the one that lacks more features when compared with torch or tf (specially the contrastive_search method, that I can only use in torch, because there is no TFGPTneo... implementation :'( )

Related to this issue, when using left padding with fp16 the model is outputting nan logits when predicting the next token, check the below snippet:

```python
outs = jax_model(input_ids, attention_mask=attention_mask).logits

DeviceArray([[[nan, nan, nan, ..., nan, nan, nan],
              [nan, nan, nan, ..., nan, nan, nan],
              [nan, nan, nan, ..., nan, nan, nan],
              ...,
              [nan, nan, nan, ..., nan, nan, nan],
              [nan, nan, nan, ..., nan, nan, nan],
              [nan, nan, nan, ..., nan, nan, nan]],

             [[nan, nan, nan, ..., nan, nan, nan],
              [nan, nan, nan, ..., nan, nan, nan],
              [nan, nan, nan, ..., nan, nan, nan],
              ...,
              [nan, nan, nan, ..., nan, nan, nan],
              [nan, nan, nan, ..., nan, nan, nan],
              [nan, nan, nan, ..., nan, nan, nan]]], dtype=float16)
```

So, the issue should be in the model call and not in the generation?

@T-Almeida thank you for your comment! It turns out that I've seen this pattern of `nan` this week, and the same fix applies here. Will open a PR soon ;)

Re TFGPTNeo, we are always open to contributions 🙏 I'd be happy to guide anyone that'd like to contribute!

Hey @T-Almeida,

Sounds like you've done a good job at assimilating the different Flax dtype terms (which isn't straightforward)! And cool to see that you're running JAX on GPU!

As you've correctly specified, `to_fp16()` will convert all the Flax model params to float16, but will leave the computations untouched (i.e. they remain in float32 precision). We need to specify `dtype=jnp.float16` to ensure our forward pass is also done in float16 precision.

Looks like @gante is on the case with fixing the Flax attention masks (which seems to be the problem here)!

@T-Almeida merged, if you install the development version of transformers it should work :)

I can confirm that now it is working properly! Thanks a lot @gante, @sanchit-gandhi for the really quick fix!