Skip to content

Decouple vision & decoder block routing in multimodal processor #4746

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
subawocit wants to merge 1 commit into
base: main
Choose a base branch
from
+252 −88
Open

Decouple vision & decoder block routing in multimodal processor #4746

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
214 changes: 126 additions & 88 deletions src/maxtext/multimodal/processor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,72 @@

"""Multimodal data preprocessor router."""

import functools
import os
from maxtext.common.common_types import DecoderBlockType, VisionEncoderBlockType
from maxtext.multimodal import utils as mm_utils
from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR
import omegaconf


@functools.lru_cache(maxsize=None)
def _get_block_name_from_model_yml(model_name: str, block_name: str) -> str | None:

@hengtaoguo hengtaoguo Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The block_name could only be chosen from ['vision_encoder_block', 'decoder_block']? Could we add a docstring to explain this?

@subawocit subawocit Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good idea - I just updated the description

"""Loads a model YAML from MAXTEXT_CONFIGS_DIR and returns the specified block name.

Args:
model_name: Name of the model configuration file (e.g., 'gemma3-4b').
block_name: The architectural attribute to retrieve, can only be
'vision_encoder_block' or 'decoder_block'.

Returns:
The string value of the block configuration if defined, otherwise None.
"""
model_yml_path = os.path.join(MAXTEXT_CONFIGS_DIR, "models", f"{model_name}.yml")
if os.path.exists(model_yml_path):
try:
loaded = omegaconf.OmegaConf.load(model_yml_path)
if block_name in loaded and loaded[block_name] is not None:
return str(loaded[block_name])
except Exception: # pylint: disable=broad-except
pass
return None


def _get_vision_block(config_or_name):
"""Extract vision encoder architecture from config."""
# If input is config, extract vision_encoder name
if hasattr(config_or_name, "vision_encoder_block"):
block = config_or_name.vision_encoder_block
if block != VisionEncoderBlockType.NONE:
return block.value.lower()
# If input is model_name (backward compatibility), find its corresponding config
elif isinstance(config_or_name, str):
val = _get_block_name_from_model_yml(config_or_name, "vision_encoder_block")
if val is not None and val.lower() != VisionEncoderBlockType.NONE.value:
return val.lower()
# If non-vision model, return None
return None


def _get_decoder_block(config_or_name):
"""Extract decoder architecture from config."""
# If input is config, extract decoder name
if hasattr(config_or_name, "decoder_block"):
block = config_or_name.decoder_block
if block != DecoderBlockType.DEFAULT:
return block.value.lower()
# If input is model_name (backward compatibility), find its corresponding config
elif isinstance(config_or_name, str):
val = _get_block_name_from_model_yml(config_or_name, "decoder_block")
if val is not None and val.lower() != DecoderBlockType.DEFAULT.value:
return val.lower()
# If decoder_block not found or is default, return model_name/default
return str(getattr(config_or_name, "model_name", config_or_name)).lower()


def preprocess_mm_data(config):
"""Preprocesses multimodal data based on the provided configuration.
Routes to the appropriate preprocessing function based on the model name.
Routes to the appropriate preprocessing function based on the vision architecture.

Args:
config: A `pyconfig.Config` object containing configuration parameters.
Expand All @@ -28,90 +88,75 @@ def preprocess_mm_data(config):
A `PreprocessorOutput` object containing the processed multimodal data.
"""
processor_outputs = mm_utils.PreprocessorOutput()
vision_block = _get_vision_block(config)

@hengtaoguo hengtaoguo Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Naive question: I wonder why we couldn't simply use block_name = _get_block_name_from_model_yml(config_or_name, "vision_encoder_block") to get either vision/decoder block name? Any reason why we want two separate functions _get_vision_block/_get_decoder_block?

@subawocit subawocit Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question & I actually spent a while thinking it through: vision and text blocks have different logic and fallback behaviors, so it's better keep them separated.

Vision Encoder is optional (only exists for multimodal models). Its default value is VisionEncoderBlockType.NONE. During processor lookup, we want to skip visual processing when no vision encoder is present.

Text Decoder is mandatory. Its default value is DecoderBlockType.DEFAULT. If the decoder block is default, we want it to falls back to standard text handling.


