Add custom MLP vision projector and tests

[subawocit]

Description

This PR adds support for configurable custom vision projectors in MaxText's multimodal models. It introduces new vision projector config options, handles checkpoint weight mismatches when attaching custom projector architectures to pre-trained base models, and adds a test script to verify projector initialization and embedding extraction.

Step 2 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.

Changes

Custom MLP Projector Module (src/maxtext/layers/encoders.py)

• Implemented MultimodalMLPProjector (nnx.Module), supporting configurable layer, activations, hidden size, and layer biases.
• Updated VisionEncoder._setup_vision_encoder_layers() to conditionally override the default vision projector with MultimodalMLPProjector whenever vision_projector_type == "customized_mlp".

Config (src/maxtext/configs/types.py)

• Extended the VisionProjector configuration class for customizing projector architecture.

Checkpoint Loading & Weight Mismatch Handling

src/maxtext/common/checkpointing.py

• Added _is_custom_vision_projector_problem() to check if missing/mismatched weight keys belong to custom vision projector layers.
• Updated _raise_on_weight_mismatch() to suppress weight mismatch errors for custom projector layers only, and raise a warning instead. This allows restoring base model parameters from standard checkpoints while leaving newly introduced custom projector weights randomly initialized.

src/maxtext/utils/model_creation_utils.py

• Modified _free_device_memory inside from_pretrained().
• Added path checks (is_custom) to skip deleting JAX arrays corresponding to custom layers (i.e., layers with "custom_linear" in their name). Reason: Pretrained base model checkpoints do not contain weights for new custom components; if their device memory were freed prior to checkpoint restoration, their initializations would be deleted with no incoming checkpoint weights to replace them. Preserving these arrays ensures their randomly initialized state remains alive in device memory when base model parameters are loaded.

Test Script (src/maxtext/experimental/omni_poc/tests/custom_vision_projector_test.py)

• Added a test script that:
  • Initializes a multimodal model (qwen3-vl-2b) with a customized MLP vision projector.
  • Checks output embeddings against standard configurations.

Example Usage

    config = pyconfig.initialize(
        model_name="qwen3-vl-2b",
        tokenizer_path="Qwen/Qwen3-VL-2B-Instruct",
        tokenizer_type="huggingface",
        load_parameters_path=<PATH_TO_QWEN3_CHECKPOINT>,
        use_multimodal=True,
        prompt="Describe this image",
        image_path="tests/assets/test_image.jpg",
        override_model_config=True,
        hf_access_token=<HF_ACCESS_TOKEN>,
        vision_projector_type="customized_mlp",
        vision_connector_num_layers=2,
        vision_connector_activation="gelu",
        vision_connector_use_bias=True,
        ... # other standard multimodal arguments
    )

Tests

1. Unit Tests

Command:

pytest tests/unit/model_creation_utils_test.py

Output:

============================= 62 passed in 45.95s ==============================

Command:

pytest tests/unit/checkpointing_nnx_load_test.py

Output:

=============================== 24 passed in 6.82s ===============================

2. Custom Vision Projector Verification

Validates that vision embeddings are the same between the standard model and
custom-projector model before the projector layer,
but differ after the projector layer due to randomly initialized projector.

Command:

python3 src/maxtext/experimental/omni_poc/tests/custom_vision_projector_test.py \
    --load_parameters_path=<your-path> \
    --hf_access_token=<your-token>

Output:

Image pixel_values shape: (1, 3, 2, 544, 640)
vision embeddings: (1, 1360, 1024)
First 10 values:
[-5.091886    1.7075098  -1.0072641  -0.16636705 -1.4430853  -5.3355365
  1.4629428   0.21186987 -2.9909272  -4.376259  ]
projector_embeddings: (1, 340, 2048)
First 10 values:
[  1.3062916   8.224888   16.922318   -7.4536076  -8.535779   33.199806
   3.0147345   2.200746  -18.904125    9.107975 ]

Original vision embeddings: (1, 1360, 1024)
First 10 values:
[-5.091886    1.7075098  -1.0072641  -0.16636705 -1.4430853  -5.3355365
  1.4629428   0.21186987 -2.9909272  -4.376259  ]
Original projector_embeddings: (1, 340, 2048)
First 10 values:
[-0.31920713  0.54748684  0.49574944 -0.88344455 -0.23951492 -0.14456059
  0.1383374   0.05032197  0.381779    0.00628588]

===================================
Comparing customized vs original embeddings
Vision embeddings match: True (Max diff: 0.0)
Projector embeddings match: False (Max diff: 313.9195556640625)

3. End-to-End Multimodal Decoding Checks

A. Baseline Multimodal Run (Default Vision Projector)

Verifies standard multimodal inference functions normally.

