Fix Flax params dtype

[patil-suraj]

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 [Flax] from_pretrained does not consider the passed dtype #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

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.

[sgugger]

I think I agree with your solution @patil-suraj but I'll defer to @patrickvonplaten 's final say as I don't know Jax/Flax enough to have a strong opinion.

Initializing a model in bfloat16 of float16 seems like the wrong thing to do.

[patrickvonplaten]

I like the design and I think it follows jax-design quite nicely (similar to how optax optimizers mask certain weights:

transformers/examples/flax/summarization/run_summarization_flax.py

Line 588 in 1fec32a

def decay_mask_fn(params):

)

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

[avital]

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.

[jheek]

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.

[patil-suraj]

just noticed a sneaky bug in some flax models, the dtype is never passed to some modules, for example here in bart,

transformers/src/transformers/models/bart/modeling_flax_bart.py

Lines 400 to 405 in 3fbb55c

	self.self_attn = FlaxBartAttention(
	config=self.config,
	embed_dim=self.embed_dim,
	num_heads=self.config.encoder_attention_heads,
	dropout=self.config.attention_dropout,
	)

attention never receives dtype, so it’s always in fp32 even if the user passed bf16 .

same with T5 here

transformers/src/transformers/models/t5/modeling_flax_t5.py

Line 1368 in 3fbb55c

self.encoder = FlaxT5Stack(encoder_config, self.shared)

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

[patil-suraj]

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

[patrickvonplaten]

Great, the functionality looks good. Left a couple of suggestions for improvement for the docstring. I wonder whether we should even leave a couple more sentences under to_bfloat16(...) and to_float16(...) that state why one would call such a method. E.g. bfloat16 for TPU to save memory and float16 for GPU to save memory. Should we maybe also adapt the docstring of the dtype input that one can pass to a Flax model to really make sure that people understand that dtype never influences the weights but just the computations?

Also would be great to check that the docstring is nicely shown in the docs .

Finally, we should leave a big statement in the PR description with 🚨 that this PR is backwards breaking.

We should really bring across the message that dtype == computation precision & to_...() converts parameters

[patrickvonplaten]

Thanks for finishing the PR!

[patrickvonplaten]

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