Decouple vision & decoder block routing in multimodal processor - #4746
Decouple vision & decoder block routing in multimodal processor#4746subawocit wants to merge 1 commit into
Conversation
90cd020 to
71f11b1
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
71f11b1 to
8bb96e9
Compare
23c17ec to
e7af676
Compare
hengtaoguo
left a comment
hengtaoguo
left a comment
There was a problem hiding this comment.
Overall I love this refactor, thanks for making codes cleaner!
|
|
||
|
|
||
| @functools.lru_cache(maxsize=None) | ||
| def _get_block_name_from_model_yml(model_name: str, block_name: str) -> str | None: |
There was a problem hiding this comment.
The block_name could only be chosen from ['vision_encoder_block', 'decoder_block']? Could we add a docstring to explain this?
There was a problem hiding this comment.
Good idea - I just updated the description
| A `PreprocessorOutput` object containing the processed multimodal data. | ||
| """ | ||
| processor_outputs = mm_utils.PreprocessorOutput() | ||
| vision_block = _get_vision_block(config) |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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 vision_block is None: | ||
| return prompt | ||
|
|
||
| decoder_block = _get_decoder_block(model_name) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
e7af676 to
0706ef6
Compare
aireenmei
left a comment
aireenmei
left a comment
There was a problem hiding this comment.
LGTM! Thanks for testing decode on different models!
| self.assertIsNone(mm_processor._get_vision_block("llama2-7b")) | ||
| self.assertEqual(mm_processor._get_decoder_block("llama2-7b"), "llama2") | ||
|
|
||
| def test_get_vision_and_decoder_block_routing_from_config(self): |
There was a problem hiding this comment.
How about adding a test for stitched config?
3 participants
Description
This PR refactors multimodal preprocessing routing in MaxText by replacing rigid model-name-based whitelists with component-based whitelists.
This change decouples vision encoder blocks from LLM decoder blocks, and gives a modular foundation for future mix-and-match implementations while eliminating the maintenance needed to manually update the model whitelisting every time a new model or a new stitched model is added.
This is Step 2.5 of a 5-step Proof-of-Concept for any-to-any multimodal alignment in MaxText. Overall goal is to align a pretrained vision encoders with another text-only LLM, and connecting the two with a dynamically configured adaptation layer (customized MLP connector), special tokenizer mapping for visual placeholder tokens, and train only the MLP using supervised fine tuning.
FIXES: b/543524347
Changes
src/maxtext/multimodal/processor.py_get_vision_block()and_get_decoder_block()to get each model's encoder and decoder block names respectively._get_block_name_from_model_yml).tests/unit/multimodal_utils_test.pyTests
Unit tests
test_get_vision_and_decoder_block_routingandtest_get_vision_and_decoder_block_routing_from_configdeselected):53 passed, 1 skipped, 2 deselected55 passed, 1 skippedModel decoding
Verify if the model decoding responses are identical before and after new processor.py implementation (to ensure that new processor changes do not affect the model's functionality).
Qwen-VL-2B
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml model_name=qwen3-vl-2b \ tokenizer_path=Qwen/Qwen3-VL-2B-Instruct tokenizer_type=huggingface use_multimodal=true \ load_parameters_path=gs://yuchenhou-maxtext-logs/checkpoints/qwen3-vl-2b-processor/0/items \ prompt='Describe this image' image_path='tests/assets/test_image.jpg' \ per_device_batch_size=1 scan_layers=false max_prefill_predict_length=512 max_target_length=768 \ ici_tensor_parallelism=4 override_model_config=true attention='dot_product' hf_access_token=<your-token>Original processor implmentation
Input
<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe this image<|im_end|>\n<|im_start|>assistant\n->This is a panoramic view of the Seattle skyline on a bright, sunny day. The image is taken from a high vantage point, looking down on the city.\n\nThe most prominent feature is the **Space Needle**, a distinctive observation tower located on the left side of the frame. It stands tall among the city's modern skyscrapers, which are a mix of glass and steel structures. The buildings are densely packed, creating a dense urban landscape.New processor implmentation
Input
<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe this image<|im_end|>\n<|im_start|>assistant\n->This is a panoramic view of the Seattle skyline on a bright, sunny day. The image is taken from a high vantage point, looking down on the city.\n\nThe most prominent feature is the **Space Needle**, a distinctive observation tower located on the left side of the frame. It stands tall among the city's modern skyscrapers, which are a mix of glass and steel structures. The buildings are densely packed, creating a dense urban landscape.Qwen3-0.6B
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml model_name=qwen3-0.6b \ tokenizer_path=src/maxtext/assets/tokenizers/qwen3-tokenizer \ load_parameters_path=gs://yuchenhou-maxtext-logs/checkpoints/qwen3-0.6b-processor/0/items \ prompt='How many r are there in the word strawberry' \ per_device_batch_size=1 scan_layers=false max_prefill_predict_length=512 max_target_length=768 \ checkpoint_storage_use_ocdbt=false checkpoint_storage_use_zarr3=false \ ici_tensor_parallelism=4 override_model_config=true attention='dot_product' hf_access_token=<your-token>Original
Input
How many r are there in the word strawberry->? Let's see, the word is strawberry. Let's break it down. The letters are S, T, R, E, W, A, R, T, E, S. So, how many times does the letter R appear? Let's count. First, S, T, R, E, W, A, R, T, E, S. So, the first R is at position 3, then another R at position 7. So that's two R's. Therefore, the answer should be 2. But wait, let me double-check. Maybe I missed something. Let me write them out again: S, T, R, E, W, A, R, T, E, S. Yes, two R's. So the answer is 2.\nThe answer is 2.New
Input
How many r are there in the word strawberry->? Let's see, the word is strawberry. Let's break it down. The letters are S, T, R, E, W, A, R, T, E, S. So, how many times does the letter R appear? Let's count. First, S, T, R, E, W, A, R, T, E, S. So, the first R is at position 3, then another R at position 7. So that's two R's. Therefore, the answer should be 2. But wait, let me double-check. Maybe I missed something. Let me write them out again: S, T, R, E, W, A, R, T, E, S. Yes, two R's. So the answer is 2.\nThe answer is 2.Gemma2-2B
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml model_name=gemma2-2b \ tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma \ load_parameters_path=gs://yuchenhou-maxtext-logs/checkpoints/gemma2-2b-processor/0/items \ prompt='How many r are there in the word strawberry' \ per_device_batch_size=1 scan_layers=false max_prefill_predict_length=512 max_target_length=768 \ checkpoint_storage_use_ocdbt=false checkpoint_storage_use_zarr3=false \ ici_tensor_parallelism=4 override_model_config=true attention='dot_product' hf_access_token=<your-token>Original
Input
How many r are there in the word strawberry->?\n[Answer 1]\nThere are <strong>10</strong> r's in the word strawberry.\n<blockquote>The word strawberry has 10 r's.</blockquote>New
Input
How many r are there in the word strawberry->?\n[Answer 1]\nThere are <strong>10</strong> r's in the word strawberry.\n<blockquote>The word strawberry has 10 r's.</blockquote>Gemma3-4B
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml model_name=gemma3-4b \ tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma3 use_multimodal=true \ load_parameters_path=gs://yuchenhou-maxtext-logs/checkpoints/gemma3-4b-processor/0/items \ prompt='Describe this image' image_path='tests/assets/test_image.jpg' \ per_device_batch_size=1 scan_layers=false max_prefill_predict_length=512 max_target_length=768 \ checkpoint_storage_use_ocdbt=false checkpoint_storage_use_zarr3=false \ ici_tensor_parallelism=4 override_model_config=true attention='dot_product' hf_access_token=<your-token>Original
Input
<start_of_turn>user\n\nDescribe this image<end_of_turn>\n<start_of_turn>model\n->Here's a description of the image:\n\n**Overall Impression:**\n\nThe image captures a stunning panoramic view of Seattle, Washington, on a bright, sunny day. The city skyline dominates the foreground, with the majestic Cascade Mountains visible in the distance.New
Input
<start_of_turn>user\n\nDescribe this image<end_of_turn>\n<start_of_turn>model\n->Here's a description of the image:\n\n**Overall Impression:**\n\nThe image captures a stunning panoramic view of Seattle, Washington, on a bright, sunny day. The city skyline dominates the foreground, with the majestic Cascade Mountains visible in the distance.Gemma4-e2b
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml model_name=gemma4-e2b \ tokenizer_path=src/maxtext/assets/tokenizers/tokenizer_gemma4.model tokenizer_type=sentencepiece \ load_parameters_path=gs://yuchenhou-maxtext-logs/checkpoints/gemma4-e2b-processor/0/items \ prompt="<bos><|turn>system\nYou are a helpful assistant.<turn|>\n<|turn>user\nHow many r are there in the word strawberry<turn|>\n<|turn>model\n" \ per_device_batch_size=1 scan_layers=false use_multimodal=false max_prefill_predict_length=512 max_target_length=768 \ checkpoint_storage_use_ocdbt=False checkpoint_storage_use_zarr3=False ici_tensor_parallelism=1 \ override_model_config=true attention=dot_product decode_sampling_strategy=composite \ decode_sampling_temperature=1.0 decode_sampling_nucleus_p=0.95 decode_sampling_top_k=64Original
Input
<bos><|turn>system You are a helpful assistant.<turn|> <|turn>user How many r are there in the word strawberry<turn|> <|turn>model->there are **3** "r"s in the word "strawberry".<turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|>New
Input
<bos><|turn>system You are a helpful assistant.<turn|> <|turn>user How many r are there in the word strawberry<turn|> <|turn>model->there are **3** "r"s in the word "strawberry".<turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|>Gemma4-26B
python3 -m maxtext.inference.decode src/maxtext/configs/base.yml model_name=gemma4-26b \ tokenizer_type="huggingface" tokenizer_path=google/gemma-4-26b-a4b-it use_multimodal=true \ load_parameters_path=gs://yuchenhou-maxtext-logs/checkpoints/gemma4-26b-processor/0/items \ prompt='Describe this image <|image|>' image_path='tests/assets/test_image.jpg' \ per_device_batch_size=1 scan_layers=false max_prefill_predict_length=512 max_target_length=768 \ checkpoint_storage_use_ocdbt=false checkpoint_storage_use_zarr3=false \ ici_fsdp_parallelism=4 parameter_memory_host_offload=True override_model_config=true attention='dot_product' hf_access_token=<your-token>Original
Input
<bos><|turn>user\nDescribe this image <|image|><turn|>\n<|turn>model\n->thought\nA wide, eye-level shot shows the Seattle skyline under a bright blue sky with white clouds. The Space Needle, a tall, white, futuristic-looking tower with a saucer-shaped top, is on the left side of the frame. Behind it, a dense cluster of skyscrapers of various heights and colors, mostly blue and gray, fills the center and left of the frame. In the background, a range of snow-capped mountains is visible under a layer of white clouds. In the foreground, several multi-story buildings in shades of white, gray, and brown are partially obscured by lush green trees. The lighting is bright and even, suggesting a sunny day.New
Input
<bos><|turn>user\nDescribe this image <|image|><turn|>\n<|turn>model\n->thought\nA wide, eye-level shot shows the Seattle skyline under a bright blue sky with white clouds. The Space Needle, a tall, white, futuristic-looking tower with a saucer-shaped top, is on the left side of the frame. Behind it, a dense cluster of skyscrapers of various heights and colors, mostly blue and gray, fills the center and left of the frame. In the background, a range of snow-capped mountains is visible under a layer of white clouds. In the foreground, several multi-story buildings in shades of white, gray, and brown are partially obscured by lush green trees. The lighting is bright and even, suggesting a sunny day.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.