Functioning decode on multimodal Gemma3-4b

[hengtaoguo]

Description

Insert the vision embeddings into text embeddings and allow a fully functioning decode forward pass on multimodal Gemma3-4b model.

• The most critical change here is that merge_mm_embeddings inserts the image_embeddings into the text_embeddings based on the image placeholder token information in the bidirectional_mask:

y = multimodal_utils.merge_mm_embeddings(
      text_embeddings=y,
      vision_embeddings=image_embeddings,
      mask=bidirectional_mask,
)

• We have a separate script golden_gemma3-4b-mm-export.py generates golden logits with fields (prompt, tokens, image_array, logits). This has been offline compared with MaxText logits with kl_div<0.01. We are not uploading this script and logits because of b/416290849. Later we need to modify MaxText/tests/forward_pass_logit_checker.py to make it accept full multimodal prompt for logits comparison (and add to end_to_end tests potentially).

Tests

A full decode forward pass command line, using prompt='Describe image <start_of_image> and image:

python -m MaxText.decode MaxText/configs/base.yml model_name=gemma3-4b tokenizer_path=assets/tokenizer.gemma3 load_parameters_path=gs://maxtext-model-checkpoints/gemma3-4b/multimodal/2025-04-25-18-06-04/checkpoints/0/items per_device_batch_size=1 run_name=ht_test max_prefill_predict_length=272 max_target_length=300 steps=1 async_checkpointing=false scan_layers=false use_multimodal=true prompt=\'Describe\ image\ \<start_of_image\>\' image_path=\'/home/hengtaoguo/projects/maxtext/MaxText/test_assets/test_image.jpg\' attention=\'dot_product\'

This yields the outcome image description logs:

RAMstats: After load_params:
        Using (GB) 31.82 / 1417.33 (2.245066%) -->  Available:1377.1
normalizer.cc(51) LOG(INFO) precompiled_charsmap is empty. use identity normalization.
Input `<start_of_turn>user
Describe image <start_of_image><end_of_turn>
<start_of_turn>model
` -> `Here's a description of the image:

**Overall Impression:** The image is a bright, expansive cityscape view of Seattle, Washington, with`

Checklist

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

• I have performed a self-review of my code.
• 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.

[aireenmei]

Great work, thanks Hengtao!

[aireenmei]

Nit: line 702 set image_embeddings to None. We can keep one of them. And we can use similar way for bidirectional mask to be consistent.

[hengtaoguo]

Done, just pass image_embeddings=image_embeddings to align with bidirectional_mask.

[aireenmei]

I think this will cause circular dependency, since I added this in BUILD file (https://source.corp.google.com/piper///depot/google3/third_party/py/maxtext/BUILD;l=360?q=BUILD%20f:maxtext). Since I only used some image related variables in gemma3.py (https://github.com/AI-Hypercomputer/maxtext/blob/main/MaxText/layers/gemma3.py#L48-L53), maybe better move them to multimodal_utils.

[hengtaoguo]

Thank you Aireen for the catch!

I have moved all the gemma-related static values from gemma3.py to multimodal_utils.py. Alongside this PR, I will amend in copybara to remove multimodal_utils.py's dependency on :layers. Let me know if this sounds good to you!

[aireenmei]

Sounds good, thank you!

[gobbleturk]

Awesome, thanks hengtao!!