Add custom MLP vision projector and tests by subawocit · Pull Request #4593 · AI-Hypercomputer/maxtext

subawocit · 2026-07-23T23:07:53Z

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_vision_projector".

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_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. Skipping deletion of 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`

Initializes a multimodal model (qwen3-vl-2b) with a customized MLP vision projector.
Checks output embeddings against standard configurations.

`tests/unit/model_creation_utils_test.py`

Tests the model creation utility with newly added customized projector parts.

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_vision_projector",
    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:

============================= 63 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:

same between the standard model and custom-projector model before the projector layer,
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_vision_projector). 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_vision_projector \
    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_vision_projector \
    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.
this 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 · 2026-07-23T23:07:56Z

Caution

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

codecov · 2026-07-23T23:16:06Z

Codecov Report

❌ Patch coverage is 64.91228% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/encoders.py	69.04%	11 Missing and 2 partials ⚠️
src/maxtext/common/checkpointing.py	40.00%	5 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!

gemini-code-assist · 2026-07-24T17:20:30Z

Caution

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

hengtaoguo · 2026-07-27T17:45:59Z

+    # If vision_projector_type is explicitly set to 'customized_mlp',
+    # override the model's default projector with a custom MLP-based projector.
+    vision_projector_type_config = getattr(self.config, "vision_projector_type", "default")
+    if vision_projector_type_config == "customized_mlp":


Maybe change this name to "customized_vision_projector"?

Because later we will have audio components and prefer explicit naming pattern.

Updated - I also changed the names in PR description accordingly

hengtaoguo · 2026-07-27T21:22:03Z

  without naming the weight.
  """
  problems = _weight_mismatches(want, have)
+  problems = [(p, why) for p, why in problems if not _is_custom_vision_projector_problem(p, want)]


Add a one-line comment explaining this line, such as "if the problem is caused by customized projector, ignore/remove this problem" if I understand correctly.

hengtaoguo · 2026-07-27T21:22:53Z

  return want, have


+def _is_custom_vision_projector_problem(path: str, want: dict) -> bool:


Rename to _is_custom_projector_problem to be more generic? (we could reuse for audio projector too in the future)

Agree - updated the code

hengtaoguo · 2026-07-27T21:24:11Z

+
+    cfg = _make_config(
+        enable_checkpointing=True,
+        load_parameters_path="gs://fake/custom_vision_ckpt",


What is this "gs://fake/custom_vision_ckpt"? If it's actually unused, can we possibly remove it from the config?

It is a dummy path required to test checkpoint loading related functions. The same flag was passed in many other unit tests (e.g., L691). I was wondering why this is added throughout the unit tests so I added some print statements in the _free_device_memory during debugging, and without this load_parameters_path="gs://fake/custom_vision_ckpt", _free_device_memory never gets executed in unit tests.

maybe use enable_checkpointing=False?

hengtaoguo · 2026-07-27T21:28:30Z

    return embeddings, deep_feats


+class MultimodalMLPProjector(nnx.Module):


I like this customizable block! Is it possible to reuse for audio tower in the future too? Asking because you named it Multimodal instead of Vision, which is great.

Yes, one just need to 1. define the input/output dimensions based on the audio config (L146), 2. define the connector hyperparameters (similar to L150-155), and 3. wire it up in AudioEncoder._setup_audio_encoder_layers

This should be very straightforward

hengtaoguo · 2026-07-27T21:30:44Z

        return node

-      jax.tree_util.tree_map(_free_device_memory, sharded_state, is_leaf=lambda n: isinstance(n, nnx.Variable))
+      jax.tree_util.tree_map_with_path(_free_device_memory, sharded_state, is_leaf=lambda n: isinstance(n, nnx.Variable))


Why is this change needed?

Standard tree_map() passes only the raw tensor (node). A tensor has no name so it's very difficult to distinguish a base model weight from a custom projector weight. We want to identify the custom projector layer so it doesnt get deleted during free memory. So I switched to tree_map_with_path() which passes pytree keys (like "custom_linear_0"), and I can filter by the name, skip freeing memory in custom_linear, and keep it randomly initialized

reference: https://docs.jax.dev/en/latest/_autosummary/jax.tree.map_with_path.html

NuojCheng · 2026-07-29T16:35:12Z

+    proj_name = parts[1]
+    proj_dict = want.get("vision_encoder", {}).get(proj_name, {})
+    if isinstance(proj_dict, dict) and any("custom_linear" in k for k in proj_dict.keys()):
+      print(


use maxtext logging

NuojCheng · 2026-07-29T16:38:45Z

+import jax.numpy as jnp
+
+# Ensure the parent directory (src/) is in sys.path so we can import maxtext
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")))


I don't think we need this. Is the maxtext installed from source properly? Follow here

NuojCheng · 2026-07-29T16:48:46Z

+      num_blocks = seq_len // self.tokens_per_block
+      x = x.reshape((batch_size, num_blocks, self.tokens_per_block * in_dim))
+
+    for i in range(self.num_layers):


in JAX we usually use jax.lax.scan instead of for-loop for faster compilation and potentially better performance

NuojCheng · 2026-07-29T16:52:39Z

+        use_multimodal=True,
+        model_name="qwen3-vl-2b",
+        override_model_config=True,
+        scan_layers=False,


why the scan is not used?

Add custom MLP vision projector and tests

c908c8e

add unit test

6cc1f5f

subawocit marked this pull request as ready for review July 24, 2026 17:20

subawocit requested review from NuojCheng, RissyRan, bvandermoon, gagika, gobbleturk, huytransformer, jiangjy1982, parambole, richjames0, shralex, shuningjin, suexu1025 and xibinliu as code owners July 24, 2026 17:20

subawocit requested review from A9isha, SurbhiJainUSC, abhinavclemson, aireenmei, darisoy, dipannita08, hengtaoguo, igorts-git, khatwanimohit and vipannalla as code owners July 24, 2026 17:20

hengtaoguo requested changes Jul 27, 2026

View reviewed changes

subawocit added 2 commits July 28, 2026 00:33

update

5145658

better description

cfa5753

NuojCheng reviewed Jul 29, 2026

View reviewed changes

		return want, have


		def _is_custom_vision_projector_problem(path: str, want: dict) -> bool:

		return embeddings, deep_feats


		class MultimodalMLPProjector(nnx.Module):

Uh oh!

Conversation

subawocit commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes

Custom MLP Projector Module

src/maxtext/layers/encoders.py

Config

src/maxtext/configs/types.py

Checkpoint Loading & Weight Mismatch Handling

src/maxtext/common/checkpointing.py

src/maxtext/utils/model_creation_utils.py

Test Script

src/maxtext/experimental/omni_poc/tests/custom_vision_projector_test.py

tests/unit/model_creation_utils_test.py

Example Usage

Tests

1. Unit Tests

2. Custom Vision Projector Verification

3. End-to-End Multimodal Decoding Checks

A. Baseline Multimodal Run (Default Vision Projector)

B. Custom Vision Projector Pipeline Run

4. Text-Only Decoding Regression Check

Step-by-Step Objectives

Checklist

Uh oh!

gemini-code-assist Bot commented Jul 23, 2026

Uh oh!

codecov Bot commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

gemini-code-assist Bot commented Jul 24, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

subawocit Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

subawocit Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

subawocit Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

subawocit Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

subawocit commented Jul 23, 2026 •

edited

Loading

`src/maxtext/layers/encoders.py`

`src/maxtext/configs/types.py`

`src/maxtext/common/checkpointing.py`

`src/maxtext/utils/model_creation_utils.py`

`src/maxtext/experimental/omni_poc/tests/custom_vision_projector_test.py`

`tests/unit/model_creation_utils_test.py`

codecov Bot commented Jul 23, 2026 •

edited

Loading

subawocit Jul 28, 2026 •

edited

Loading

subawocit Jul 28, 2026 •

edited

Loading

subawocit Jul 28, 2026 •

edited

Loading

subawocit Jul 27, 2026 •

edited

Loading