Add omni_checkpoint_stitcher and omni-gemma3-qwen3 configuration

[subawocit]

Description

This PR implements the multi-directory checkpoint loader. It loads pretrained parameter subtrees from two distinct checkpoints (e.g., a Gemma 3 vision tower and a Qwen 3 decoder as the initial setup) and merging them into a unified multimodal checkpoint.

Step 1 of a 5-step Proof-of-Concept for any-to-any multimodal alignment in MaxText. Overall goal is to align a pretrained vision encoders with another text-only LLM, and connecting the two with a dynamically configured adaptation layer (MLP connector), special tokenizer mapping for visual placeholder tokens, and train only the MLP using supervised fine tuning.

Files

• prepare_checkpoint.sh
  Main script: download original vision and language models' checkpoints from Hugging Face, convert to MaxText format, and stitch two checkpoints into a single multimodal checkpoint.
• omni-gemma3-qwen3.yml
  Configuration file combining the Qwen 3 4B language backbone with the Gemma 3 4B vision tower.
• utils/stitch_checkpoint.py
  Checkpoint stitcher utility that loads pretrained parameter subtrees from two separate checkpoints and merges them into a unified checkpoint.
• tests/stitch_checkpoint_test.py
  Unit tests.

Usage

HF_TOKEN="<your_token>" BASE_OUTPUT_DIRECTORY="<your_dir>/omni_checkpoints" ./prepare_checkpoint.sh

The checkpoints from the stitched model will then be saved in output directory.

Tests

Unit tests:

1. Layer Count Matching: Check if layer counts for the vision encoder blocks and language decoder blocks in the stitched checkpoint match the original source checkpoints.
2. Component-level Forward Pass Comparison:
   • Compare output logits from the vision encoder and language token embedder between the stitched checkpoint and original checkpoints. They must match exactly.
   • Compare output logits from the newly added vision projector in the stitched checkpoint with the original values. They must differ (given that the projector is randomly initialized in stitched model).

OMNI_TEST_BASE_DIR=<path_to_omni_checkpoints> pytest src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py

Unit Test Results

================================================================================
TEST 1: MODEL LAYER COUNT MATCHING
================================================================================
  - Verified LLM layers count (qwen3-4b): orig=36 and stitched=36
  - Verified Vision layers count (gemma3-4b): orig=27 and stitched=27
================================================================================

================================================================================
TEST 2: FORWARD LOGITS COMPARISON
================================================================================
orig vision embedding kernel shape: (14, 14, 3, 1152)
orig vision embedding kernel first 10 values: 
[-0.02490234 -0.00927734 -0.00124359  0.00150299  0.01397705 -0.00263977
  0.02539062 -0.01239014 -0.00058365  0.01257324]
stitched vision embedding kernel shape: (14, 14, 3, 1152)
stitched vision embedding kernel first 10 values: 
[-0.02490234 -0.00927734 -0.00124359  0.00150299  0.01397705 -0.00263977
  0.02539062 -0.01239014 -0.00058365  0.01257324]

orig llm token embedder shape: (151936, 2560)
orig llm token embedder first 10 values: 
[-0.02868652  0.01171875  0.01043701 -0.07177734 -0.03540039  0.00247192
 -0.05639648  0.00848389  0.00408936 -0.00668335]
stitched llm token embedder shape: (151936, 2560)
stitched llm token embedder first 10 values: 
[-0.02868652  0.01171875  0.01043701 -0.07177734 -0.03540039  0.00247192
 -0.05639648  0.00848389  0.00408936 -0.00668335]

orig vision projector weights first 10 values: 
[-0.00213623 -0.0703125   0.01025391 -0.02880859 -0.00473022  0.01574707
  0.03637695  0.04467773  0.00118256  0.01574707]
