Fix Flax params dtype
Author: patil-suraj
Pull request: yes
Labels: name: WIP

# What does this PR do?

The `dtype` argument in flax models is used ambiguously. This argument is actually supposed to specify the `dtype` of computation and not the `dtype` of model parameters. But in some models/modules, it's passed to kernel_initializers  which causes the `kernel` parameters to be initialized with that `dtype`. This causes the following issues

- in flax models, we don't pass `bias_init` to `Dense` layers since the default value is as expected by our models. So if we pass `dtype=jnp.bfloat16` it's only passed to `kernel_init`, so for a dense layer the kernel params are in `bfloat16` while the `bias` params are in `fp32`
-  This also causes issues with saving and loading models as explained in #12534

This PR corrects the usage of `dtype` in flax models and  adds `to_bf16`, `to_fp16` and `to_fp32` methods in `FlaxPreTrainedModel`. These methods could accept any arbitrary params tree and change its `dtype`. So if users want they could keep certain params in bf16 and certain others in fp32 however they like, by just passing the right parameters to these methods.

To allow keeping only certain params in half-precision the  `to_bf16` method accepts a mask that specifies what params to keep in `bf16` and what params in `fp32`. For example

```python3
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import freeze, unfreeze
from flax.traverse_util import flatten_dict, unflatten_dict
from transformers import FlaxBertModel, BertConfig

config = BertConfig(num_hidden_layers=1)
model = FlaxBertModel(config, dtype=jnp.dtype("bfloat16"))

# keep layer norm in fp32
def mask_fn(params):
    flat_params = flatten_dict(params)
    flat_mask = {path: (path[-2:] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) for path in flat_params}
    return unflatten_dict(flat_mask)

mask = mask_fn(model.params)
params = model.to_bf16(model.params, mask)

jax.eval_shape(lambda x : x, freeze(params)) # view the dtypes
```

- This PR also fixes an issue in some models where the `dtype` was never passed to some modules, so those modules were always doing computation in fp32 even if the user passed `bf16` or `fp16` dtype .

- This should now help enable mixed-precision training in flax models as we can keep the params and computation `dtype` separate. 

---

🚨  **BREAKING CHANGE** 🚨
**Note that: this will be a breaking change since the meaning of `dtype` is now changed and it's only used to specify the data type of computation and does not influence the data type of model parameters.**

Comments:
I like the design and I think it follows jax-design quite nicely (similar to how optax optimizers mask certain weights: https://github.com/huggingface/transformers/blob/1fec32adc6a4840123d5ec5ff5cf419c02342b5a/examples/flax/summarization/run_summarization_flax.py#L588)

This PR will necessarly have some breaking changes as after it loading a model with `dtype=bfloat16` won't convert the weights into bfloat16 anymore, so we should announce it well. 

Also it would be great if @avital could maybe quickly give his opinion on the API here

Hi folks, sorry for the radio silence, I'm back now. @jheek has thought carefully about the exact meaning of dtypes in Flax modules so I'd like to hear his take on this confusion. 

I think masking is the right approach here. The right dtype is very context dependent. During inference half precision is almost always better while during training it's pretty much never worth it. And then of course there is fine-tuning where the masked weights are basically in inference mode. The mask based API captures this complexity really well. 

just noticed a sneaky bug in some flax models, the `dtype` is never passed to some modules, for example here in bart, 
https://github.com/huggingface/transformers/blob/3fbb55c75779824aacfc43067f0892674a9cfbc6/src/transformers/models/bart/modeling_flax_bart.py#L400-L405
attention never receives  `dtype`, so it’s always in `fp32` even if the user passed `bf16` .

same with T5 here
https://github.com/huggingface/transformers/blob/3fbb55c75779824aacfc43067f0892674a9cfbc6/src/transformers/models/t5/modeling_flax_t5.py#L1368

@patrickvonplaten @sgugger I propose we make `dtype` required for all modules except user-facing once? So all main model classes will have a default type (`fp32`) but for all other submodules make it required to avoid such bugs.

And I think the flax template has to be adapted here as well

Hey @patrickvonplaten !

- added a couple more tests as you suggested
- updated the templates
- ran tests on both GPU and TPU and they pass 

Would be awesome if you could take quick final look :) 


Thanks for finishing the PR!

Think we just need to update the Flax templates now and we're good to go :-)