# How to properly UPCAST the model weights to float32?
**URL:**
**Author:**  [@selamw](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fselamw)%5C
**Post date:** [April 11, 2024, 12:19am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-properly-upcast-the-model-weights-to-float32%2F81159%2F1 "2024-04-11T00:19:52Z")
```Python
from transformers import FlaxGemmaForCausalLM
import jax.numpy as jnp
model_name = "google/gemma-1.1-2b-it"
# Load the model with desired data type (bfloat16 for reduced memory usage)
model, params = FlaxGemmaForCausalLM.from_pretrained(model_name, revision="flax", _do_init=False, dtype=jnp.bfloat16, token=access_token)
# model, params = FlaxGemmaForCausalLM.from_pretrained(model_name, revision="flax", _do_init=False, dtype=jnp.float32, token=access_token)
# model, params = FlaxGemmaForCausalLM.from_pretrained(model_name, revision="flax", _do_init=False, token=access_token)
```
WARNING:
```auto
Some of the weights of FlaxGemmaForCausalLM were initialized in bfloat16 precision from the model checkpoint at google/gemma-1.1-2b-it: [('model', 'embed_tokens', 'embedding'), ......
('model', 'norm', 'weight')]
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.
```
I’m encountering a warning message every time I try to load the model. The message says: “Some of the weights of the FlaxGemmaForCausalLM were initialized in bfloat16 precision…”. I have two questions:
1. **How can I suppress this warning message?** Is it safe to ignore it, or are there potential consequences?
2. **How can I ensure all model weights are explicitly upcast to float32 for potentially higher accuracy?** `model = model.to_fp32(params)` is this correct, and how can I use it effectively?
**Additionally, I’d like to know how to get a summary of the model architecture and key parameters.**
---
**Author:**  [@swtb](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fswtb)%5C
**Post date:** [April 11, 2024, 11:45am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-properly-upcast-the-model-weights-to-float32%2F81159%2F2 "2024-04-11T11:45:37Z")
summary of model:
print(model)
Your usage seems correct, does the warning disappear when you run it?
---
**Author:**  [@selamw](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fselamw)%5C
**Post date:** [April 11, 2024, 5:44pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-properly-upcast-the-model-weights-to-float32%2F81159%2F3 "2024-04-11T17:44:48Z")
- After applying `model = model.to_fp32(params)` the warning message disappear. However, it changed it in to dictionary and generate raises below error:
```auto
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
in ()
----> 1 _ = p_generate(inputs, params, max_new_tokens)
[... skipping hidden 12 frame]
in generate(inputs, params, max_new_tokens)
1 def generate(inputs, params, max_new_tokens):
----> 2 generated_ids = model.generate(
3 inputs["input_ids"],
4 attention_mask=inputs["attention_mask"],
5 params=params,
AttributeError: 'dict' object has no attribute 'generate'
```
- Before applying `model = model.to_fp32(params)`, print(model) : \
- After applying `model = model.to_fp32(params)`, print(model) :
list of dictionaries, however, i’m looking for summary of the model architecture.
|