Unexpected Behavior: Gemma INT4 Quantization and BF16 dtype Loading

Hi Gemma manitainers,

I am attempting to load and quantize the Gemma 3 model using the gemma library, specifically targeting INT4 quantization to reduce GPU memory footprint. I've referred to this [page](https://gemma-llm.readthedocs.io/en/latest/colab_quantization_sampling.html) for the quantization method. However, after modifying my code to explicitly use `gm.nn.IntWrapper` and `peft.quantize`, I observe no decrease in GPU RAM usage compared to loading the model without explicit quantization steps.

```python
import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"

from gemma import gm
from gemma import peft
import jax.numpy as jnp
import jax

# original loading method
# on two GPUs, each occupies approximately 43xx MiB of RAM.
model = gm.nn.Gemma3_1B()
params = gm.ckpts.load_params("/path/to/gemma3-1b-it-int4-v2/gemma3-1b-it-int4/")
tokenizer = gm.text.Gemma3Tokenizer("/path/to/gemma3-1b-it-int4-v2/tokenizer.model")

# my quantization code,
# on two GPUs, each still occupies approximately 43xx MiB of RAM.
# though I have noticed that some params is in int4 format
model = gm.nn.IntWrapper(model=gm.nn.Gemma3_1B(), dtype=jnp.int4)
original = gm.ckpts.load_params("/path/to/gemma3-1b-it-int4-v2/gemma3-1b-it-int4/")
params = peft.quantize(original, method=peft.QuantizationMethod.INT4, checkpoint_kernel_key='w')
tokenizer = gm.text.Gemma3Tokenizer("/path/to/gemma3-1b-it-int4-v2/tokenizer.model")
```

My questions are:

1. Am I misunderstanding the correct way to load and apply INT4 quantization using the `gemma` and `peft` method? If my current approach is incorrect, could you please provide guidance on the proper method?
2. Can the quantized model be used using `gm.text.ChatSampler`?
3. Additionally, I've noticed that even when attempting to explicitly set dtype=jnp.bfloat16 during model initialization (e.g., model = gm.nn.Gemma3_1B(dtype=jnp.bfloat16)), the model appears to be loaded in FP32 format. Given that Gemma is trained in BF16, is this expected behavior, or am I misconfiguring something?

Any insights or assistance would be greatly appreciated. Thank you!