[Flax] Torch fp16 model weights not upcast when loaded in Flax
Author: sanchit-gandhi
Pull request: no

In some scenarios, one may want to load a Flax model directly from pre-trained PyTorch model weights. In this process, the original dtype of the PyTorch model weights is maintained when loaded into Flax. For models such as [bart-large](https://huggingface.co/facebook/bart-large), which has it's PyTorch weights stored in fp16 on the Hub, this can result in a Flax model with weights in an undesirable dtype. This is highlighted by the following code snippet, which first loads a FlaxSpeechEncoderDecoderModel from entirely fp32 PyTorch weights, and then again from fp32 encoder weights and fp16 decoder weights:
```python
from transformers import FlaxSpeechEncoderDecoderModel

# fp32 PyTorch weights
encoder_id = 'hf-internal-testing/tiny-random-wav2vec2'
decoder_id = 'hf-internal-testing/tiny-random-bart'

model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(encoder_id, decoder_id, encoder_from_pt=True, decoder_from_pt=True)
print("-----------From fp32 PyTorch weights-----------")
print(f"Encoder dtype: {model.params['encoder']['masked_spec_embed'].dtype}")
print(f"Decoder dtype: {model.params['decoder']['model']['decoder']['embed_tokens']['embedding'].dtype}")

# same decoder as previously, but with weights downcasted to fp16
decoder_id = 'sanchit-gandhi/tiny-random-bart-fp16'

model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(encoder_id, decoder_id, encoder_from_pt=True, decoder_from_pt=True)
print("---------From fp32/fp16 PyTorch weights---------")
print(f"Encoder dtype: {model.params['encoder']['masked_spec_embed'].dtype}")
print(f"Decoder dtype: {model.params['decoder']['model']['decoder']['embed_tokens']['embedding'].dtype}")
```
Output:
```
-----------From fp32 PyTorch weights-----------
Encoder dtype: float32
Decoder dtype: float32

---------From fp32/fp16 PyTorch weights---------
Encoder dtype: float32
Decoder dtype: float16
```
Having a model stored in two different dtype raises issues with training - Optax optimisers expect the model to maintain one uniform dtype. Furthermore, the default assumption is that all Flax model weights are in fp32.

This weight conversion is handled by the general conversion script: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flax_pytorch_utils.py. Would it be wise to inform the user of the potentially erroneous model dtype in this scenario? If informed, they could then call the `to_fp32` method from `modeling_flax_utils` to upcast the weights to fp32:
https://github.com/huggingface/transformers/blob/a9604067225219e132abdff2793f78ead798453b/src/transformers/modeling_flax_utils.py#L231

Comments:
cc @patrickvonplaten @patil-suraj 

Great catch! Gosh, we should have never uploaded `bart-large` with fp16 weights - I think this happened accidentally and a long time ago :-/ Usually we want all weights to be stored as full fp32 weights. 

To be honest, for now I think this is really an edge-case - I don't know any model besides bart that has its weights uploaded in fp16, so I think we could do three things here:

- 1. Don't do anything - it's an edge case
- 2. Make `from_pretrained(...)` error out
- 3. Automatically convert to fp32

I strongly advocate fro 1. or 2. here. 

I'll upload the original weights of `bart-large` in full fp32 probably in a separate repo now.

What do you think ? @patil-suraj @sanchit-gandhi 



Might be related to https://github.com/huggingface/transformers/issues/15559

If this is solely an issue concerning `bart-large` and this truly is an edge-case, then 1 or 2 seem reasonable. 3 could cause some serious ramifications for instances where the fp16 model is currently used (e.g. new OOM's with training). In 2, would the error out completely prohibit the user from loading weights in fp16, or just provide them with a warning and the advice to upcast the weights/load from fp32?

In 2. I'd completely error out and state that the two checkpoints have different precision and can't be combined

My worry with a complete error out would be that it prevents the user from ever being able to load the model, even if they have the intent of upcasting/correcting for the dtype. My suggestion would be to add a warning to the [`from_pretrained`](https://github.com/huggingface/transformers/blob/b24201fa44e1a14e83be890dcbc231e926c37ec1/src/transformers/modeling_flax_utils.py#L298) method in [`modeling_flax_utils.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flax_utils.py) if the Flax model weights are loaded in a dtype other than fp32. This could then be proceeded by the advice that the user should upcast to fp32 using the provided method `to_fp32`. By displaying a warning instead of an error out, the user is still able to load the model and then subsequently rectify any dtype mismatch.

Guess adding a warning and no automatic upcasting is fine as well! Just not in favor of automatic upcasting :-)

Agree with @sanchit-gandhi here. I'm in favour of adding a warning and letting the user know that weights are not `fp32`.

The user warning for the Flax `.from_pretrained` method was implemented in #16762. As an extreme edge case and following an extensive offline discussion, it was decided that the fp16 PyTorch weights for [bart-large](https://huggingface.co/facebook/bart-large) will remain as is. The original checkpoint has been reconverted and uploaded it in fp32 to another repo for those wishing to explicitly use full-precision weights: https://huggingface.co/patrickvonplaten/bart-large-fp32 Note that the fp16 weights should not be an issue for any PyTorch models: the PyTorch `.from_pretrained` method automatically upcasts model weights to fp32.