Add custom MLP vision projector and tests

src/maxtext/common/checkpointing.py

@@ -88,6 +88,22 @@ def _expected_and_restored_params(abstract_nnx_state, restored_linen):
   return want, have
 
 
+def _is_custom_vision_projector_problem(path: str, want: dict) -> bool:
+  """Returns True if a weight mismatch belongs to a newly attached custom vision projector."""
+  parts = [p for p in path.replace(".", "/").split("/") if p and p != "params"]
+  if len(parts) >= 2 and parts[0] == "vision_encoder":
+    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(
+          f"===Warning: weight mismatch found in custom vision projector: {proj_name}.\n"
+          f"Path: {path}\n"
+          "This custom vision projector will be initialized with random weights.==="
+      )
+      return True
+  return False
+
+
 def _raise_on_weight_mismatch(want, have):
   """Raises if the restored weights (`have`) don't match what the model expects (`want`).
 
@@ -98,6 +114,7 @@ def _raise_on_weight_mismatch(want, have):
   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)]
   if not problems:
     return
   lines = "\n".join(f"  - '{p}': {why}" for p, why in problems)

src/maxtext/configs/types.py

@@ -2147,6 +2147,16 @@ class VisionProjector(BaseModel):
   projector_output_dim_for_vit: int = Field(4096, description="Output dimension for the vision projector.")
   pixel_shuffle_ratio_for_vit: float = Field(0.5, description="Pixel shuffle ratio for the Vision Transformer.")
   projector_dropout_for_vit: float = Field(0.0, description="Dropout rate for the vision projector.")
+  vision_projector_type: str = Field(
+      "default", description="Type of the vision projector to use. Supported: 'default', 'customized_mlp'."
+  )
+  vision_connector_num_layers: int = Field(2, description="Number of layers in custom mlp projector.")
+  vision_connector_hidden_size: int = Field(
+      0,
+      description=("Hidden size for custom mlp projector intermediate layers. 0 defaults to LLM hidden size."),
+  )
+  vision_connector_activation: str = Field("gelu", description="Activation function for custom mlp projector.")
+  vision_connector_use_bias: bool = Field(True, description="Whether to use bias in custom mlp projector layers.")
 
 
 class AudioEncoder(BaseModel):

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

@@ -0,0 +1,155 @@
+"""Tests for custom vision projector implementation in MaxText.
+
+"""
+
+import argparse
+import gc
+import os
+import sys
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+
+# Ensure the parent directory is in sys.path so we can import maxtext
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+from maxtext.configs import pyconfig
+from maxtext.inference.maxengine import maxengine
+from maxtext.multimodal import processor as mm_processor
+
+
+def test_load_custom_vision_projector(hf_access_token, param_path):
+  base_yml_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../configs/base.yml"))
+
+  config = pyconfig.initialize(
+      base_config=base_yml_path,
+      model_name="qwen3-vl-2b",
+      tokenizer_path="Qwen/Qwen3-VL-2B-Instruct",
+      tokenizer_type="huggingface",
+      load_parameters_path=param_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=hf_access_token,
+      vision_projector_type="customized_mlp",
+      vision_connector_num_layers=2,
+      vision_connector_activation="gelu",
+      vision_connector_use_bias=True,
+  )
+
+  print("Customized model")
+  print(f"Active Vision Projector Type: {config.vision_projector_type}")
+  print(f"Num Layers: {config.vision_connector_num_layers}")
+
+  engine = maxengine.MaxEngine(config)
+
+  # Load the parameters from the checkpoint
+  rng = jax.random.PRNGKey(1234)
+  params = engine.load_params(rng=rng)
+  nnx.update(engine.model, params)
+
+  # Preprocess the example image using the provided image_path in the config
+  processor_outputs = mm_processor.preprocess_mm_data(config)
+
+  images = processor_outputs.pixel_values
+  print(f"Image pixel_values shape: {images.shape}")
+
+  # Extract the vision wrapper and its underlying parts
+  vision_wrapper = engine.model.vision_encoder
+  pure_encoder = getattr(vision_wrapper, vision_wrapper.encoder_name)
+  projector = getattr(vision_wrapper, vision_wrapper.projector_name)
+
+  # Pass the image through the pure vision tower (before projector)
+  encoder_output = pure_encoder(images)
+
+  vision_embeddings_customized = encoder_output[0] if isinstance(encoder_output, tuple) else encoder_output
+
+  print("vision embeddings:", vision_embeddings_customized.shape)
+  print(f"First 10 values:\n{vision_embeddings_customized.flatten()[:10]}")
+
+  # Pass the image through the vision tower (including projector)
+  projected_embeddings_customized = projector(vision_embeddings_customized)
+
+  print(f"projector_embeddings: {projected_embeddings_customized.shape}")
+  print(f"First 10 values:\n{projected_embeddings_customized.flatten()[:10]}")
+
+  del engine, params, vision_wrapper, pure_encoder, projector
+  gc.collect()
+
+  print("\nOriginal model")
+  config_original = pyconfig.initialize(
+      base_config=base_yml_path,
+      model_name="qwen3-vl-2b",
+      tokenizer_path="Qwen/Qwen3-VL-2B-Instruct",
+      tokenizer_type="huggingface",
+      load_parameters_path=param_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=hf_access_token,
+  )
+
+  print(f"Original Vision Projector Type: {getattr(config_original, 'vision_projector_type', None)}")
+
+  engine_original = maxengine.MaxEngine(config_original)
+
+  params_original = engine_original.load_params(rng=rng)
+  nnx.update(engine_original.model, params_original)
+
+  processor_outputs_original = mm_processor.preprocess_mm_data(config_original)
+  images_original = processor_outputs_original.pixel_values
+
+  vision_wrapper_original = engine_original.model.vision_encoder
+  pure_encoder_original = getattr(vision_wrapper_original, vision_wrapper_original.encoder_name)
+  projector_original = getattr(vision_wrapper_original, vision_wrapper_original.projector_name)
+
+  encoder_output_original = pure_encoder_original(images_original)
+
+  vision_embeddings_original = (
+      encoder_output_original[0] if isinstance(encoder_output_original, tuple) else encoder_output_original
+  )
+
+  print("Original vision embeddings:", vision_embeddings_original.shape)
+  print(f"First 10 values:\n{vision_embeddings_original.flatten()[:10]}")
+
+  projected_embeddings_original = projector_original(vision_embeddings_original)
+
+  print(f"Original projector_embeddings: {projected_embeddings_original.shape}")
+  print(f"First 10 values:\n{projected_embeddings_original.flatten()[:10]}")
+
+  print("\n===================================")
+  print("Comparing customized vs original embeddings")
+  vision_embeddings_match = jnp.allclose(vision_embeddings_customized, vision_embeddings_original, atol=1e-5)
+  vision_max_diff = jnp.max(jnp.abs(vision_embeddings_customized - vision_embeddings_original))
+  print(f"Vision embeddings match: {vision_embeddings_match} (Max diff: {vision_max_diff})")
+
+  projector_embeddings_match = jnp.allclose(projected_embeddings_customized, projected_embeddings_original, atol=1e-5)
+  projector_max_diff = jnp.max(jnp.abs(projected_embeddings_customized - projected_embeddings_original))
+  print(f"Projector embeddings match: {projector_embeddings_match} (Max diff: {projector_max_diff})")
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--load_parameters_path", required=True)
+  parser.add_argument("--hf_access_token", default=os.getenv("HF_TOKEN", ""))
+
+  args = parser.parse_args()
+  test_load_custom_vision_projector(
+      hf_access_token=args.hf_access_token,
+      param_path=args.load_parameters_path,
+  )

src/maxtext/layers/encoders.py

@@ -21,6 +21,7 @@
 from maxtext.common.common_types import Config, VisionEncoderBlockType
 from maxtext.layers import nnx_wrappers
 from maxtext.layers import initializers
+from maxtext.layers import linears
 
 
 class VisionEncoder(nnx.Module):
@@ -43,23 +44,20 @@ def _setup_vision_encoder_layers(self):
       projector_name = "VisionEmbedder_0"
       setattr(self, encoder_name, gemma3.Gemma3VisionEncoderLayer(config=self.config, mesh=self.mesh, rngs=self.rngs))
       setattr(self, projector_name, gemma3.VisionEmbedder(config=self.config, mesh=self.mesh, rngs=self.rngs))
-      return encoder_name, projector_name
     elif self.vision_encoder_block == VisionEncoderBlockType.LLAMA4:
       from maxtext.models import llama4  # pylint: disable=import-outside-toplevel
 
       encoder_name = "Llama4VisionModel_0"
       projector_name = "Llama4MultiModalProjector_0"
       setattr(self, encoder_name, llama4.Llama4VisionModel(config=self.config, mesh=self.mesh, rngs=self.rngs))
       setattr(self, projector_name, llama4.Llama4MultiModalProjector(config=self.config, mesh=self.mesh, rngs=self.rngs))
-      return encoder_name, projector_name
     elif self.vision_encoder_block == VisionEncoderBlockType.QWEN3_OMNI:
       from maxtext.models import qwen3  # pylint: disable=import-outside-toplevel
 
       encoder_name = "Qwen3OmniMoeVisionEncoder_0"
       projector_name = "Qwen3OmniMoeVisionProjector_0"
       setattr(self, encoder_name, qwen3.Qwen3OmniMoeVisionEncoder(config=self.config, mesh=self.mesh, rngs=self.rngs))
       setattr(self, projector_name, qwen3.Qwen3OmniMoeVisionProjector(config=self.config, rngs=self.rngs))
-      return encoder_name, projector_name
     elif self.vision_encoder_block == VisionEncoderBlockType.GEMMA4:
       from maxtext.models import gemma4_vision  # pylint: disable=import-outside-toplevel
 
@@ -71,7 +69,6 @@ def _setup_vision_encoder_layers(self):
       setattr(
           self, projector_name, gemma4_vision.Gemma4VisionProjector(config=self.config, mesh=self.mesh, rngs=self.rngs)
       )
-      return encoder_name, projector_name
     elif self.vision_encoder_block == VisionEncoderBlockType.QWEN3_5:
       from maxtext.models import qwen3_5_vision  # pylint: disable=import-outside-toplevel
 
@@ -81,7 +78,6 @@ def _setup_vision_encoder_layers(self):
           self, encoder_name, qwen3_5_vision.Qwen3_5MoeVisionEncoder(config=self.config, mesh=self.mesh, rngs=self.rngs)
       )
       setattr(self, projector_name, qwen3_5_vision.Qwen3_5MoeVisionProjector(config=self.config, rngs=self.rngs))
-      return encoder_name, projector_name
     elif self.vision_encoder_block == VisionEncoderBlockType.QWEN3_VL:
       from maxtext.models import qwen3_vl_vision  # pylint: disable=import-outside-toplevel
 
@@ -91,14 +87,25 @@ def _setup_vision_encoder_layers(self):
           self, encoder_name, qwen3_vl_vision.Qwen3VLVisionEncoder(config=self.config, mesh=self.mesh, rngs=self.rngs)
       )
       setattr(self, projector_name, qwen3_vl_vision.Qwen3VLVisionProjector(config=self.config, rngs=self.rngs))
-      return encoder_name, projector_name
     else:
       supported_blocks = [block.value for block in VisionEncoderBlockType if block != VisionEncoderBlockType.NONE]
       raise ValueError(
           f"Unsupported vision_encoder_block={self.vision_encoder_block.value!r} "
           f"for model_name={self.config.model_name!r}. Supported values are: {supported_blocks}."
       )
 
+    # 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":
+      setattr(
+          self,
+          projector_name,
+          MultimodalMLPProjector(config=self.config, mesh=self.mesh, rngs=self.rngs),
+      )
+
+    return encoder_name, projector_name
+
   def __call__(self, input_images, input_masks=None, video_grid_thw=None, deterministic=False):
     # vision encoder output, frozen params in many cases
     encoder = getattr(self, self.encoder_name)
@@ -127,6 +134,83 @@ def __call__(self, input_images, input_masks=None, video_grid_thw=None, determin
     return embeddings, deep_feats
 
 
+class MultimodalMLPProjector(nnx.Module):
+  """A multi-layer perceptron (MLP) projector for multimodal models."""
+
+  def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs):
+    self.config = config
+    self.mesh = mesh
+    self.rngs = rngs
+
+    # Input and output dimensions
+    in_features = config.hidden_size_for_vit
+    out_features = config.emb_dim
+
+    # Custom MLP projector hyperparameters
+    self.num_layers = getattr(config, "vision_connector_num_layers", 2)
+    self.hidden_size = getattr(config, "vision_connector_hidden_size", out_features)
+    if self.hidden_size == 0:
+      self.hidden_size = out_features
+    self.activation_name = getattr(config, "vision_connector_activation", "gelu")
+    self.use_bias = getattr(config, "vision_connector_use_bias", True)
+
+    vision_block_str = str(config.vision_encoder_block).lower()
+
+    # Qwen model family uses 2x2 spatial patch merging (need extra reshape in forward pass)
+    if "qwen" in vision_block_str:
+      spatial_merge_size = getattr(config, "spatial_merge_size_for_vit", 2)
+    # Gemma, LLaMA 4, and other model families use 1:1 token projection in visual embedding
+    else:
+      spatial_merge_size = 1
+
+    self.tokens_per_block = spatial_merge_size**2
+
+    # Activations
+    activations = {
+        "gelu": jax.nn.gelu,
+        "silu": jax.nn.silu,
+        "swish": jax.nn.silu,
+        "relu": jax.nn.relu,
+        "sigmoid": jax.nn.sigmoid,
+        "tanh": jax.nn.tanh,
+    }
+    self.activation = activations.get(self.activation_name.lower(), jax.nn.gelu)
+
+    current_in = in_features * self.tokens_per_block
+    for i in range(self.num_layers):
+      current_out = out_features if i == self.num_layers - 1 else self.hidden_size
+      layer = linears.DenseGeneral(
+          in_features_shape=current_in,
+          out_features_shape=current_out,
+          dtype=config.dtype_mm,
+          weight_dtype=config.weight_dtype,
+          matmul_precision=config.matmul_precision,
+          use_bias=self.use_bias,
+          kernel_init=lambda key, shape, dtype, *args, **kwargs: jax.nn.initializers.normal(stddev=0.02)(
+              key, shape, dtype
+          ),
+          kernel_axes=("embed", "mlp"),
+          rngs=rngs,
+      )
+
+      setattr(self, f"custom_linear_{i}", layer)
+      current_in = current_out
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    # for qwen3 models, concatenate tokens per block (e.g. 2x2 spatial patches) along feature dimension
+    if self.tokens_per_block > 1 and x.ndim == 3 and x.shape[1] % self.tokens_per_block == 0:
+      batch_size, seq_len, in_dim = x.shape
+      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):
+      linear_layer = getattr(self, f"custom_linear_{i}")
+      x = linear_layer(x)
+      if i < self.num_layers - 1:
+        x = self.activation(x)
+    return x
+
+
 class AudioEncoder(nnx.Module):
   """Audio encoder to encode audio features into soft tokens."""
 

src/maxtext/utils/model_creation_utils.py

@@ -1094,9 +1094,15 @@ def _build_value_target(v):
       # Free memory used by initial sharded_state before restore, to make room for the incoming checkpoint arrays.
       # Skip transient runtime variables (RngState, Cache, Intermediate, BatchStat) — they hold runtime state
       # that is not present in the checkpoint and must remain valid after the restore.
-      def _free_device_memory(node):
-        if isinstance(node, nnx.Variable) and not isinstance(
-            node, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat)
+      def _free_device_memory(path, node):
+        # Check if the path contains any name containing 'custom' which indicates a customized layer.
+        # If yes, skip freeing device memory for this node and its children so they can be
+        # initialized with random values.
+        is_custom = any("custom_linear" in str(getattr(p, "key", p)) for p in path)
+        if (
+            isinstance(node, nnx.Variable)
+            and not isinstance(node, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat))
+            and not is_custom
         ):
           inner = node.get_value() if hasattr(node, "get_value") else node[...]
           # AQT serve-mode `qrhs.frozen` wraps a QTensor (composite pytree) rather
@@ -1105,12 +1111,12 @@ def _free_device_memory(node):
           for leaf in jax.tree_util.tree_leaves(inner):
             if isinstance(leaf, jax.Array) and not leaf.is_deleted():
               leaf.delete()
-        elif isinstance(node, jax.Array) and not node.is_deleted():
+        elif isinstance(node, jax.Array) and not node.is_deleted() and not is_custom:
           node.delete()
 
         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))
 
       restored = ckptr.restore(
           epath.Path(config.load_parameters_path),