stitched vision projector weights first 10 values: 
[ 0.00064966 -0.00556512  0.00634653  0.01351321  0.00568023  0.00013807
 -0.0145482   0.0056528   0.00071098 -0.01255925]

  - Vision encoder output logits mean of abs: orig=0.09, stitched=0.09 (Mean abs diff = 0.00e+00, max = 0.00e+00)
  - Token embedder output logits mean of abs: orig=0.02, stitched=0.02 (Mean abs diff = 0.00e+00, max = 0.00e+00)
  - Projector output logits mean of abs:     orig=0.80, stitched=0.27 (Mean abs diff = 8.50e-01, max = 3.25e+01)
================================================================================

Step-by-Step Objectives

• √ Multi-Directory Checkpoint Restoration: Implemented selective sub-tree parameter loading using Orbax such that we can initiaze a multimodal model with parameters from multiple checkpoints.
• next Dynamic MLP Connector: Add omni modal adapter layer (MLP) to connect vision tower output to the LLM decoder
• next Special Tokenizer & Placeholder Masking: Add special token <|image|> to tokenizer and make sure the masking is correct in the decoder.
• COCO-Narratives Data Pipeline Processing
• SFT Execution & Evaluation

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 Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[hengtaoguo]

Thanks for the awesome groundwork!

[hengtaoguo]

This test is great! I would slightly prefer this file hierarchy, so later you won't need to create new folders for every new helper functions:

omni_poc/
├── tests/
│   └── stitch_test.py
├── utils/
│   └── stitch_checkpoint.py
├── prepare_checkpoint.py (maybe convert to bash instead?)
└── omni-gemma3-qwen3.yml

[subawocit]

Good idea - I've reorganized the files

[hengtaoguo]

Given its main job is converting checkpoints and calling stitch, would a simple Bash script be a cleaner fit here than Python? WDYT?

[subawocit]

SG!

[hengtaoguo]

Thanks! Since Gemma3 was the first multimodal model in MaxText, all vision-related flags defaulted to Gemma3 values in base.yml and types.py. Just to confirm, is that why you didn't need to add any multimodal configs here?

Even though the values would remain the same as the defaults, I'd highly recommend explicitly adding those specific vision configs below. It'll make this YAML file much more self-contained and easier to read.

[subawocit]

SG - I updated the yml and also did another test run to make sure the results are the same

[hengtaoguo]

Thanks for the great work!

[hengtaoguo]

nit: lower-case vision-language to align with text-only below?

[hengtaoguo]

One comment line above for something like # Multimodal configs for gemma3 vision encoder?

[subawocit]

I put Multimodal config for vision model for now, but may need to modify the name and parameter structure during step 3

[hengtaoguo]

Can you make this BASE_OUTPUT_DIRECTORY configurable in the usage example above? So later the user don't need to manually change any codes inside the bash script. Instead, they can easily set an env variable for the path.

ref:

maxtext/tests/end_to_end/tpu/deepseek/v3-671b/1_test_deepseek_mtp.sh

Lines 18 to 24 in fe5032a

	if [ -z "${BASE_OUTPUT_PATH}" ]; then
	# Non-Googlers please remember to point `BASE_OUTPUT_PATH` to GCS buckets that you own, this script uses internal buckets for testing.
	# this bucket will store all the files generated by MaxText during a run
	export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/$(date +%Y-%m-%d-%H-%M)
	echo "BASE_OUTPUT_PATH is not set"
	fi
	BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/}

[subawocit]

Updated!

[abhinavclemson]

Using a personal GCS bucket (yuchenhou-maxtext-logs) as the default fallback will cause automated CI runs or external developers to encounter permission errors if OMNI_TEST_BASE_DIR is not set in their environment.

Moreover it is not recommended, to use internal gcs bucket names in OOS repos.

[abhinavclemson]

You can use mock for unit tests.

[subawocit]

Output path: Good catch - I just updated the the code such that the out dir must be set manually

OMNI_TEST_BASE_DIR=<path_to_omni_checkpoints> pytest src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py

Mock for unit tests: Yes, good point, I had used mock models, and all the unit tests were passed. I later switched to use full checkpoints because we specifically needed to verify the vision and language output logits match the original source models exactly after stitching