From c908c8e8193f4dfab6da94179f078a25d09c12c3 Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Thu, 23 Jul 2026 23:06:47 +0000 Subject: [PATCH 1/2] Add custom MLP vision projector and tests --- src/maxtext/common/checkpointing.py | 17 ++ src/maxtext/configs/types.py | 10 ++ .../tests/custom_vision_projector_test.py | 155 ++++++++++++++++++ src/maxtext/layers/encoders.py | 96 ++++++++++- src/maxtext/utils/model_creation_utils.py | 16 +- 5 files changed, 283 insertions(+), 11 deletions(-) create mode 100644 src/maxtext/experimental/omni_poc/tests/custom_vision_projector_test.py diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index cf51083597..8846f77b30 100644 --- a/src/maxtext/common/checkpointing.py +++ b/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) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index dc08063d7a..f358125b29 100644 --- a/src/maxtext/configs/types.py +++ b/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): diff --git a/src/maxtext/experimental/omni_poc/tests/custom_vision_projector_test.py b/src/maxtext/experimental/omni_poc/tests/custom_vision_projector_test.py new file mode 100644 index 0000000000..342ffe7dd9 --- /dev/null +++ b/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, + ) diff --git a/src/maxtext/layers/encoders.py b/src/maxtext/layers/encoders.py index 0eed92ae3e..162f50c540 100644 --- a/src/maxtext/layers/encoders.py +++ b/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,7 +44,6 @@ 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 @@ -51,7 +51,6 @@ def _setup_vision_encoder_layers(self): 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 @@ -59,7 +58,6 @@ def _setup_vision_encoder_layers(self): 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,7 +87,6 @@ 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( @@ -99,6 +94,18 @@ def _setup_vision_encoder_layers(self): 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.""" diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index 7385408379..2326bff6b5 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/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), From 6cc1f5f82da2c8464ce53cd0f65df7dcfde00d3d Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Thu, 23 Jul 2026 23:22:38 +0000 Subject: [PATCH 2/2] add unit test --- tests/unit/model_creation_utils_test.py | 26 +++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/unit/model_creation_utils_test.py b/tests/unit/model_creation_utils_test.py index b0ec54f694..6123e3eda5 100644 --- a/tests/unit/model_creation_utils_test.py +++ b/tests/unit/model_creation_utils_test.py @@ -711,6 +711,32 @@ def test_load_linen_checkpoint(self, mock_ocp): model = model_creation_utils.from_pretrained(cfg, self.mesh) self.assertIsInstance(model, models.Transformer) + @patch("maxtext.utils.model_creation_utils.ocp") + def test_load_checkpoint_with_custom_vision_projector(self, mock_ocp): + """NNX checkpoint loading with customized_mlp vision projector runs _free_device_memory cleanly.""" + mock_ckptr = MagicMock() + mock_ckptr.metadata.return_value = self._make_nnx_metadata_mock() + mock_ckptr.restore.side_effect = lambda path, item=None, **kw: item + mock_ocp.Checkpointer.return_value = mock_ckptr + mock_ocp.PyTreeCheckpointHandler.return_value = MagicMock() + mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} + mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs + + cfg = _make_config( + enable_checkpointing=True, + load_parameters_path="gs://fake/custom_vision_ckpt", + use_multimodal=True, + model_name="qwen3-vl-2b", + override_model_config=True, + scan_layers=False, + vision_projector_type="customized_mlp", + vision_connector_num_layers=2, + vision_connector_activation="gelu", + vision_connector_use_bias=True, + ) + model = model_creation_utils.from_pretrained(cfg, self.mesh) + self.assertIsInstance(model, models.Transformer) + @patch("maxtext.utils.model_creation_utils.ocp") def test_checkpoint_load_error_propagates(self, mock_ocp): """Non-mismatch exceptions during checkpoint loading should propagate directly."""