Fix Gemma 3 safetensor loading

[copybara-service]

Fix Gemma 3 safetensor loading

Loading Gemma 3 from safetensor weights fails due to two reasons:

• Safetensor weights follow the pattern "language_model*". "language_model" was missing.
• When we load weights, weights from the vision part of the model are automatically checked.
  Since they are absent from the key dict, an error is thrown.

Colab: https://colab.research.google.com/gist/abheesht17/9e85a2f7d08b60712f42bd1b44481ee6/bug-fix-gemma-3-loading-tunix.ipynb

Code snippet:

from tunix.models.gemma3 import params_safetensors as params_safetensors_lib
from tunix.models.gemma3 import model as gemma_lib
from flax import nnx

model_config = gemma_lib.ModelConfig.gemma3_4b_it()
MODEL_CP_PATH = local_model_path
gemma3 = params_safetensors_lib.create_model_from_safe_tensors(
    MODEL_CP_PATH, (model_config),
)

Error:

ValueError: Only one key should be found: [] for language_model.model.embed_tokens.weight