Generating with Flax fails when using padding
Author: haoliuhl
Pull request: no

### System Info

- `transformers` version: 4.25.1
- Platform: Linux-5.10.133+-x86_64-with-glibc2.27
- Python version: 3.8.15
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.13.0+cu116 (False)
- Tensorflow version (GPU?): 2.9.2 (False)
- Flax version (CPU?/GPU?/TPU?): 0.6.2 (cpu)
- Jax version: 0.3.25
- JaxLib version: 0.3.25
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No


### Who can help?

@patil-suraj, @patrickvonplaten, @LysandreJik

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Generation with GPT-2 and OPT doesn't work when using padding token.
Specifically, the issues are

1. In GPT-2, using builtin |endoftext| as pad token works with padding, but using customized token 
 doesn't work, the generation only repeats `!`, e.g., `
 My cat is cute!!!!!!!!!!!!!!!`.
2. In GPT-2, despite pad_token_id is optional in generate() function, one has to provide pad_token_id=tokenizer.pad_token_id otherwise error. 
3. In OPT, even using builtin <pad> token as padding token doesn't work, the generation only repeats `
`, e.g., `<pad>
My cat is cute













`.


The colab to reproduce the problems 1-3 is [here](https://colab.research.google.com/drive/1pcwTRU3snLjz8wTJx_Z5t4Z7IXnxHnDI?usp=sharing). 

A related issue to problem 2 is #18884. 

### Expected behavior

Expect the generation works with padding. E.g., `<pad>
My cat is super cute and I love her so much. I love her so much. I love`.

Comments:
cc @sanchit-gandhi and @gante 

Hey @lhao499!

> using builtin |endoftext| as pad token works with padding, but using customized token 
 doesn't work

Note that GPT2 was not pre-trained with padding tokens. We can use the `
` token as a substitute, but really we should specify an [attention mask](https://huggingface.co/transformers/glossary.html#attention-mask) to the model so that it doesn't attend to padded indices, therefore ignoring the value of the token. It's good to see that you've done this in your Colab for the GPT2 examples! (`model.generate(**inputs, ...)`)

In using a customised token (such as `
`), you are providing the model with a format entirely different to that seen during pre-training. We cannot expect our model to understand this different format given that it has never seen it before. This is likely the reason for the unexpected behaviour when using `
` as the padding token.

As for the OPT generation, it looks like this is fixed by setting the padding side: https://colab.research.google.com/drive/1pcwTRU3snLjz8wTJx_Z5t4Z7IXnxHnDI?usp=sharing

Hi Sanchit, when attention mask is provided, it's expected that customized token won't be attended at all, so effectively there is no differences between pretraining and inference formats?  

For the OPT, it looks like you linked my Colab, where the OPT generation does not work. 

Hey @lhao499! Sorry for the delay in getting back to you. I checked against PyTorch GPT-2, and we see the same phenomenon here, so we can exclude it as being a Flax specific issue: https://colab.research.google.com/drive/1qK2t8YNKLnX-oednqiVRDSFkFs5DPnY4#scrollTo=eGKOy0NbzCAW

cc'ing @gante here who might be able to provide some more insight!

Context: when we pass an attention mask for auto-regressive generation, it's expected that any padded tokens won't be attended to. Meaning we should be able to specify any arbitrary token as our pad token? 

For GPT-2, we see that generation works when the pad token is set to the default pad token. But it breaks when we set it to some arbitrary pad token.

Hey @lhao499 @sanchit-gandhi 👋 

In practice, adding new tokens and using the model straight away has very unpredictable results that depend on the framework used. For instance, the current version of JAX has the problem flagged above. However, if you use an older version of JAX, everything seems to work fine 🤷 If you try to do the same in TF CPU it might work, whereas on TF GPU it will crash unless you explicitly expand the vocabulary. PT also works if you expand the vocabulary, although with bad results.

Adding a new token corresponds to initializing a random entry in the embedding matrix, which has unforeseeable consequences. I highly advise against it unless you fine-tune the model afterward.

See [this colab](https://colab.research.google.com/drive/1Qly3125Q2happG1dGqdOpl1Q8MVNP_Nq?usp=sharing) for examples -- if you really want to go down this path, you might get away with an older Jax version ;)

P.S.: this Jax finding was a happy coincidence, my local desktop env had an older version Jax version which happened to work 👀 That's how unreliable this strategy is.

Thanks @sanchit-gandhi @gante for looking into the issues. 

Hi @gante, 

It is a surprising finding that adding new tokens has such unpredictable results depends framework and hardware. Thanks for trying different frameworks. 

It is a good strategy to resize the embedding matrix and fine-tune the model after adding new padding tokens. It would be great if transformers library could get rid of the strange results and/or raise errors. 

The attention mask still doesn't work as expected though. For instance, in OPT, even using builtin <pad> padding token only leads to repeating`
`, e.g., `<pad>
My cat is cute













`, as shown in the last block of https://colab.research.google.com/drive/1pcwTRU3snLjz8wTJx_Z5t4Z7IXnxHnDI?usp=sharing

👍 will have a look at OPT

@lhao499 there was indeed a problem with the attention masking in OPT, causing all but the longest input to fail. #21150 fixes it :)