[gemma4] resize_token_embeddings does not effect to embed_tokens_per_layer or output_embeddings

[KoichiYasuoka]

System Info

• transformers version: 5.5.0
• Platform: Linux-6.6.113+-x86_64-with-glibc2.35
• Python version: 3.12.13
• Huggingface_hub version: 1.8.0
• Safetensors version: 0.7.0
• Accelerate version: 1.13.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.10.0+cu128 (CUDA)
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: Tesla T4

(Google Colaboratory GPU)

Who can help?

@zucchini-nlp @Cyrilvallez

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Reproduce the behavior:

from transformers import Gemma4ForConditionalGeneration
mdl = Gemma4ForConditionalGeneration.from_pretrained("google/gemma-4-E2B-it")
e = mdl.get_input_embeddings()
f = mdl.model.language_model.embed_tokens_per_layer
g = mdl.get_output_embeddings()
print(e.num_embeddings, f.num_embeddings, g.out_features)
assert e.num_embeddings == f.num_embeddings == g.out_features
e = mdl.resize_token_embeddings(e.num_embeddings + 1)
f = mdl.model.language_model.embed_tokens_per_layer
g = mdl.get_output_embeddings()
print(e.num_embeddings, f.num_embeddings, g.out_features)
assert e.num_embeddings == f.num_embeddings == g.out_features

Expected behavior

All e.num_embeddings, f.num_embeddings and g.out_features should be increased to 262145.

[KoichiYasuoka]

I made PR #45236 but failed in test.

[zucchini-nlp]

I remember seeing a similar issue in gemma3n in the past which couldn't be reproduced unfortunately #39921. Looking at prev gemma3n I see that they zero out all token beyond vocab_size_per_layer which I believe was supposed to be added in Gemma4 as well 🤔 Not really sure if resizing per-layer-embeddings is intended for FT tbh

transformers/src/transformers/models/gemma3n/modeling_gemma3n.py

Lines 2048 to 2052 in b9f0fbf

	# Prepare per-layer inputs from inputs_ids
	per_layer_inputs_mask = torch.logical_and(input_ids >= 0, input_ids < self.vocab_size_per_layer_input)
	per_layer_inputs_tokens = torch.where(per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids))
	per_layer_inputs = self.language_model.get_per_layer_inputs(per_layer_inputs_tokens)

[KoichiYasuoka]

Thank you @zucchini-nlp for the information about #39921.

>>> from transformers import pipeline
>>> nlp=pipeline("image-text-to-text","google/gemma-3n-E2B-it",token="hf_...")
>>> nlp.model.get_input_embeddings()
Gemma3nTextScaledWordEmbedding(262400, 2048, padding_idx=0)
>>> nlp.model.get_output_embeddings()
Linear(in_features=2048, out_features=262400, bias=False)
>>> nlp.model.model.language_model.embed_tokens_per_layer
Gemma3nTextScaledWordEmbedding(262144, 7680, padding_idx=0)

Oops, embed_tokens_per_layer differs in its sizes...

>>> nlp.model.model.language_model.embed_tokens_per_layer.weight
Parameter containing:
tensor([[ 0.0613,  0.0200,  0.0603,  ..., -0.0703,  0.0552, -0.0486],
        [ 0.0150, -0.0282,  0.0469,  ..., -0.0530, -0.0645, -0.1309],
        [ 0.0049, -0.0160, -0.0308,  ..., -0.0562,  0.0176, -0.0562],
        ...,
        [ 0.0564,  0.0076, -0.0435,  ...,  0.0410, -0.0223,  0.0042],
        [ 0.0282, -0.0075, -0.0503,  ...,  0.0264,  0.0254, -0.0294],
        [-0.0064,  0.0013,  0.0262,  ..., -0.1357, -0.0215, -0.0679]],
       dtype=torch.bfloat16, requires_grad=True)
>>> nlp.model.model.language_model.embed_tokens_per_layer.weight.size()
torch.Size([262144, 7680])
>>> len(nlp.tokenizer)
262400

