[Flax `.from_pretrained`] Raise a warning if model weights are not in float32
The Flax `.from_pretrained` method respects the dtype of the model weights from which it is loaded. For model weights stored in bfloat16/float16, Flax models are instantiated with parameter weights in bfloat16/float16 respectively (see #16736). The general assumption is that all Flax model weights are in float32. Loading and storing model weights in a lower precision (bfloat16/float16) is likely to lead to undesirable behaviour and model instabilities. This PR adds a warning to the `.from_pretrained` method should any of the model weights not be in float32, and advices the user to upcast the weights to float32 prior to use.
_The documentation is not available anymore as the PR was closed or merged._<|||||>As an example, loading a set of PyTorch float16 Bart model weights into a FlaxBartForCausalLM model produces the following warning:
```python
from transformers import FlaxBartForCausalLM
model = FlaxBartForCausalLM.from_pretrained('sanchit-gandhi/tiny-random-bart-fp16', from_pt=True)
```
```
Some of the weights of FlaxBartForCausalLM were initialized in float16 precision from the model checkpoint at sanchit-gandhi/tiny-random-bart-fp16:
[('model', 'decoder', 'embed_positions', 'embedding'), ('model', 'decoder', 'embed_tokens', 'embedding'), ('model', 'decoder', 'layernorm_embedding', 'bias'), ('model', 'decoder', 'layernorm_embedding', 'scale'), ('model', 'decoder', 'layers', '0', 'encoder_attn', 'k_proj', 'bias'), ('model', 'decoder', 'layers', '0', 'encoder_attn', 'k_proj', 'kernel'), ('model', 'decoder', 'layers', '0', 'encoder_attn', 'out_proj', 'bias'), ('model', 'decoder', 'layers', '0', 'encoder_attn', 'out_proj', 'kernel'), ('model', 'decoder', 'layers', '0', 'encoder_attn', 'q_proj', 'bias'), ('model', 'decoder', 'layers', '0', 'encoder_attn', 'q_proj', 'kernel'), ('model', 'decoder', 'layers', '0', 'encoder_attn', 'v_proj', 'bias'), ('model', 'decoder', 'layers', '0', 'encoder_attn', 'v_proj', 'kernel'), ('model', 'decoder', 'layers', '0', 'encoder_attn_layer_norm', 'bias'), ('model', 'decoder', 'layers', '0', 'encoder_attn_layer_norm', 'scale'), ('model', 'decoder', 'layers', '0', 'fc1', 'bias'), ('model', 'decoder', 'layers', '0', 'fc1', 'kernel'), ('model', 'decoder', 'layers', '0', 'fc2', 'bias'), ('model', 'decoder', 'layers', '0', 'fc2', 'kernel'), ('model', 'decoder', 'layers', '0', 'final_layer_norm', 'bias'), ('model', 'decoder', 'layers', '0', 'final_layer_norm', 'scale'), ('model', 'decoder', 'layers', '0', 'self_attn', 'k_proj', 'bias'), ('model', 'decoder', 'layers', '0', 'self_attn', 'k_proj', 'kernel'), ('model', 'decoder', 'layers', '0', 'self_attn', 'out_proj', 'bias'), ('model', 'decoder', 'layers', '0', 'self_attn', 'out_proj', 'kernel'), ('model', 'decoder', 'layers', '0', 'self_attn', 'q_proj', 'bias'), ('model', 'decoder', 'layers', '0', 'self_attn', 'q_proj', 'kernel'), ('model', 'decoder', 'layers', '0', 'self_attn', 'v_proj', 'bias'), ('model', 'decoder', 'layers', '0', 'self_attn', 'v_proj', 'kernel'), ('model', 'decoder', 'layers', '0', 'self_attn_layer_norm', 'bias'), ('model', 'decoder', 'layers', '0', 'self_attn_layer_norm', 'scale'), ('model', 'decoder', 'layers', '1', 'encoder_attn', 'k_proj', 'bias'), ('model', 'decoder', 'layers', '1', 'encoder_attn', 'k_proj', 'kernel'), ('model', 'decoder', 'layers', '1', 'encoder_attn', 'out_proj', 'bias'), ('model', 'decoder', 'layers', '1', 'encoder_attn', 'out_proj', 'kernel'), ('model', 'decoder', 'layers', '1', 'encoder_attn', 'q_proj', 'bias'), ('model', 'decoder', 'layers', '1', 'encoder_attn', 'q_proj', 'kernel'), ('model', 'decoder', 'layers', '1', 'encoder_attn', 'v_proj', 'bias'), ('model', 'decoder', 'layers', '1', 'encoder_attn', 'v_proj', 'kernel'), ('model', 'decoder', 'layers', '1', 'encoder_attn_layer_norm', 'bias'), ('model', 'decoder', 'layers', '1', 'encoder_attn_layer_norm', 'scale'), ('model', 'decoder', 'layers', '1', 'fc1', 'bias'), ('model', 'decoder', 'layers', '1', 'fc1', 'kernel'), ('model', 'decoder', 'layers', '1', 'fc2', 'bias'), ('model', 'decoder', 'layers', '1', 'fc2', 'kernel'), ('model', 'decoder', 'layers', '1', 'final_layer_norm', 'bias'), ('model', 'decoder', 'layers', '1', 'final_layer_norm', 'scale'), ('model', 'decoder', 'layers', '1', 'self_attn', 'k_proj', 'bias'), ('model', 'decoder', 'layers', '1', 'self_attn', 'k_proj', 'kernel'), ('model', 'decoder', 'layers', '1', 'self_attn', 'out_proj', 'bias'), ('model', 'decoder', 'layers', '1', 'self_attn', 'out_proj', 'kernel'), ('model', 'decoder', 'layers', '1', 'self_attn', 'q_proj', 'bias'), ('model', 'decoder', 'layers', '1', 'self_attn', 'q_proj', 'kernel'), ('model', 'decoder', 'layers', '1', 'self_attn', 'v_proj', 'bias'), ('model', 'decoder', 'layers', '1', 'self_attn', 'v_proj', 'kernel'), ('model', 'decoder', 'layers', '1', 'self_attn_layer_norm', 'bias'), ('model', 'decoder', 'layers', '1', 'self_attn_layer_norm', 'scale')]
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.
```<|||||>Sorry this a super nitty question, but I just wanted to ask to make sure we're all on the same page for best practice! Should one not ideally merge their own PR's rather than the reviewer?<|||||>Aah, yes! One should merge their own PRs, I rushed a bit this one.