BART-large + JAX produce nan loss during training/eval
Author: fen-deepscribe
Pull request: no

### System Info

- `transformers` version: 4.25.0.dev0
- Platform: Linux-5.15.0-1023-aws-x86_64-with-glibc2.35
- Python version: 3.9.12
- Huggingface_hub version: 0.10.1
- PyTorch version (GPU?): 1.13.0+cu117 (True)
- Tensorflow version (GPU?): 2.10.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.6.1 (gpu)
- Jax version: 0.3.24
- JaxLib version: 0.3.24
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No
- CUDA: NVIDIA-SMI 520.61.05    Driver Version: 520.61.05    CUDA Version: 11.8

### Who can help?

@patil-suraj

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I first experienced with a customized dataset with facebook/bart-large. The model trained well under torch+fp16. For speed purpose, I was trying to use jax using the provided example, but experienced nan loss in the evaluation loop. Sometime in the training loop as well. There is no issue with facebook/bart-base model.

After that, I switched to the original example provided in the main repo and trained on a public dataset. Yet no luck. Here is the process of reproduction:

1. Install transformer under development mode with `pip install -e`.
2. going into example folder `cd transformers/examples/flax/summarization`.
3. run training script using following command:
```
python run_summarization_flax.py \
--output_dir ~/data/test-cnn_dailymail \
--model_name_or_path facebook/bart-large \
--per_device_train_batch_size 2 \
--per_device_eval_batch_size 2 \
--dataset_name cnn_dailymail \
--dataset_config_name 3.0.0 \
--do_train \
--do_eval \
--max_train_samples 10000 \
--max_eval_samples 100 \
--max_source_length 512 \
--max_target_length 200
```
4. adding fp16 with `--dtype float16` resulting in the same issue

see log below:
```
INFO:__main__:***** Running training *****
INFO:__main__:  Num examples = 10000
INFO:__main__:  Num Epochs = 3
INFO:__main__:  Instantaneous batch size per device = 2
INFO:__main__:  Total train batch size (w. parallel & distributed) = 2
INFO:__main__:  Total optimization steps = 15000
Epoch... (1/3 | Loss: nan, Learning Rate: 3.3336666092509404e-05)                                                                                                                                        
Epoch... (1/3 | Eval Loss: nan | )                                                                                                                                                                       
```

I also noticed that the torch model is saved in float16. I've tried to convert the model to float32 and load with `from_pt=True`. Receive same nan loss problem. Not sure if it related to this issue or not.
```
You should probably UPCAST the model weights to float32 if this was not intended. See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this.
```

I will keep digging on this and share more information. Please kindly let me know if there is any recommend steps to debug this. :)

### Expected behavior

bart-large can be trained and evaluated with normal loss with JAX.

Comments:
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.

cc @sanchit-gandhi 

Hey @fen-deepscribe! Sorry for the late reply here!

Looks like there could be two issues at play:

1. The BART-large weights are indeed stored in fp16 on the HF Hub. The Flax `.from_pretrained` method respects the dtype of the stored params (no upcast/downcast operations), so when we load from the checkpoint at [facebook/bart-large](https://huggingface.co/facebook/bart-large), we load the weights in fp16. You can read more about this here: https://github.com/huggingface/transformers/issues/16736. Loading the weights in fp16 precision might be causing undesirable behaviour during training: Flax doesn't expect a dtype of fp16 (only fp32 or bf16). This could be messing up the dtypes of the activations, giving exploding grads and updates. If you want to load the weights in fp32, you can use the checkpoint at [patrickvonplaten/bart-large-fp32](https://huggingface.co/patrickvonplaten/bart-large-fp32).

2. BART-large is know to have numerical instabilities during fine-tuning (see https://github.com/huggingface/transformers/issues/15559 and https://github.com/huggingface/transformers/issues/15559#issuecomment-1062880564 and https://github.com/huggingface/transformers/issues/15559#issuecomment-1294457798). If you're fine-tuning the model for summarisation, you can try loading the checkpoint [facebook/bart-large-cnn](https://huggingface.co/facebook/bart-large-cnn) -> this checkpoint is stable and should be less prone to exploding gradients! I would give this checkpoint a go with your fine-tuning experiments. It tends to be a much easier fix than those linked in the aforementioned thread!

Let me know how you get on! More than happy to dig into this further if the exploding loss persists!

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.

Well, I have the same problem. After a period of training ,Bart's forecast output is starting to get really weird. No matter what the model inputs, Bart will produce the same text. Whether I use FP16 or not. 

Have you tried using the checkpoint [facebook/bart-large-cnn](https://huggingface.co/facebook/bart-large-cnn)?