Skip to content

Fix Gemma-3 4B Multimodal SFT test script#4268

Merged
copybara-service[bot] merged 1 commit into
mainfrom
fix-gemma3-multimodal-sft-tokenizer
Jun 25, 2026
Merged

Fix Gemma-3 4B Multimodal SFT test script#4268
copybara-service[bot] merged 1 commit into
mainfrom
fix-gemma3-multimodal-sft-tokenizer

Conversation

@RexBearIU

@RexBearIU RexBearIU commented Jun 25, 2026

Copy link
Copy Markdown
Collaborator

Description

This PR fixes multiple bugs and configuration issues in Step 4 of the Gemma-3 4B Multimodal SFT end-to-end test script:

  1. Fixes Syntax Typo: Corrects the typo run_name=${run_id}} (with a trailing curly bracket) to run_name=${run_id}.
  2. Adds Missing Tokenizer Configuration: Adds tokenizer_type=huggingface to Step 4. Without this, MaxText defaults to a SentencePiece tokenizer and attempts to load a local .model file at the Hugging Face repo path google/gemma-3-4b-it, resulting in the following error:
    tensorflow.python.framework.errors_impl.NotFoundError: google/gemma-3-4b-it; No such file or directory
    

Description of the issue we found

Without the tokenizer_type=huggingface flag, Step 4 of the end-to-end Gemma-3 SFT test crashes upon parameter loading when it tries to find a local SentencePiece model path under google/gemma-3-4b-it instead of fetching the proper Hugging Face config.

Tests

We tested this change end-to-end on a v6e TPU VM:

  1. Re-ran the end-to-end script tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh using a unique run ID.
  2. Verified that Step 2 successfully executed inference.
  3. Verified that Step 3 completed all 5 SFT training steps and successfully saved the checkpoints to GCS.
  4. Verified that Step 4 successfully completed loading parameters from the newly trained SFT checkpoint on GCS, initialized the Hugging Face tokenizer, and correctly decoded/described the image:
    Input `<start_of_turn>user
    Describe image <img><end_of_turn>
    <start_of_turn>model
    ` -> `Seattle skyline`
    

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@codecov

codecov Bot commented Jun 25, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@RexBearIU RexBearIU force-pushed the fix-gemma3-multimodal-sft-tokenizer branch from a4e1ea6 to d2658b8 Compare June 25, 2026 08:36
@RexBearIU RexBearIU force-pushed the fix-gemma3-multimodal-sft-tokenizer branch from d2658b8 to fe4cb35 Compare June 25, 2026 08:51
@RexBearIU RexBearIU changed the title Fix Gemma-3 4B Multimodal SFT test script Step 4 decode inference Fix Gemma-3 4B Multimodal SFT test script Jun 25, 2026
@copybara-service copybara-service Bot merged commit 94f60c3 into main Jun 25, 2026
49 checks passed
@copybara-service copybara-service Bot deleted the fix-gemma3-multimodal-sft-tokenizer branch June 25, 2026 18:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants