Make vision_encoder_block configurable

hengtaoguo · hengtaoguo · commit faf38eb75e2b · 2026-07-17T18:57:27.000Z
diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py
@@ -116,6 +116,18 @@ class DecoderBlockType(enum.Enum):
   DEEPSEEK4 = "deepseek4"
 
 
+class VisionEncoderBlockType(enum.Enum):
+  """Vision encoder block types."""
+
+  NONE = "none"
+  GEMMA3 = "gemma3"
+  LLAMA4 = "llama4"
+  QWEN3_OMNI = "qwen3_omni"
+  GEMMA4 = "gemma4"
+  QWEN3_5 = "qwen3_5"
+  QWEN3_VL = "qwen3_vl"
+
+
 class AttentionType(enum.Enum):
   GLOBAL = "global"  # default, with causality
   LOCAL_SLIDING = "local_sliding"
diff --git a/src/maxtext/configs/models/gemma3-12b.yml b/src/maxtext/configs/models/gemma3-12b.yml
@@ -33,6 +33,7 @@ rope_max_timescale: 1_000_000
 rope_linear_scaling_factor: 8.0
 
 # Multimodal flags (need to set use_multimodal=true)
+vision_encoder_block: "gemma3"
 image_size_for_vit: 896
 num_channels_for_vit: 3
 patch_size_for_vit: 14
diff --git a/src/maxtext/configs/models/gemma3-27b.yml b/src/maxtext/configs/models/gemma3-27b.yml
@@ -33,6 +33,7 @@ rope_max_timescale: 1_000_000
 rope_linear_scaling_factor: 8.0
 
 # Multimodal flags (need to set use_multimodal=true)
+vision_encoder_block: "gemma3"
 image_size_for_vit: 896
 num_channels_for_vit: 3
 patch_size_for_vit: 14
diff --git a/src/maxtext/configs/models/gemma3-4b.yml b/src/maxtext/configs/models/gemma3-4b.yml
@@ -33,6 +33,7 @@ rope_max_timescale: 1_000_000
 rope_linear_scaling_factor: 8.0
 
 # Multimodal flags (need to set use_multimodal=true)
+vision_encoder_block: "gemma3"
 image_size_for_vit: 896
 num_channels_for_vit: 3
 patch_size_for_vit: 14
diff --git a/src/maxtext/configs/models/gemma4-26b.yml b/src/maxtext/configs/models/gemma4-26b.yml
@@ -48,6 +48,7 @@ norm_topk_prob: true
 load_balance_loss_weight: 0.001
 
 # Multimodal flags (need to set use_multimodal=true)
+vision_encoder_block: "gemma4"
 rope_theta_for_vit: 100
 image_size_for_vit: [672, 960]
 num_channels_for_vit: 3
diff --git a/src/maxtext/configs/models/gemma4-31b.yml b/src/maxtext/configs/models/gemma4-31b.yml
@@ -44,6 +44,7 @@ local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
 
 # Multimodal flags (need to set use_multimodal=true)
+vision_encoder_block: "gemma4"
 rope_theta_for_vit: 100
 image_size_for_vit: [672, 960]
 num_channels_for_vit: 3
diff --git a/src/maxtext/configs/models/gemma4-e2b.yml b/src/maxtext/configs/models/gemma4-e2b.yml
@@ -45,6 +45,7 @@ local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
 
 # Vision encoder flags — multimodal not yet supported for E2B / E4B.
+vision_encoder_block: "gemma4"
 rope_theta_for_vit: 100
 image_size_for_vit: [672, 960]
 num_channels_for_vit: 3
diff --git a/src/maxtext/configs/models/gemma4-e4b.yml b/src/maxtext/configs/models/gemma4-e4b.yml
@@ -46,6 +46,7 @@ local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
 
 # Vision encoder flags — multimodal not yet supported for E2B / E4B.
+vision_encoder_block: "gemma4"
 rope_theta_for_vit: 100
 image_size_for_vit: [672, 960]
 num_channels_for_vit: 3
diff --git a/src/maxtext/configs/models/llama4-17b-128e.yml b/src/maxtext/configs/models/llama4-17b-128e.yml
@@ -43,6 +43,7 @@ temperature_tuning: true
 chunk_attn_window_size: 8192
 
 # Multimodal flags (need to set use_multimodal=true)
+vision_encoder_block: "llama4"
 image_size_for_vit: 336
 num_channels_for_vit: 3
 patch_size_for_vit: 14
diff --git a/src/maxtext/configs/models/llama4-17b-16e.yml b/src/maxtext/configs/models/llama4-17b-16e.yml
@@ -43,6 +43,7 @@ temperature_tuning: true
 chunk_attn_window_size: 8192
 
 # Multimodal flags (need to set use_multimodal=true)
+vision_encoder_block: "llama4"
 image_size_for_vit: 336
 num_channels_for_vit: 3
 patch_size_for_vit: 14
diff --git a/src/maxtext/configs/models/qwen3-omni-30b-a3b.yml b/src/maxtext/configs/models/qwen3-omni-30b-a3b.yml
@@ -42,6 +42,7 @@ enable_dropout: false
 scan_layers: false  # deepstack does not support scan_layers
 
 # Vision Encoder Configuration
+vision_encoder_block: "qwen3_omni"
 # Based on https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
 image_size_for_vit: 768
 hidden_size_for_vit: 1152
diff --git a/src/maxtext/configs/models/qwen3-vl-2b.yml b/src/maxtext/configs/models/qwen3-vl-2b.yml
@@ -36,6 +36,7 @@ rope_max_timescale: 5000000
 enable_dropout: false
 
 # Vision Encoder Configuration
+vision_encoder_block: "qwen3_vl"
 # Based on HuggingFace AutoConfig for Qwen/Qwen3-VL-2B-Instruct
 use_multimodal: true
 image_size_for_vit: 768
diff --git a/src/maxtext/configs/models/qwen3-vl-4b.yml b/src/maxtext/configs/models/qwen3-vl-4b.yml
@@ -36,6 +36,7 @@ rope_max_timescale: 5000000
 enable_dropout: false
 
 # Vision Encoder Configuration
+vision_encoder_block: "qwen3_vl"
 # Based on HuggingFace AutoConfig for Qwen/Qwen3-VL-4B-Instruct
 use_multimodal: true
 image_size_for_vit: 768
diff --git a/src/maxtext/configs/models/qwen3.5-35b-a3b.yml b/src/maxtext/configs/models/qwen3.5-35b-a3b.yml
@@ -50,6 +50,7 @@ partial_rotary_factor: 0.25
 enable_dropout: False
 
 # Vision Encoder Configuration (need to set use_multimodal=true)
+vision_encoder_block: "qwen3_5"
 # Based on Qwen3.5 MoE Vision Model Config
 image_size_for_vit: 768
 hidden_size_for_vit: 1152
diff --git a/src/maxtext/configs/models/qwen3.5-397b-a17b.yml b/src/maxtext/configs/models/qwen3.5-397b-a17b.yml
@@ -51,6 +51,7 @@ partial_rotary_factor: 0.25
 enable_dropout: false
 
 # Vision Encoder Configuration (need to set use_multimodal=true)
+vision_encoder_block: "qwen3_5"
 # Based on Qwen3.5 MoE Vision Model Config
 image_size_for_vit: 768
 hidden_size_for_vit: 1152
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
@@ -29,7 +29,7 @@
 from typing import Any, Literal, NewType, Optional
 
 import jax
-from maxtext.common.common_types import AttentionType, DecoderBlockType, ReorderStrategy, ShardMode, CustomRule
+from maxtext.common.common_types import AttentionType, DecoderBlockType, ReorderStrategy, ShardMode, CustomRule, VisionEncoderBlockType
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_utils
 from maxtext.utils import elastic_utils
@@ -2022,6 +2022,10 @@ class MultimodalGeneral(BaseModel):
 
   use_multimodal: bool = Field(False, description="Enable multimodal capabilities.")
   attention_for_vit: str = Field("dot_product", description="The attention algorithm to use for vision encoder.")
+  vision_encoder_block: VisionEncoderBlockType = Field(
+      VisionEncoderBlockType.NONE,
+      description="The style of VisionEncoderBlock to use (e.g., 'gemma3', 'llama4').",
+  )
   freeze_vision_encoder_params: bool = Field(True, description="Freeze the parameters of the vision encoder.")
   freeze_audio_encoder_params: bool = Field(True, description="Freeze the parameters of the audio encoder.")
   use_audio: bool = Field(False, description="Enable audio encoder for multimodal models.")
diff --git a/src/maxtext/layers/encoders.py b/src/maxtext/layers/encoders.py
@@ -18,7 +18,7 @@
 from flax import nnx
 from jax.sharding import Mesh
 
-from maxtext.common.common_types import Config
+from maxtext.common.common_types import Config, VisionEncoderBlockType
 from maxtext.layers import nnx_wrappers
 from maxtext.layers import initializers
 
@@ -34,31 +34,33 @@ def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs):
 
   def _setup_vision_encoder_layers(self):
     """Setup vision encoder layers specific to the model, instantiate NNX modules."""
-    if self.config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]:
+    self.vision_encoder_block = self.config.vision_encoder_block
+
+    if self.vision_encoder_block == VisionEncoderBlockType.GEMMA3:
       from maxtext.models import gemma3  # pylint: disable=import-outside-toplevel
 
       encoder_name = "Gemma3VisionEncoderLayer_0"
       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.config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]:
+    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.config.model_name in ["qwen3-omni-30b-a3b"]:
+    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.config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
+    elif self.vision_encoder_block == VisionEncoderBlockType.GEMMA4:
       from maxtext.models import gemma4_vision  # pylint: disable=import-outside-toplevel
 
       encoder_name = "Gemma4VisionEncoderLayer_0"
@@ -70,7 +72,7 @@ def _setup_vision_encoder_layers(self):
           self, projector_name, gemma4_vision.Gemma4VisionProjector(config=self.config, mesh=self.mesh, rngs=self.rngs)
       )
       return encoder_name, projector_name
-    elif self.config.model_name in ["qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
+    elif self.vision_encoder_block == VisionEncoderBlockType.QWEN3_5:
       from maxtext.models import qwen3_5_vision  # pylint: disable=import-outside-toplevel
 
       encoder_name = "Qwen3_5MoeVisionEncoder_0"
@@ -80,7 +82,7 @@ def _setup_vision_encoder_layers(self):
       )
       setattr(self, projector_name, qwen3_5_vision.Qwen3_5MoeVisionProjector(config=self.config, rngs=self.rngs))
       return encoder_name, projector_name
-    elif self.config.model_name in ["qwen3-vl-4b", "qwen3-vl-2b"]:
+    elif self.vision_encoder_block == VisionEncoderBlockType.QWEN3_VL:
       from maxtext.models import qwen3_vl_vision  # pylint: disable=import-outside-toplevel
 
       encoder_name = "Qwen3VLVisionEncoder_0"
@@ -91,12 +93,16 @@ def _setup_vision_encoder_layers(self):
       setattr(self, projector_name, qwen3_vl_vision.Qwen3VLVisionProjector(config=self.config, rngs=self.rngs))
       return encoder_name, projector_name
     else:
-      raise ValueError(f"No VisionEncoder implemented for {self.config.model_name} yet")
+      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}."
+      )
 
   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)
-    if self.config.model_name.startswith("qwen3") and input_masks is not None:
+    if self.vision_encoder_block.value.startswith("qwen3") and input_masks is not None:
       encoder_output = encoder(
           input_images, video_mask=input_masks, video_grid_thw=video_grid_thw, deterministic=deterministic
       )
