Skip to content

Fix OOM in CI by reducing intermediate_size and image token budget for tiny Gemma4 #5760

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
albertvillanova merged 4 commits into from
May 15, 2026
Merged

Fix OOM in CI by reducing intermediate_size and image token budget for tiny Gemma4 #5760

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9d325ff
Reduce intermediate_size and image token budget for tinny gemma4
albertvillanova May 13, 2026
d89829c
Fix Gemma4Processor top-level image_seq_length
albertvillanova May 13, 2026
5133955
Simplify processor instantiation
albertvillanova May 13, 2026
e0971af
Revert "Simplify processor instantiation"
albertvillanova May 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions scripts/generate_tiny_models/for_conditional_generation/gemma4_for_conditional_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,20 +29,35 @@
processor = AutoProcessor.from_pretrained(MODEL_ID)
generation_config = GenerationConfig.from_pretrained(MODEL_ID)

# Gemma4 image processor uses aspect-ratio-preserving resizing, not a fixed image size. max_soft_tokens controls
# the output token budget and must be one of (70, 140, 280, 560, 1120). The smallest value (70) gives
# max_patches = 70 × pooling_kernel_size² = 70 × 9 = 630, so position_embedding_size must be at least 630.
# intermediate_size mirrors Gemma3: without it the production value (text: 6144, vision: 3072) is inherited, causing
# training activations [batch, patches, intermediate_size] to dominate GPU memory and OOM in CI.
IMAGE_TOKENS = 70 # minimum supported max_soft_tokens
MAX_PATCHES = IMAGE_TOKENS * 3**2 # 630

text_config = {
"num_hidden_layers": 2,
"hidden_size": 16,
"num_attention_heads": 4,
"num_key_value_heads": 2,
"intermediate_size": 32,
}
vision_config = {
"num_hidden_layers": 2,
"hidden_size": 16,
"num_attention_heads": 4,
"num_key_value_heads": 2,
"embed_dim": 64,
"intermediate_size": 32,
"position_embedding_size": MAX_PATCHES, # 630
"default_output_length": IMAGE_TOKENS, # 70
}

processor.image_processor.image_seq_length = IMAGE_TOKENS
Comment thread
cursor[bot] marked this conversation as resolved.
Outdated
processor.image_processor.max_soft_tokens = IMAGE_TOKENS

@cursor cursor Bot May 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Processor attributes leak into tokenizer_config.json

Low Severity

Setting processor.image_seq_length and processor.image_processor.max_soft_tokens before push_to_hub causes image_seq_length and max_soft_tokens to be serialized into tokenizer_config.json as well. These are image-processing parameters that don't belong in the tokenizer configuration. The PR reviewer confirmed that after this commit, the generated tokenizer_config.json contains "image_seq_length": 70 and "max_soft_tokens": 70, which is incorrect.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit e0971af. Configure here.

@albertvillanova albertvillanova May 13, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was the case after 5133955, but I have reverted it with e0971af.

I have checked that the generated tokenizer_config.json does not contain any image attributes: https://huggingface.co/trl-internal-testing/tiny-Gemma4ForConditionalGeneration/blob/refs%2Fpr%2F4/tokenizer_config.json


config = AutoConfig.from_pretrained(MODEL_ID)
for k, v in text_config.items():
setattr(config.text_config, k, v)
Expand Down
Loading