FlaxGPT2LMHeadModel doesn't fail on mismatch between tokenizer and model vocab size
## Environment info
- `transformers` version: 4.9.0
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.11
- PyTorch version (GPU?): 1.9.0+cu102 (True)
- Tensorflow version (GPU?): 2.5.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (gpu)
- Jax version: 0.2.18
- JaxLib version: 0.1.69
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help
@patil-suraj
@patrickvonplaten

## Information

Model I am using (Bert, XLNet ...): FlaxGPT2LMHeadModel

The problem arises when using:
* run_clm_flax.py script
* other situations as well as the problem is lack of error handling

The tasks I am working on is:
* training GPT2

## To reproduce

Run the colab linked below: 
https://colab.research.google.com/drive/1G07rCEZyquu5CFVvQZlkK1Mgjd1gT86g?usp=sharing

## Expected behavior

We found an issue with our Pytorch model where tokenizer vocab size was not aligned with our model embedding size. This resulted in an error - as expected. However, the model was trained in Flax with that mismatch, and on testing, it appears to work while with vocab size larger than embedding size. I think this should raise an exception instead. We use FlaxGPT2LMHeadModel. 
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/master/CONTRIBUTING.md) are likely to be ignored.<|||||>I think it's difficult to ensure that a model which has a mismatch in vocab size with its tokenizer runs correctly. We could think about showing a better error message than:

```
IndexError: index out of range in self
```

If it's the word embeddings matrix that throws that error! I don't think we can make sure though that the model works with an incorrect tokenizer vocab size<|||||>Hi Patrick, thanks for commenting on the issue! I actually have no concerns
with the error in model inference due to the mismatch in vocab size with
its tokenizer. The concern was that the model trained with this mismatch, I
think it should throw an error instead during training. Thanks, Darek

On Tue, 24 Aug 2021 at 11:16, Patrick von Platen ***@***.***>
wrote:

> I think it's difficult to ensure that a model which has a mismatch in
> vocab size with its tokenizer runs correctly. We could think about showing
> a better error message than:
>
> IndexError: index out of range in self
>
> If it's the word embeddings matrix that throws that error! I don't think
> we can make sure though that the model works with an incorrect tokenizer
> vocab size
>
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/12842#issuecomment-904471318>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ADYSLBN5A54ZS36YNXWSND3T6NPQNANCNFSM5AZ4HV5Q>
> .
> Triage notifications on the go with GitHub Mobile for iOS
> <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675>
> or Android
> <https://play.google.com/store/apps/details?id=com.github.android&utm_campaign=notification-email>
> .
>
<|||||>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/master/CONTRIBUTING.md) are likely to be ignored.