Add multi-directory checkpoint loader for multimodal alignment poc

[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.py
  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.
• checkpoint_stitcher/stitch.py
  Checkpoint stitcher utility that loads pretrained parameter subtrees from two separate checkpoints and merges them into a unified checkpoint.
• checkpoint_stitcher/unittest_stitch.py
  Unit tests.

Usage

Define hf_token and base_output_directory in prepare_checkpoint.py, then:

python3 -m maxtext.experimental.omni_poc.prepare_checkpoint

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).

$ pytest src/maxtext/experimental/omni_poc/checkpoint_stitcher/unittest_stitch.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
• 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!