if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
if vision_block in ["gemma3"]:
from maxtext.multimodal.processor_gemma3 import preprocess_mm_data_gemma3 # pylint: disable=import-outside-toplevel

images = [mm_utils.load_image_from_path(p) for p in config.image_path.split(",")]
processor_outputs = preprocess_mm_data_gemma3(images)
elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
elif vision_block in ["gemma4"]:
from maxtext.multimodal.processor_gemma4 import preprocess_mm_data_gemma4 # pylint: disable=import-outside-toplevel

images = [mm_utils.load_image_from_path(p) for p in config.image_path.split(",")]
processor_outputs = preprocess_mm_data_gemma4(images)
elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
elif vision_block in ["llama4"]:
from maxtext.multimodal.processor_llama4 import preprocess_mm_data_llama4 # pylint: disable=import-outside-toplevel

images = [mm_utils.load_image_from_path(p) for p in config.image_path.split(",")]
processor_outputs = preprocess_mm_data_llama4(images)
elif config.model_name in [
"qwen3-omni-30b-a3b",
"qwen3-vl-2b",
"qwen3-vl-4b",
"qwen3-vl-30b-a3b",
"qwen3.5-35b-a3b",
"qwen3.5-397b-a17b",
]:
elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
from maxtext.multimodal.processor_qwen3_omni import preprocess_mm_data_qwen3_omni # pylint: disable=import-outside-toplevel

processor_outputs = preprocess_mm_data_qwen3_omni(config)
else:
raise ValueError(f"Model {config.model_name} not supported for multimodal preprocessing.")
raise ValueError(
f"Model {config.model_name} (vision block {vision_block}) not supported for multimodal preprocessing."
)

return processor_outputs


def preprocess_image_for_training(image, config):
"""Preprocesses a single image for training based on the model name."""
if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
"""Preprocesses a single image for training based on the vision architecture."""
vision_block = _get_vision_block(config)
if vision_block in ["gemma3"]:
from maxtext.multimodal.processor_gemma3 import preprocess_mm_data_gemma3 # pylint: disable=import-outside-toplevel

return preprocess_mm_data_gemma3(image)
elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
elif vision_block in ["gemma4"]:
from maxtext.multimodal.processor_gemma4 import preprocess_mm_data_gemma4 # pylint: disable=import-outside-toplevel

return preprocess_mm_data_gemma4(image)
elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
elif vision_block in ["llama4"]:
from maxtext.multimodal.processor_llama4 import preprocess_mm_data_llama4 # pylint: disable=import-outside-toplevel

return preprocess_mm_data_llama4(image)
elif config.model_name in [
"qwen3-omni-30b-a3b",
"qwen3-vl-2b",
"qwen3-vl-4b",
"qwen3-vl-30b-a3b",
"qwen3.5-35b-a3b",
"qwen3.5-397b-a17b",
]:
elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
from maxtext.multimodal.processor_qwen3_omni import preprocess_mm_data_qwen3_omni_for_training # pylint: disable=import-outside-toplevel

return preprocess_mm_data_qwen3_omni_for_training(image, config)
else:
raise ValueError(f"Model {config.model_name} not supported for image preprocessing.")
raise ValueError(f"Model {config.model_name} (vision block {vision_block}) not supported for image preprocessing.")


def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | None):
"""Get the increase in total token count after inserting image token placeholders"""
if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
vision_block = _get_vision_block(config)

if 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)
elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
elif vision_block in ["gemma4"]:
from maxtext.multimodal.processor_gemma4 import get_image_offsets_gemma4 # pylint: disable=import-outside-toplevel

return get_image_offsets_gemma4(processor_output)
elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
elif vision_block in ["llama4"]:
from maxtext.multimodal.processor_llama4 import get_image_offsets_llama4 # pylint: disable=import-outside-toplevel

return get_image_offsets_llama4(processor_output)
elif config.model_name in [
"qwen3-omni-30b-a3b",
"qwen3-vl-2b",
"qwen3-vl-4b",
"qwen3-vl-30b-a3b",
"qwen3.5-35b-a3b",
"qwen3.5-397b-a17b",
]:
elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
from maxtext.multimodal.processor_qwen3_omni import get_mm_offsets_qwen3_omni # pylint: disable=import-outside-toplevel

