

TITLE
Fix Flax params dtype

COMMENTS
8

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
# 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.**