Umm, it's different from embed_tokens_per_layer in gemma4.

>>> nlp.model.config.vision_config
Gemma3nVisionConfig {
  "architecture": "mobilenetv5_300m_enc",
  "do_pooling": false,
  "dtype": "bfloat16",
  "hidden_size": 2048,
  "initializer_range": 0.02,
  "label_names": [
    "LABEL_0",
    "LABEL_1"
  ],
  "model_args": null,
  "model_type": "gemma3n_vision",
  "num_classes": 2,
  "rms_norm_eps": 1e-06,
  "transformers_version": "5.5.0",
  "vocab_offset": 262144,
  "vocab_size": 128
}
>>> nlp.model.model.embed_vision
Gemma3nMultimodalEmbedder(
  (embedding): Embedding(128, 2048)
  (hard_embedding_norm): Gemma3nRMSNorm()
  (soft_embedding_norm): Gemma3nRMSNorm()
  (embedding_projection): Linear(in_features=2048, out_features=2048, bias=False)
  (embedding_post_projection_norm): Gemma3nRMSNorm()
)
>>> nlp.tokenizer.convert_ids_to_tokens(262144)
'<end_of_image>'
>>> nlp("<end_of_image>")
Both `max_new_tokens` (=256) and `max_length`(=20) seem to have been set. `max_new_tokens` will take precedence. Please refer to the documentation for more information. (https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)
[{'input_text': '<end_of_image>', 'generated_text': "<end_of_image>\n,\nI have a project where I'd like to use a RESTful API.  I am unsure where to start.  Can you help me break down the process?\n\nOkay, let's break down the process of using a RESTful API, from understanding the basics to getting it working in practice. Here's a structured guide.\n\n**1. Understanding the Basics of REST (Representational State Transfer)**\n\n*   **What is REST?** REST is an architectural style for designing networked applications. It's based on using standard HTTP methods (like GET, POST, PUT, DELETE) to interact with resources over a network.\n*   **Key Concepts:**\n    *   **Resources:**  Representable things (data, objects, etc.) on a server.  Examples: a user, a product, a blog post.\n    *   **HTTP Methods:**  Define the actions you can perform on a resource.\n        *   `GET`:  Retrieve a resource. (e.g., `GET /users/123`)\n        *   `POST`:  Create a new resource. (e.g., `POST /users`)\n        *   `PUT`:  Update an existing resource completely"}]

Well, I think that google/gemma-3n-E2B-it uses token ids 0-262399 in Gemma3nTextModel, while embed_tokens_per_layer uses 0-262143 and embed_vision uses 262144-262271. Token ids 262272-262399 seem to be used in audio_tower but I'm not sure. But now I understand resize_token_embeddings is too difficult for gemma3n.

[KoichiYasuoka]

On the other hand, gemma4 uses token ids easier way.

[KoichiYasuoka]

I've just found this bug partially comes from _tied_weights_keys at

transformers/src/transformers/models/gemma4/modeling_gemma4.py

Line 2365 in 7f6cc4b

_tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}

and g.out_features can be fixed by PR #45311. But how do I fix f.num_embeddings?

[zucchini-nlp]

Haha, is that an agent reasoning output? Yeah, that makes sense if gemma3n already has reserved tokens in per_layer_input which is bigger than model's vocab size

My point is that Gemma4 might not need resizing the per-layer input when FT, which would be consistent with 3N release. Anyway, lemme ask the research team about intended behavior

We'll either override resize_token_embeddings or mask-out overflowing tokens

[KoichiYasuoka]

Haha, is that an agent reasoning output?

No, it's by my head and hands. Through checking #39921 (and #45292) I've found that this bug partially occurs from the confliction between resize_token_embeddings and #41580. And I've also found that Gemma3N has vocab_offset but Gemma4 not.

[zucchini-nlp]

Yep, I asked the Gemma team and we need to resize all per-layer embeddings that are not in vision/audio offset range. Linked a PR fixing Gemma4