Command:

python3 -m maxtext.inference.decode \
    src/maxtext/configs/base.yml \
    model_name=qwen3-vl-2b \
    tokenizer_path=Qwen/Qwen3-VL-2B-Instruct \
    tokenizer_type=huggingface \
    load_parameters_path=<your-path> \
    per_device_batch_size=1 \
    scan_layers=false \
    use_multimodal=true \
    prompt='Describe this image' \
    image_path='tests/assets/test_image.jpg' \
    max_prefill_predict_length=512 \
    max_target_length=768 \
    ici_tensor_parallelism=4 \
    override_model_config=true \
    attention='dot_product' \
    hf_access_token=<your-token>

Output:

Input `<|im_start|>user
<|vision_start|><|image_pad|><|vision_end|>Describe this image<|im_end|>
<|im_start|>assistant
` -> `This is a panoramic view of the Seattle skyline on a bright, sunny day. The image is taken from a high vantage point, looking down on the city.

The most prominent feature is the **Space Needle**, a distinctive observation tower located on the left side of the frame. It stands tall among the city's modern skyscrapers, which are a mix of glass and steel structures. The buildings are densely packed, creating a dense urban landscape.

In the background, a range of mountains is visible, with their peaks covered in snow, suggesting a high elevation. The sky is a clear, vibrant blue with a few scattered white clouds.

In the foreground, there are lush green trees and foliage, which add a natural element to the urban scene. The overall impression is one of a vibrant, modern city with a beautiful natural backdrop.

B. Custom Vision Projector Pipeline Run

Verifies end-to-end data flow with custom projector flags (customized_mlp). Output text is bad due to random weight initialization, as expected.

Command:

python3 -m maxtext.inference.decode \
    src/maxtext/configs/base.yml \
    model_name=qwen3-vl-2b \
    tokenizer_path=Qwen/Qwen3-VL-2B-Instruct \
    tokenizer_type=huggingface \
    load_parameters_path=<your-path> \
    per_device_batch_size=1 \
    run_name=runner_image_2026-06-26-20-40 \
    scan_layers=false \
    use_multimodal=true \
    prompt='Describe this image' \
    image_path='tests/assets/test_image.jpg' \
    max_prefill_predict_length=512 \
    max_target_length=768 \
    ici_tensor_parallelism=4 \
    override_model_config=true \
    attention='dot_product' \
    hf_access_token=<your-token> \
    vision_projector_type=customized_mlp \
    vision_connector_num_layers=2 \
    vision_connector_activation=gelu \
    vision_connector_use_bias=true

Output:

Input `<|im_start|>user
<|vision_start|><|image_pad|><|vision_end|>Describe this image<|im_end|>
<|im_start|>assistant
` -> `The image is a detailed and intricate depiction of a complex, multi-layered, and highly detailed scene that appears to be a digital or conceptual representation rather than a physical photograph...

4. Text-Only Decoding Regression Check

Verifies custom vision projector configurations do not affect text-only inference mode (use_multimodal=false).

Command:

python3 -m maxtext.inference.decode \
    src/maxtext/configs/base.yml \
    model_name=qwen3-vl-2b \
    tokenizer_path=Qwen/Qwen3-VL-2B-Instruct \
    tokenizer_type=huggingface \
    load_parameters_path=<your-path> \
    per_device_batch_size=1 \
    max_prefill_predict_length=128 \
    max_target_length=256 \
    scan_layers=false \
    use_multimodal=false \
    prompt='How many times does the letter r appear in the word Strawberry?' \
    override_model_config=true \
    ici_fsdp_parallelism=1 \
    ici_tensor_parallelism=-1 \
    hf_access_token=<your-token> \
    vision_projector_type=customized_mlp \
    vision_connector_num_layers=2 \
    vision_connector_activation=gelu \
    vision_connector_use_bias=true

Output:

Input `How many times does the letter r appear in the word Strawberry?` -> ` To determine how many times the letter **r** appears in the word **Strawberry**, let's go step by step.

1. Write down the word: **Strawberry**
2. Break it down into individual letters:
   - S - t - r - a - w - b - e - r - r - y
3. Now, count how many times the letter **r** appears:
   - The first **r** is in the 3rd position.
   - The second **r** is in the 8`

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.
• √ 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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 20.96774% with 49 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/encoders.py	9.52%	38 Missing ⚠️
src/maxtext/common/checkpointing.py	40.00%	5 Missing and 1 partial ⚠️
...rc/maxtext/checkpoint_conversion/to_huggingface.py	0.00%	2 Missing ⚠️
src/maxtext/layers/quantizations.py	33.33%	1 Missing and 1 partial ⚠️
src/maxtext/utils/model_creation_utils.py	80.00%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!