custom processor for omni poc

[subawocit]

Description

This PR introduces the implementations for hybrid multimodal processor.

Specifically, it implements modular hybrid multimodal support in the preprocessing and fusion pipelines:

• Text Prompt: Formatted with text decoder (Qwen 3) chat template and special vision tags (<|vision_start|>, <|image_pad|>, <|vision_end|>) and tokenized by the text decoder's tokenizer.
• Image Inputs: Preprocessed by the vision encoder (Gemma 3) processor.
• Fusion Stage: Handled by a custom processor that expands the placeholder tokens into 256 copies of <|image_pad|>. This ensures the text decoder knows precisely where to inject the 256 projected visual embeddings within the token stream without corrupting text subword tokenization or array indexing.

This is Step 3 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 (customized MLP connector), special tokenizer mapping for visual placeholder tokens, and train only the MLP using supervised fine tuning.

Files

1. Omni Multimodal Processor

   • experimental/omni_poc/utils/processor_omni_gemma3_qwen3.py
     • Implemented visual placeholder token management (<|vision_start|>, <|image_pad|>, <|vision_end|>) combining Qwen 3's vocabulary with Gemma 3's visual token (256 tokens per image).
     • Added functions for offset calculation and placeholder token expansion (get_image_offsets_omni, add_extra_tokens_for_omni).

2. Name Registration & Routing

   • multimodal/processor.py, layers/decoders.py, and layers/nnx_decoders.py
     • Integrated stitched architecture routing into processor.py for models with vision_block == "gemma3" and decoder_block == "qwen3".
     • Registered omni-gemma3-qwen3 model name.

3. Unit Tests

   • experimental/omni_poc/tests/processor_omni_gemma3_qwen3_test.py
     • Added unit tests for Omni visual token expansion, offset calculations, and dtype preservation.
   • tests/unit/multimodal_utils_test.py
     • Added unit tests for omni-gemma3-qwen3 routing and fusion masking in the multimodal processor.

4. End-to-End Decoding

   • experimental/omni_poc/utils/decode_omni.py
     • Added a decoding script to run prefill and autoregressive generation on ChartQA multimodal validation samples.

5. Checkpoint Stitching (Minor Change)

   • experimental/omni_poc/utils/stitch_checkpoint.py
     • Modified the checkpoint stitch script to be compatible with current codebase.
   • experimental/omni_poc/tests/stitch_checkpoint_test.py
     • Modified unit tests for checkpoint stitching to be compatible with current codebase.

6. Configuration (Minor Change)

   • experimental/omni_poc/omni-gemma3-qwen3.yml
     • Modified model configuration to separately define Gemma 3 vision tower parameters, Qwen 3 4B LLM backbone parameters, and custom projector.

Tests

1. Checkpoint Stitching Generation

Generated a stitched Orbax checkpoint from converted Gemma 3 4B vision weights and Qwen 3 4B language weights:

JAX_PLATFORMS=cpu python -m maxtext.experimental.omni_poc.utils.stitch_checkpoint \
    src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml \
    --vision_load_path=gs://yuchenhou-maxtext-logs/omni_checkpoints/gemma3-4b_converted/0/items \
    --llm_load_path=gs://yuchenhou-maxtext-logs/omni_checkpoints/qwen3-4b_converted/0/items \
    --stitched_output_path=gs://yuchenhou-maxtext-logs/omni_checkpoints/omni_stitched_gemma3-4b_qwen3-4b/0/items

2. Checkpoint Stitching Unit Test

Verified that the stitched checkpoint accurately preserves original weights from both source models:

OMNI_TEST_BASE_DIR='gs://yuchenhou-maxtext-logs/omni_checkpoints' JAX_PLATFORMS=cpu \
    /home/yuchenhou_google_com/maxtext_env/bin/python -m unittest src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py

Output:

----------------------------------------------------------------------
Ran 2 tests in 135.804s

OK

3. Omni Processor Unit Test

Verified token ID mappings, sequence expansions, and offsets:

python -m unittest src/maxtext/experimental/omni_poc/tests/processor_omni_gemma3_qwen3_test.py

Output:

----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK

4. Multimodal Processor Unit Tests

Verified hybrid multimodal routing and fusion:

JAX_PLATFORMS=cpu python3 -m pytest -v tests/unit/multimodal_utils_test.py tests/unit/qwen3_omni_layers_test.py

Output:

============================= 56 passed, 1 skipped in 192.32s =============================

5. End-to-End Multimodal Decoding

Ran end-to-end decoding on ChartQA validation samples using the stitched checkpoint:

JAX_PLATFORMS=tpu /home/yuchenhou_google_com/maxtext_env/bin/python src/maxtext/experimental/omni_poc/utils/decode_omni.py \
    --checkpoint_path="gs://yuchenhou-maxtext-logs/omni_checkpoints/omni_stitched_gemma3-4b_qwen3-4b/0/items" \
    --num_samples=2 \
    --max_new_tokens=64

Output:

Sample 1 (Index 1309):
  Question: What was the average annual total income per household of those in the top decile group?
  Ground Truth: ['186600']
  Model Response: Okay, the user is asking for the average annual total income per household in the top decile group. First, I need to figure out what data they're referring to. The question mentions "top decile group," which typically refers to the highest 10% of income earners. But the user hasn

Sample 2 (Index 228):
  Question: Which gender is represented by the red color line?
  Ground Truth: ['Male']
  Model Response: Okay, the user is asking about the gender represented by the red color line in a given text. Let me start by understanding the context. The text provided seems to be a mix of Chinese and English characters, possibly with some typos or errors. The user is asking which gender is represented by the red color

This ensures the data is flowing end to end. The response itself does not make sense due to the randomly initialized projector, as expected.

Step-by-Step Objectives

• PR 4485 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.
• PR 4593 Dynamic MLP Connector: Add omni modal adapter layer (MLP) to connect vision tower output to the LLM decoder
• PR 4746 and Current 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.

[gemini-code-assist]

Code Review

This pull request introduces support for the 'omni-gemma3-qwen3' stitched model, combining the Gemma 3 vision encoder with the Qwen 3 LLM decoder. It includes configuration updates, routing logic in the multimodal processor, a decoding evaluation script, and corresponding unit tests. The review feedback highlights a 2-token discrepancy in the image offset calculation that needs alignment with the token expansion logic, along with a performance optimization in the decoding loop to use jnp.argmax instead of np.argmax to avoid costly device-to-host transfers.

[github-actions]

🤖 Hi @subawocit, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request introduces a custom multimodal processor and evaluation script for "omni-gemma3-qwen3," a stitched pilot model combining a Gemma 3 Vision Encoder with a Qwen 3 LLM Decoder. The implementation is of very high quality, with robust unit testing, flexible routing, and extremely clean integration with MaxText's layers and checkpointers.

🔍 General Feedback

• Component Design: The routing and preprocessing are designed elegantly, utilizing local/dynamic imports (import-outside-toplevel) in the generic multimodal processor, which cleanly isolates the experimental model pilot from the core codebase.
• Robustness of Stitching Tests: The checkpoint stitching tests are very comprehensive and cover both structural correctness (layer counts) and parameter-restore verification.
• Evaluation Tooling: The inclusion of decode_omni.py as a standalone verification and evaluation tool on the ChartQA dataset is excellent for model validation.

[codecov]

Codecov Report

❌ Patch coverage is 64.28571% with 5 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/multimodal/processor.py	64.28%	2 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!