return get_mm_offsets_qwen3_omni(config, processor_output)
Expand All @@ -121,26 +166,25 @@ def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | No

def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_placeholder="<|video|>", num_videos=0):
"""Reformat prompt for different models."""
if model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
vision_block = _get_vision_block(model_name)
if vision_block is None:
return prompt

decoder_block = _get_decoder_block(model_name)

@hengtaoguo hengtaoguo Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you have any preliminary idea what we should do if the decoder itself doesn't natively support multimodal? In that case, we don't have (1) image special token/id, (2) image chat template. This could be a follow-up design too.

@subawocit subawocit Aug 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, additional processor functions are required if we have an encoder and a text-only decoder. I already implemented a potential solution and will include it under experimental folder in my next PR!

In short, for (1), split the prompt on <image> tag, tokenize the text pieces with the standard text-only tokenizer, and inject the image token sequence; for (2), use standard text chat templates to handle conversation structure (e.g., user/assistant turns), and use the custom function in (1) to inject/prepend the visual token slots


if decoder_block in ["gemma3"]:
from maxtext.multimodal.processor_gemma3 import reformat_prompt_gemma3 # pylint: disable=import-outside-toplevel

return reformat_prompt_gemma3(prompt, image_placeholder, num_images)
elif model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
elif decoder_block in ["gemma4", "gemma4_small"]:
from maxtext.multimodal.processor_gemma4 import reformat_prompt_gemma4 # pylint: disable=import-outside-toplevel

return reformat_prompt_gemma4(prompt, image_placeholder, num_images)
elif model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
elif decoder_block in ["llama4"]:
from maxtext.multimodal.processor_llama4 import reformat_prompt_llama4 # pylint: disable=import-outside-toplevel

return reformat_prompt_llama4(prompt, image_placeholder, num_images)
elif model_name in [
"qwen3-omni-30b-a3b",
"qwen3-vl-2b",
"qwen3-vl-4b",
"qwen3-vl-30b-a3b",
"qwen3.5-35b-a3b",
"qwen3.5-397b-a17b",
]:
elif decoder_block in ["qwen3", "qwen3_moe", "qwen3_5"]:
from maxtext.multimodal.processor_qwen3_omni import reformat_prompt_qwen3_omni # pylint: disable=import-outside-toplevel

