diff --git a/src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml b/src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml index fbbcc3116b..9aceb27d64 100644 --- a/src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml +++ b/src/maxtext/experimental/omni_poc/omni-gemma3-qwen3.yml @@ -1,12 +1,15 @@ # model config for omni-gemma3-qwen3 # Combines LLM backbone from Qwen 3 4B with Vision features from Gemma 3 4B -# NOTE: model_name is set to "gemma3-4b" so encoders.py loads Gemma 3's Vision Tower without modification. -# Meanwhile, decoders.py can build Qwen 3 LLM decoder directly from reading `decoder_block: "qwen3"` below. -model_name: "gemma3-4b" +# NOTE: vision_encoder_block is set to "gemma3" so encoders.py loads Gemma 3's Vision Tower without modification. +# Meanwhile, decoders.py builds Qwen 3 LLM decoder directly from reading `decoder_block: "qwen3"` below. +# base_config is currently a placeholder flag for future SFT runs. +base_config: "base.yml" +model_name: "omni-gemma3-qwen3" use_multimodal: true -# Multimodal config for vision model (from gemma3-4b.yml) +# Vision encoder config (from gemma3-4b.yml) +vision_encoder_block: "gemma3" image_size_for_vit: 896 num_channels_for_vit: 3 patch_size_for_vit: 14 @@ -15,7 +18,14 @@ hidden_size_for_vit: 1152 intermediate_size_for_vit: 4304 num_hidden_layers_for_vit: 27 num_attention_heads_for_vit: 16 -image_placeholder: +image_placeholder: "" + +# Vision projector config +vision_projector_type: "customized_vision_projector" +vision_connector_hidden_size: 2560 +vision_connector_num_layers: 2 +vision_connector_activation: "gelu" +vision_connector_use_bias: true # LLM backbone config (from qwen3-4b.yml) base_emb_dim: 2560 @@ -34,6 +44,7 @@ logits_via_embedding: true normalize_embedding_logits: false enable_dropout: false tokenizer_type: "huggingface" +tokenizer_path: "Qwen/Qwen2.5-7B-Instruct" # Ensure no pretrained weights or checkpoints are loaded initially load_parameters_path: "" diff --git a/src/maxtext/experimental/omni_poc/tests/processor_omni_gemma3_qwen3_test.py b/src/maxtext/experimental/omni_poc/tests/processor_omni_gemma3_qwen3_test.py new file mode 100644 index 0000000000..e3a596e17c --- /dev/null +++ b/src/maxtext/experimental/omni_poc/tests/processor_omni_gemma3_qwen3_test.py @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for Omni multimodal processor.""" + +from types import SimpleNamespace +import unittest +import numpy as np + +from maxtext.experimental.omni_poc.utils import processor_omni_gemma3_qwen3 as omni_processor + + +class TestProcessorOmniGemma3Qwen3(unittest.TestCase): + """Concise test suite for processor_omni_gemma3_qwen3.""" + + def test_constants_and_offsets(self): + self.assertEqual(omni_processor.IMAGE_PAD_ID, 151655) + self.assertEqual(omni_processor.get_image_offsets_omni(None), 255) + self.assertEqual( + omni_processor.get_image_offsets_omni(SimpleNamespace(pixel_values=np.zeros((3, 1, 1, 3)))), + 255 * 3, + ) + + def test_add_extra_tokens_placeholder_expansion(self): + pad_id = omni_processor.IMAGE_PAD_ID + # Expand single placeholder with custom count + tokens = [10, pad_id, 20] + out = omni_processor.add_extra_tokens_for_omni(tokens, num_tokens_per_image=3) + np.testing.assert_array_equal(out, [10, pad_id, pad_id, pad_id, 20]) + + # 2D array and dtype preservation + tokens_i64 = np.array([[pad_id], [99]], dtype=np.int64) + out_i64 = omni_processor.add_extra_tokens_for_omni(tokens_i64, num_tokens_per_image=2) + self.assertEqual(out_i64.dtype, np.int64) + np.testing.assert_array_equal(out_i64, [pad_id, pad_id, 99]) + + def test_add_extra_tokens_prepend_and_text_only(self): + # Prepend vision block when pixel values exist + proc_out = SimpleNamespace(pixel_values=np.zeros((1, 1, 1, 3))) + out = omni_processor.add_extra_tokens_for_omni([1, 2], processor_output=proc_out, num_tokens_per_image=2) + expected = [ + omni_processor.VISION_START_ID, + omni_processor.IMAGE_PAD_ID, + omni_processor.IMAGE_PAD_ID, + omni_processor.VISION_END_ID, + 1, + 2, + ] + np.testing.assert_array_equal(out, expected) + + # Text-only returns original + np.testing.assert_array_equal(omni_processor.add_extra_tokens_for_omni([1, 2]), [1, 2]) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py b/src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py index dc82b3e118..cadd0006b4 100644 --- a/src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py +++ b/src/maxtext/experimental/omni_poc/tests/stitch_checkpoint_test.py @@ -52,12 +52,13 @@ from maxtext.configs import pyconfig from maxtext.experimental.omni_poc.utils import stitch_checkpoint +from maxtext.layers import encoders from maxtext.layers.embeddings import Embed from maxtext.models import gemma3 from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils -from maxtext.utils.globals import MAXTEXT_REPO_ROOT +from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR, MAXTEXT_PKG_DIR warnings.filterwarnings("ignore", category=DeprecationWarning) warnings.filterwarnings("ignore", category=FutureWarning) @@ -120,9 +121,7 @@ def tearDown(self): def _create_config(self): """Creates a MaxText config from the Omni config file and defaults.""" omni_config_path = os.path.join( - MAXTEXT_REPO_ROOT, - "src", - "maxtext", + MAXTEXT_PKG_DIR, "experimental", "omni_poc", "omni-gemma3-qwen3.yml", @@ -136,10 +135,12 @@ def _create_config(self): "stitched_output_path", "vision_model_name", "llm_model_name", + "base_config", + "model_name", } yaml_overrides = {k: v for k, v in custom_cfg.items() if k not in omni_keys} - base_config_path = os.path.join(MAXTEXT_REPO_ROOT, "src", "maxtext", "configs", "base.yml") + base_config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml") config = pyconfig.initialize( ["", base_config_path], override_model_config=True, @@ -151,6 +152,7 @@ def _create_config(self): **yaml_overrides, ) # Use object.__setattr__ to bypass the read-only check on _HyperParameters + object.__setattr__(config, "model_name", custom_cfg.get("model_name", "omni-gemma3-qwen3")) object.__setattr__(config, "vision_model_name", "gemma3-4b") object.__setattr__(config, "llm_model_name", "qwen3-4b") object.__setattr__(config, "ici_context_autoregressive_parallelism", 1) @@ -206,7 +208,11 @@ def to_concrete_zeros(leaf): stitched_inner = stitch_checkpoint._restore_subtrees_from_path(self.output_ckpt_dir, concrete_template, ckptr) # Load weights from the ORIGINAL Vision checkpoint - vision_abstract = concrete_template["vision_encoder"] + vision_abstract = { + k: v + for k, v in concrete_template["vision_encoder"].items() + if not ("projector" in k.lower() or "embedder" in k.lower()) + } vision_restored = stitch_checkpoint._restore_subtrees_from_path( self.vision_ckpt_dir, {"vision_encoder": vision_abstract}, ckptr ) @@ -323,18 +329,14 @@ def test_2_stage_outputs_original_vs_stitched(self): "stitched llm token embedder first 10 values:", stitched_inner["token_embedder"]["embedding"].reshape(-1)[:10], ) - print( - "orig vision projector weights first 10 values:", - vision_restored["vision_encoder"]["VisionEmbedder_0"]["mm_input_projection"]["w"].reshape(-1)[:10], - ) print( "stitched vision projector weights first 10 values:", - stitched_inner["vision_encoder"]["VisionEmbedder_0"]["mm_input_projection"]["w"].reshape(-1)[:10], + stitched_inner["vision_encoder"]["VisionEmbedder_0"]["custom_linear_0"]["kernel"].reshape(-1)[:10], ) # Instantiate and update original vision & token embedder modules orig_vision_tower = gemma3.Gemma3VisionEncoderLayer(config, self.test_mesh, rngs=rngs) - orig_projector = gemma3.VisionEmbedder(config, self.test_mesh, rngs=rngs) + orig_projector = encoders.MultimodalMLPProjector(config, self.test_mesh, rngs=nnx.Rngs(1)) # Manually define qwen3's token embedder to match MaxText's structures orig_token_embedder = Embed( num_embeddings=config.vocab_size, @@ -348,10 +350,6 @@ def test_2_stage_outputs_original_vs_stitched(self): orig_vision_tower, vision_restored["vision_encoder"]["Gemma3VisionEncoderLayer_0"], ) - nnx.update( - orig_projector, - vision_restored["vision_encoder"]["VisionEmbedder_0"], - ) nnx.update( orig_token_embedder, {"embedding": llm_restored["token_embedder"]["embedding"]}, @@ -359,7 +357,7 @@ def test_2_stage_outputs_original_vs_stitched(self): # Instantiate and update stitched vision & token embedder modules stitched_vision_tower = gemma3.Gemma3VisionEncoderLayer(config, self.test_mesh, rngs=rngs) - stitched_projector = gemma3.VisionEmbedder(config, self.test_mesh, rngs=rngs) + stitched_projector = encoders.MultimodalMLPProjector(config, self.test_mesh, rngs=rngs) stitched_token_embedder = Embed( num_embeddings=config.vocab_size, num_features=config.emb_dim, diff --git a/src/maxtext/experimental/omni_poc/utils/decode_omni.py b/src/maxtext/experimental/omni_poc/utils/decode_omni.py new file mode 100644 index 0000000000..ace0175a80 --- /dev/null +++ b/src/maxtext/experimental/omni_poc/utils/decode_omni.py @@ -0,0 +1,302 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Retrieve random examples from ChartQA dataset and evaluate the stitched model +(pre-SFT and post-SFT) checkpoints to verify if: + 1. the data flows end-to-end + 2. the response is reasonable given the random vision projector +ChartQA sample Q&A, ground truth and model responses are logged. + +Example usage: + +python src/maxtext/experimental/omni_poc/utils/decode_omni.py \ + --checkpoint_path="gs://YOUR_BUCKET_NAME/omni_stitched_gemma3-4b_qwen3-4b/0/items" \ + --num_samples=3 \ + --max_new_tokens=128 + +""" + +import functools +import os +import random +import sys +from absl import app, flags +import datasets +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +import numpy as np +import omegaconf +from transformers import AutoTokenizer + +from maxtext.common import checkpointing +from maxtext.configs import pyconfig +from maxtext.multimodal import processor as mm_processor +from maxtext.utils import max_logging +from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils +from maxtext.utils import model_creation_utils +from maxtext.utils.globals import MAXTEXT_PKG_DIR + +FLAGS = flags.FLAGS + + +def _define_flag(fn, name, default, help_str): + if name not in FLAGS: + fn(name, default, help_str) + + +_define_flag(flags.DEFINE_string, "config_path", "", "Path to the config YAML file.") +_define_flag(flags.DEFINE_string, "checkpoint_path", "", "Path to checkpoint parameters directory.") +_define_flag(flags.DEFINE_integer, "num_samples", 5, "Number of random ChartQA validation samples to evaluate.") +_define_flag(flags.DEFINE_string, "description", "Omni Model", "Description for evaluation logs.") +_define_flag(flags.DEFINE_integer, "max_new_tokens", 128, "Maximum number of new tokens to generate.") + + +def initialize_model_and_weights(config, checkpoint_path, mesh): + """Instantiates the Omni model and loads checkpoint. + + Args: + config: The MaxText Omni configuration. + checkpoint_path: Path to the checkpoint parameters directory. + mesh: The JAX mesh for sharding. + + Returns: + model: The stitched model. + params: The restored model parameters. + """ + with jax.set_mesh(mesh): + model = model_creation_utils.from_config(config, mesh=mesh) + + abstract_vars = maxtext_utils.get_abstract_param(model, config) + target_params_abstract = max_utils.unbox_logicallypartioned(abstract_vars["params"]) + + max_logging.log(f"Checkpoint parameters path: {checkpoint_path}") + restored = checkpointing.load_params_from_path( + checkpoint_path, + {"params": target_params_abstract}, + config.checkpoint_storage_concurrent_gb, + use_ocdbt=config.checkpoint_storage_use_ocdbt, + use_zarr3=config.checkpoint_storage_use_zarr3, + ) + return model, restored.get("params", restored) + + +def load_omni_config(yaml_path, checkpoint_path): + """Loads custom omni config YAML and converts overrides onto base.yml.""" + custom_cfg = omegaconf.OmegaConf.to_container(omegaconf.OmegaConf.load(yaml_path), resolve=True) + base_yml = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml") + argv = [ + sys.argv[0], + base_yml, + "override_model_config=True", + "skip_jax_distributed_system=True", + "ici_fsdp_parallelism=1", + "ici_tensor_parallelism=-1", + f"load_parameters_path={checkpoint_path}", + ] + + omni_skip_keys = { + "vision_load_path", + "llm_load_path", + "stitched_output_path", + "vision_model_name", + "llm_model_name", + "base_config", + "model_name", + } + for k, v in custom_cfg.items(): + if k in omni_skip_keys: + continue + if isinstance(v, str): + argv.append(f"{k}='{v}'") + else: + argv.append(f"{k}={v}") + + # Initialize config + config = pyconfig.initialize( + argv, + override_model_config=True, + skip_jax_distributed_system=True, + log_config=False, + ) + # Explicitly set model_name on the frozen config for multimodal routing in processor.py. + # This prevents pyconfig search errors or inheriting conflicting base model hyperparameters. + # TODO: works for experimental piloting. For production, relocate model config to the + # main folder, remove this line so pyconfig can initialize model_name directly. + object.__setattr__(config, "model_name", "omni-gemma3-qwen3") + return config + + +@functools.partial(jax.jit, static_argnums=(0,)) +def _forward_step(model, params, tokens, positions, segment_ids, images): + """Executes a single forward pass of the model and returns output logits.""" + return model.apply( + {"params": params}, + decoder_input_tokens=tokens, + decoder_positions=positions, + decoder_segment_ids=segment_ids, + encoder_images=images, + enable_dropout=False, + model_mode="prefill", + ) + + +def decode_omni_sample(model, params, config, mesh, tokenizer, prompt_str, pil_image, max_new_tokens=128): + """Runs prefill and autoregressive decoding for a single multimodal sample.""" + # Preprocess image + image_np = np.array(pil_image.convert("RGB"), dtype=np.uint8) + processed_image = mm_processor.preprocess_image_for_training(image_np, config) + image_pixels = ( + processed_image.pixel_values + if hasattr(processed_image, "pixel_values") and processed_image.pixel_values is not None + else processed_image + ) + image_pixels = np.asarray(image_pixels) + if image_pixels.ndim == 4: + image_pixels = np.expand_dims(image_pixels, axis=0) + mock_image = jnp.array(image_pixels, dtype=jnp.bfloat16 if config.dtype == "bfloat16" else jnp.float32) + + # Format prompt and expand image placeholder tokens + formatted_prompt = mm_processor.reformat_prompt( + prompt=prompt_str, + image_placeholder=config.image_placeholder, + model_name=config, + num_images=1, + ) + initial_tokens = tokenizer.encode(formatted_prompt, add_special_tokens=False) + combined_tokens = mm_processor.prepare_text_for_image_fusion( + tokens=initial_tokens, + config=config, + processor_output=processed_image, + ).tolist() + + # Pad token sequence to fixed config.max_target_length + true_length = len(combined_tokens) + seq_len = config.max_target_length + if true_length > seq_len: + raise ValueError( + f"The combined length of expanded prompt and vision tokens ({true_length}) " + f"exceeds config.max_target_length ({seq_len}). " + "Please increase max_target_length in your model config." + ) + pad_token = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + padded_tokens = combined_tokens + [pad_token] * (seq_len - true_length) + + tokens = np.array([padded_tokens[:seq_len]], dtype=np.int32) + positions = np.tile(np.arange(seq_len, dtype=np.int32), (1, 1)) + segment_ids = np.zeros((1, seq_len), dtype=np.int32) + segment_ids[:, :true_length] = 1 + + # Autoregressive decoding loop + generated_tokens = [] + curr_len = true_length + + with jax.set_mesh(mesh): + # Initial prefill pass + logits = _forward_step(model, params, tokens, positions, segment_ids, mock_image) + + for _ in range(max_new_tokens): + next_token_id = int(jnp.argmax(logits[0, curr_len - 1, :])) + if next_token_id == tokenizer.eos_token_id: + break + generated_tokens.append(next_token_id) + if curr_len >= seq_len: + break + tokens[0, curr_len] = next_token_id + curr_len += 1 + + # Update segment IDs and positions and do the next autoregressive step + segment_ids[0, :curr_len] = 1 + logits = _forward_step(model, params, tokens, positions, segment_ids, mock_image) + + return tokenizer.decode(generated_tokens, skip_special_tokens=True).strip() + + +def run_evaluation(checkpoint_path, config, num_samples=5, description="Omni Model", max_new_tokens=128): + """Runs evaluation on ChartQA random validation samples. + Args: + checkpoint_path: Path to the checkpoint parameters directory. + config: The MaxText Omni configuration. + num_samples: Number of random ChartQA validation samples to evaluate. + description: Description for evaluation logs. + max_new_tokens: Maximum number of new tokens to generate. + """ + max_logging.log("=" * 60) + max_logging.log(f"Running Evaluation for {description}...") + max_logging.log(f"Loading checkpoint from: {checkpoint_path}") + max_logging.log("=" * 60) + + # Load the dataset + try: + ds = datasets.load_dataset("HuggingFaceM4/ChartQA", split="val") + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.log(f"Error loading ChartQA evaluation dataset: {e}") + return + + # Initialize the model + if jax.local_devices()[0].platform == "cpu": + mesh = Mesh(np.array([jax.devices("cpu")[0]]), axis_names=("data",)) + else: + mesh = maxtext_utils.get_mesh_from_config(config) + + model, restored_params = initialize_model_and_weights(config, checkpoint_path, mesh) + + tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path) + + # Run decode on samples from the dataset + random.seed(42) + total_samples = len(ds) + sample_indices = random.sample(range(total_samples), min(num_samples, total_samples)) + + for idx, i in enumerate(sample_indices): + sample = ds[i] + model_response = decode_omni_sample( + model=model, + params=restored_params, + config=config, + mesh=mesh, + tokenizer=tokenizer, + prompt_str=f" {sample['query']}", + pil_image=sample["image"], + max_new_tokens=max_new_tokens, + ) + + max_logging.log( + f""" +Sample {idx+1} (Index {i}): + Question: {sample['query']} + Ground Truth: {sample['label']} + Model Response: {model_response} +""" + ) + + +def main(argv): + config_path = FLAGS.config_path or os.path.join(MAXTEXT_PKG_DIR, "experimental", "omni_poc", "omni-gemma3-qwen3.yml") + assert FLAGS.checkpoint_path, "Must specify --checkpoint_path" + config = load_omni_config(config_path, FLAGS.checkpoint_path) + + run_evaluation( + checkpoint_path=FLAGS.checkpoint_path, + config=config, + num_samples=FLAGS.num_samples, + description=FLAGS.description, + max_new_tokens=FLAGS.max_new_tokens, + ) + + +if __name__ == "__main__": + app.run(main) diff --git a/src/maxtext/experimental/omni_poc/utils/processor_omni_gemma3_qwen3.py b/src/maxtext/experimental/omni_poc/utils/processor_omni_gemma3_qwen3.py new file mode 100644 index 0000000000..ab886d951f --- /dev/null +++ b/src/maxtext/experimental/omni_poc/utils/processor_omni_gemma3_qwen3.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multimodal processor for Omni (Gemma 3 Vision Encoder + Qwen 3 LLM Decoder). + +Handles Qwen 3 text-only tokenizer limitations by explicitly mapping and expanding +visual token sequences (<|vision_start|>, <|image_pad|>*256, <|vision_end|>), +where 256 is the default number of placeholder tokens per image for Gemma 3. +""" + +import numpy as np +from maxtext.multimodal.processor_gemma3 import GEMMA_NUM_PLACEHOLDER_TOKENS_PER_IMAGE +from maxtext.multimodal.processor_qwen3_omni import QWEN_SPECIAL_TOKEN_CONFIGS + +# Load visual token IDs from Qwen 3 config +_QWEN_TOKENS = QWEN_SPECIAL_TOKEN_CONFIGS["qwen3-omni-30b-a3b"] +VISION_START_ID = _QWEN_TOKENS["vision_start"] # 151652 (<|vision_start|>) +VISION_END_ID = _QWEN_TOKENS["vision_end"] # 151653 (<|vision_end|>) +IMAGE_PAD_ID = _QWEN_TOKENS["image_pad"] # 151655 (<|image_pad|>) +QWEN_IMAGE_TAG = "<|vision_start|><|image_pad|><|vision_end|>" + +# Load vision token count per image from Gemma 3's config +DEFAULT_NUM_TOKENS_PER_IMAGE = GEMMA_NUM_PLACEHOLDER_TOKENS_PER_IMAGE # 256 + + +def get_image_offsets_omni(processor_output=None): + """Calculate the increase in total token count after inserting visual token sequences.""" + has_images = processor_output is not None and processor_output.pixel_values is not None + num_images = processor_output.pixel_values.shape[0] if has_images else 1 + # +256 for <|image_pad|>, -1 for original placeholder (<|vision_start|> and <|vision_end|> already in prompt) + return (DEFAULT_NUM_TOKENS_PER_IMAGE - 1) * num_images + + +def add_extra_tokens_for_omni( + tokens, + config=None, + processor_output=None, + num_tokens_per_image=DEFAULT_NUM_TOKENS_PER_IMAGE, +): + """Expands <|image_pad|> placeholders or prepends vision tokens if missing. + + - If <|image_pad|> is present in `tokens`, expands each placeholder into `num_tokens_per_image` copies. + - If <|image_pad|> is absent but `processor_output` contains image data, + automatically prepends the vision sequence: [<|vision_start|>, <|image_pad|>*N, <|vision_end|>]. + - Otherwise, returns the original tokens unchanged. + """ + dtype = tokens.dtype if isinstance(tokens, np.ndarray) else np.int32 + token_list = np.asarray(tokens).flatten().tolist() + + # Case 1: Prompt contains <|image_pad|> placeholders, expand them + if IMAGE_PAD_ID in token_list: + expanded_tokens = [] + for token in token_list: + if token == IMAGE_PAD_ID: + expanded_tokens.extend([IMAGE_PAD_ID] * num_tokens_per_image) + else: + expanded_tokens.append(token) + return np.array(expanded_tokens, dtype=dtype) + + # Case 2: No placeholder in text, but image data is present, prepend vision block + if processor_output is not None and processor_output.pixel_values is not None: + num_images = processor_output.pixel_values.shape[0] + vision_block = [] + for _ in range(num_images): + vision_block.extend([VISION_START_ID] + [IMAGE_PAD_ID] * num_tokens_per_image + [VISION_END_ID]) + return np.array(vision_block + token_list, dtype=dtype) + + # Case 3: Text-only sequence, return original tokens + return np.array(token_list, dtype=dtype) diff --git a/src/maxtext/experimental/omni_poc/utils/stitch_checkpoint.py b/src/maxtext/experimental/omni_poc/utils/stitch_checkpoint.py index a8e2518c0f..f2808cecc5 100644 --- a/src/maxtext/experimental/omni_poc/utils/stitch_checkpoint.py +++ b/src/maxtext/experimental/omni_poc/utils/stitch_checkpoint.py @@ -21,7 +21,8 @@ Example usage: -python -m maxtext.experimental.omni_poc.utils.stitch_checkpoint \ +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://YOUR_BUCKET_NAME/checkpoints/gemma3-4b_converted/0/items \ --llm_load_path=gs://YOUR_BUCKET_NAME/checkpoints/qwen3-4b_converted/0/items \ --stitched_output_path=gs://YOUR_BUCKET_NAME/checkpoints/omni-gemma3-qwen3-4b/0/items @@ -30,13 +31,18 @@ import os from typing import Any, Dict -from absl import app +from absl import app, flags from etils import epath from flax import nnx import jax import omegaconf from orbax import checkpoint as ocp +FLAGS = flags.FLAGS +flags.DEFINE_string("vision_load_path", "", "Path to the vision model checkpoint.") +flags.DEFINE_string("llm_load_path", "", "Path to the LLM checkpoint.") +flags.DEFINE_string("stitched_output_path", "", "Path to save the stitched checkpoint.") + from maxtext.common import checkpointing from maxtext.configs import pyconfig as pyconfig_mod from maxtext.trainers.pre_train.train import initialize @@ -139,13 +145,16 @@ def stitch_and_save_checkpoints( max_logging.log("=" * 60) max_logging.log("Starting Omni Multi-Directory Checkpoint Stitching...") - vision_model_name = getattr(config, "model_name", None) + vision_model_name = getattr(config, "vision_encoder_block", None) llm_model_name = getattr(config, "decoder_block", None) - assert vision_model_name, "model_name must be configured for vision component." + vision_projector_type = getattr(config, "vision_projector_type", None) + assert vision_model_name, "vision_encoder_block must be configured for vision component." assert llm_model_name, "decoder_block must be configured for LLM component." + assert vision_projector_type, "vision_projector_type must be configured for vision component." max_logging.log(f" Vision (Model {vision_model_name}) Path: {vision_checkpoint_path}") max_logging.log(f" LLM (Model {llm_model_name}) Path: {llm_checkpoint_path}") + max_logging.log(f" Projector Type: {vision_projector_type}") max_logging.log(f" Output Stitched Path: {output_checkpoint_path}") max_logging.log("=" * 60) @@ -206,11 +215,7 @@ def stitch_and_save_checkpoints( # 4. Assemble: Vision (Model A) + LLM (Model B) + Random Init Projector stitched_inner = {k: _assemble(k, v, stitched_subtrees) for k, v in inner_params.items()} - final_params = ( - {"params": stitched_inner} - if "params" in params_dict and isinstance(params_dict["params"], dict) - else stitched_inner - ) + final_params = {"params": stitched_inner} # 5. Save unified parameter tree to output_checkpoint_path max_logging.log(f"Saving stitched checkpoint to: {output_checkpoint_path}") @@ -241,13 +246,21 @@ def _load_custom_yaml_overrides(yaml_path: str, omni_keys: set[str]): def main(argv): - # Extract omni stitching arguments directly from argv before passing to pyconfig.initialize - omni_keys = {"vision_load_path", "llm_load_path", "stitched_output_path", "vision_model_name", "llm_model_name"} + omni_keys = { + "vision_load_path", + "llm_load_path", + "stitched_output_path", + "vision_model_name", + "llm_model_name", + "base_config", + "model_name", + } omni_kwargs = {} cleaned_argv = [] for arg in argv: - if "=" in arg and arg.split("=", 1)[0] in omni_keys: - k, v = arg.split("=", 1) + cleaned_arg = arg.lstrip("-") + if "=" in cleaned_arg and cleaned_arg.split("=", 1)[0] in omni_keys: + k, v = cleaned_arg.split("=", 1) omni_kwargs[k] = v else: cleaned_argv.append(arg) @@ -275,12 +288,16 @@ def main(argv): cleaned_argv[1] = os.path.join(pyconfig_mod.MAXTEXT_CONFIGS_DIR, "base.yml") + if not any(arg.startswith("skip_jax_distributed_system=") for arg in cleaned_argv): + cleaned_argv.append("skip_jax_distributed_system=True") + # Initialize MaxText config using standard train.initialize config, _ = initialize(cleaned_argv) - # Extract paths from command-line arguments - vision_path = omni_kwargs.get("vision_load_path") - llm_path = omni_kwargs.get("llm_load_path") - output_path = omni_kwargs.get("stitched_output_path") + object.__setattr__(config, "model_name", "omni-gemma3-qwen3") + # Extract paths from command-line arguments or FLAGS + vision_path = FLAGS.vision_load_path or omni_kwargs.get("vision_load_path") + llm_path = FLAGS.llm_load_path or omni_kwargs.get("llm_load_path") + output_path = FLAGS.stitched_output_path or omni_kwargs.get("stitched_output_path") assert ( vision_path and llm_path and output_path ), "Must specify vision_load_path, llm_load_path, and stitched_output_path" diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 42753eb752..b2f6293774 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -722,6 +722,7 @@ def _apply_embedding( "qwen3-vl-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "omni-gemma3-qwen3", ]: y = mm_utils.merge_mm_embeddings( text_embeddings=y, diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index ee02407502..e80af834c8 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1361,6 +1361,7 @@ def _apply_embedding( "qwen3-vl-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", + "omni-gemma3-qwen3", }: y = mm_utils.merge_mm_embeddings( text_embeddings=y, diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py index 025d8ab40b..8ee0e7143e 100644 --- a/src/maxtext/multimodal/processor.py +++ b/src/maxtext/multimodal/processor.py @@ -37,6 +37,8 @@ "qwen3-vl-30b-a3b": ("qwen3_vl", "qwen3_moe"), "qwen3.5-35b-a3b": ("qwen3_5", "qwen3_5"), "qwen3.5-397b-a17b": ("qwen3_5", "qwen3_5"), + # Stitched model + "omni-gemma3-qwen3": ("gemma3", "qwen3"), } @@ -138,8 +140,13 @@ def preprocess_image_for_training(image, config): def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | None): """Get the increase in total token count after inserting image token placeholders""" vision_block = _get_vision_block(config) + decoder_block = _get_decoder_block(config) - if vision_block in ["gemma3"]: + if vision_block == "gemma3" and decoder_block == "qwen3": + from maxtext.experimental.omni_poc.utils import processor_omni_gemma3_qwen3 # pylint: disable=import-outside-toplevel + + return processor_omni_gemma3_qwen3.get_image_offsets_omni(processor_output) + elif vision_block in ["gemma3"]: from maxtext.multimodal.processor_gemma3 import get_image_offsets_gemma3 # pylint: disable=import-outside-toplevel return get_image_offsets_gemma3(processor_output) @@ -220,7 +227,13 @@ def reformat_response(response, model_name): def prepare_text_for_image_fusion(tokens, config, processor_output=None): """Prepare text by adding extra tokens for image fusion based on the model.""" vision_block = _get_vision_block(config) - if vision_block in ["gemma3"]: + decoder_block = _get_decoder_block(config) + + if vision_block == "gemma3" and decoder_block == "qwen3": + from maxtext.experimental.omni_poc.utils import processor_omni_gemma3_qwen3 # pylint: disable=import-outside-toplevel + + return processor_omni_gemma3_qwen3.add_extra_tokens_for_omni(tokens, config=config, processor_output=processor_output) + elif vision_block in ["gemma3"]: from maxtext.multimodal.processor_gemma3 import add_extra_tokens_for_images_gemma3 # pylint: disable=import-outside-toplevel return add_extra_tokens_for_images_gemma3( @@ -296,7 +309,11 @@ def get_bidirectional_mask_vision(config, decoder_input_tokens, is_video: bool = decoder_block = _get_decoder_block(config) - if decoder_block in ["gemma3"]: + if vision_block == "gemma3" and decoder_block == "qwen3": + from maxtext.experimental.omni_poc.utils import processor_omni_gemma3_qwen3 # pylint: disable=import-outside-toplevel + + bidirectional_mask_vision = decoder_input_tokens == processor_omni_gemma3_qwen3.IMAGE_PAD_ID + elif decoder_block in ["gemma3"]: from maxtext.multimodal.processor_gemma3 import GEMMA_TOKEN_PLACEHOLDER # pylint: disable=import-outside-toplevel bidirectional_mask_vision = decoder_input_tokens == GEMMA_TOKEN_PLACEHOLDER diff --git a/tests/unit/multimodal_utils_test.py b/tests/unit/multimodal_utils_test.py index efcd58db89..9a94d4600d 100644 --- a/tests/unit/multimodal_utils_test.py +++ b/tests/unit/multimodal_utils_test.py @@ -500,6 +500,29 @@ def test_preprocess_image_for_training(self): with self.assertRaises(ValueError): mm_processor.preprocess_image_for_training([dummy_image], config_text_only) + def test_omni_gemma3_qwen3_processor_routing(self): + # pylint: disable=protected-access,import-outside-toplevel + self.assertEqual(mm_processor._get_vision_block("omni-gemma3-qwen3"), "gemma3") + self.assertEqual(mm_processor._get_decoder_block("omni-gemma3-qwen3"), "qwen3") + + omni_config = types.SimpleNamespace( + vision_encoder_block=VisionEncoderBlockType.GEMMA3, + decoder_block=DecoderBlockType.QWEN3, + model_name="omni-gemma3-qwen3", + ) + self.assertEqual(mm_processor.get_image_offsets(omni_config, None), 255) + + from maxtext.experimental.omni_poc.utils import processor_omni_gemma3_qwen3 + + tokens = [1, processor_omni_gemma3_qwen3.IMAGE_PAD_ID, 2] + fused = mm_processor.prepare_text_for_image_fusion(tokens, config=omni_config) + self.assertEqual(len(fused), 1 + 256 + 1) + + mask = mm_processor.get_bidirectional_mask_vision( + omni_config, np.array([[10, processor_omni_gemma3_qwen3.IMAGE_PAD_ID]]) + ) + np.testing.assert_array_equal(mask, [[False, True]]) + if __name__ == "__main__": unittest.main()