[Flax] Torch fp16 model weights not upcast when loaded in Flax

[sanchit-gandhi]

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, 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:

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:

transformers/src/transformers/modeling_flax_utils.py

Line 231 in a960406

def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None):

[sanchit-gandhi]

cc @patrickvonplaten @patil-suraj

[patrickvonplaten]

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

[patrickvonplaten]

Might be related to #15559

[sanchit-gandhi]

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?

[patrickvonplaten]

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

[sanchit-gandhi]

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 method in 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.

[patrickvonplaten]

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

[patil-suraj]

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

[sanchit-gandhi]

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