return reformat_prompt_qwen3_omni(
Expand All @@ -156,23 +200,22 @@ def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_pla

def reformat_response(response, model_name):
"""Reformat response for different models."""
if model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
vision_block = _get_vision_block(model_name)
if vision_block is None:
return response

decoder_block = _get_decoder_block(model_name)

if decoder_block in ["llama4"]:
formatted_response = f"{response}<|eot|>"
return formatted_response
elif model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
elif decoder_block in ["gemma3"]:
formatted_response = f"{response}<end_of_turn>"
return formatted_response
elif model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
elif decoder_block in ["gemma4", "gemma4_small"]:
formatted_response = f"{response}<turn|>"
return formatted_response
elif model_name in [
"qwen3-omni-30b-a3b",
"qwen3-vl-2b",
"qwen3-vl-4b",
"qwen3-vl-30b-a3b",
"qwen3.5-35b-a3b",
"qwen3.5-397b-a17b",
]:
elif decoder_block in ["qwen3", "qwen3_moe", "qwen3_5"]:
formatted_response = f"{response}<|im_end|>"
return formatted_response
else:
Expand All @@ -181,53 +224,48 @@ 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."""
if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
vision_block = _get_vision_block(config)
if 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(
tokens, max_num_images=processor_output.num_images # pyrefly: ignore[missing-attribute]
) # pyrefly: ignore[missing-attribute]
elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
elif vision_block in ["gemma4"]:
from maxtext.multimodal.processor_gemma4 import add_extra_tokens_for_images_gemma4 # pylint: disable=import-outside-toplevel

return add_extra_tokens_for_images_gemma4(
tokens, max_num_images=processor_output.num_images # pyrefly: ignore[missing-attribute]
) # pyrefly: ignore[missing-attribute]
elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
elif vision_block in ["llama4"]:
from maxtext.multimodal.processor_llama4 import add_extra_tokens_for_images_llama4 # pylint: disable=import-outside-toplevel
# pyrefly: ignore[bad-argument-type]
return add_extra_tokens_for_images_llama4(tokens, processor_output)
elif config.model_name in [
"qwen3-omni-30b-a3b",
"qwen3-vl-2b",
"qwen3-vl-4b",
"qwen3-vl-30b-a3b",
"qwen3.5-35b-a3b",
"qwen3.5-397b-a17b",
]:

return add_extra_tokens_for_images_llama4(tokens, processor_output) # pyrefly: ignore[bad-argument-type]
elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
from maxtext.multimodal.processor_qwen3_omni import add_extra_tokens_for_qwen3_omni # pylint: disable=import-outside-toplevel

return add_extra_tokens_for_qwen3_omni(tokens, config, processor_output)
else:
raise ValueError(f"Model {config.model_name} does not support multimodal inference.")
raise ValueError(f"Model {config.model_name} (vision block {vision_block}) does not support multimodal inference.")


def get_dummy_image_shape_for_init(model_name, batch_size=1, num_image_per_sequence=1):
"""Return the shape of the dummy image for specific model's initialization."""
image_shape = ()
if model_name.startswith("gemma3"):
vision_block = _get_vision_block(model_name)
if vision_block in ["gemma3"]:
from maxtext.multimodal.processor_gemma3 import get_dummy_image_shape_for_init_gemma3 # pylint: disable=import-outside-toplevel

image_shape = get_dummy_image_shape_for_init_gemma3(batch_size, num_image_per_sequence)
elif model_name.startswith("gemma4"):
elif vision_block in ["gemma4"]:
from maxtext.multimodal.processor_gemma4 import get_dummy_image_shape_for_init_gemma4 # pylint: disable=import-outside-toplevel

image_shape = get_dummy_image_shape_for_init_gemma4(batch_size, num_image_per_sequence)
elif model_name.startswith("llama4"):
elif vision_block in ["llama4"]:
from maxtext.multimodal.processor_llama4 import get_dummy_image_shape_for_init_llama4 # pylint: disable=import-outside-toplevel

image_shape = get_dummy_image_shape_for_init_llama4(batch_size, num_image_per_sequence)
elif model_name.startswith("qwen3-omni") or model_name.startswith("qwen3-vl") or model_name.startswith("qwen3.5"):
elif vision_block in ["qwen3_omni", "qwen3_vl", "qwen3_5"]:
from maxtext.multimodal.processor_qwen3_omni import get_dummy_image_shape_for_init_qwen3_omni # pylint: disable=import-outside-toplevel

image_shape = get_dummy_image_shape_for_init_qwen3_omni(batch_size)
Expand Down Expand Up @@ -256,26 +294,26 @@ def get_dummy_audio_shape_for_init(config):
def get_bidirectional_mask_vision(config, decoder_input_tokens, is_video: bool = False):
"""Get the bidirectional mask for specific models."""
bidirectional_mask_vision = None
if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:

vision_block = _get_vision_block(config)
if vision_block is None:
return bidirectional_mask_vision

decoder_block = _get_decoder_block(config)

if 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
elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
elif decoder_block in ["gemma4", "gemma4_small"]:
from maxtext.multimodal.processor_gemma4 import GEMMA4_TOKEN_PLACEHOLDER # pylint: disable=import-outside-toplevel

bidirectional_mask_vision = decoder_input_tokens == GEMMA4_TOKEN_PLACEHOLDER
elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
elif decoder_block in ["llama4"]:
from maxtext.multimodal.processor_llama4 import LLAMA4_PATCH_TOKEN # pylint: disable=import-outside-toplevel

bidirectional_mask_vision = decoder_input_tokens == LLAMA4_PATCH_TOKEN
elif config.model_name in [
"qwen3-omni-30b-a3b",
"qwen3-vl-2b",
"qwen3-vl-4b",
"qwen3-vl-30b-a3b",
"qwen3.5-35b-a3b",
"qwen3.5-397b-a17b",
]:
elif decoder_block in ["qwen3", "qwen3_moe", "qwen3_5"]:
from maxtext.multimodal.processor_qwen3_omni import QwenTokens # pylint: disable=import-outside-toplevel

tokens = QwenTokens(config)
Expand Down
Loading
Loading