[Flax] from_pretrained does not consider the passed dtype

Environment info

When loading a flax model with from_pretrained the type argument is not used. The weights are initialized with the dtype of saved weights.

So if you do

model = FlaxGPT2ForCausalLM.from_pretrained("gpt2", dtype=jnp.dtype("bfloat16"))

# check the dtype of one of the params 
model.params["transformer"]["wpe"]["embedding"].dtype
=> dtype("float32")

We should probably cast the weights to self.dtype.

As a workaround for bf16, one could manually cast the weighs with

def to_bf16(t):
    return jax.tree_map(lambda x: x.astype(jnp.bfloat16) if x.dtype == jnp.float32 else x, t)

model.params = to_bf16(model.params)

cc @patrickvonplaten

[avital]

I think it’s fine to manually port weights to bfloat16 if you want to. In general all Flax layers accept a dtype attribute when it’s safe to do intermediate computation in bloat16 and you can set dtype=bfloat16 for those layers. Keeping parameters as bfloat16 should only be necessary if the model is huge and the parameters can’t fit on device memory, from what I know. I think it’s tricky to get that right and requires careful attention to which parameters are safe to keep in bfloat16, but I don’t have too much personal context here. I can ask others if that helps.

So I’m first curious whether indeed it’s necessary to keep parameters as bfloat16 in this case, and if so, why

[patrickvonplaten]

will soon be taken care of by @patil-suraj 😃