Qwen3.5 vision layers ckpt conversion and decode

Author: hengtaoguo

benchmarks/multimodal/multimodal_eval.py

@@ -230,6 +230,7 @@ def main(config, local_args):
             second_per_grids=processor_output.video_second_per_grid,  # pytype: disable=attribute-error
             spatial_merge_size=config.spatial_merge_size_for_vit,  # pytype: disable=attribute-error
             position_id_per_seconds=config.position_id_per_seconds,
+            config=config,
         )
 
     if true_length > max_prefill_predict_length:

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -1042,6 +1042,64 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
               ): f"model.language_model.layers.{i}.mlp.experts.gate_up_proj",
           }
       )
+
+  # Vision mapping for Qwen3.5
+  if maxtext_config.use_multimodal and "vision_config" in config:
+    vision_config = config["vision_config"]
+    n_vision_layers = vision_config["depth"]
+
+    # Vision patch embedding
+    mapping["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-kernel"] = (
+        "model.visual.patch_embed.proj.weight"
+    )
+    mapping["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-bias"] = (
+        "model.visual.patch_embed.proj.bias"
+    )
+
+    # Vision positional embedding
+    mapping["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-pos_embed_interpolate-pos_embed"] = (
+        "model.visual.pos_embed.weight"
+    )
+
+    # Vision blocks
+    for i in range(n_vision_layers):
+      prefix = f"params-vision_encoder-Qwen3_5MoeVisionEncoder_0-blocks_{i}"
+      hf_prefix = f"model.visual.blocks.{i}"
+
+      # Layer norms
+      mapping[f"{prefix}-ln1-scale"] = f"{hf_prefix}.norm1.weight"
+      mapping[f"{prefix}-ln1-bias"] = f"{hf_prefix}.norm1.bias"
+      mapping[f"{prefix}-ln2-scale"] = f"{hf_prefix}.norm2.weight"
+      mapping[f"{prefix}-ln2-bias"] = f"{hf_prefix}.norm2.bias"
+
+      # Attention
+      mapping[f"{prefix}-attn-attn-query-kernel"] = f"{hf_prefix}.attn.qkv.weight"
+      mapping[f"{prefix}-attn-attn-query-bias"] = f"{hf_prefix}.attn.qkv.bias"
+      mapping[f"{prefix}-attn-attn-key-kernel"] = f"{hf_prefix}.attn.qkv.weight"
+      mapping[f"{prefix}-attn-attn-key-bias"] = f"{hf_prefix}.attn.qkv.bias"
+      mapping[f"{prefix}-attn-attn-value-kernel"] = f"{hf_prefix}.attn.qkv.weight"
+      mapping[f"{prefix}-attn-attn-value-bias"] = f"{hf_prefix}.attn.qkv.bias"
+      mapping[f"{prefix}-attn-attn-out-kernel"] = f"{hf_prefix}.attn.proj.weight"
+      mapping[f"{prefix}-attn-attn-out-bias"] = f"{hf_prefix}.attn.proj.bias"
+
+      # MLP
+      mapping[f"{prefix}-mlp-kernel"] = f"{hf_prefix}.mlp.linear_fc1.weight"
+      mapping[f"{prefix}-mlp-bias"] = f"{hf_prefix}.mlp.linear_fc1.bias"
+      mapping[f"{prefix}-mlp_out-kernel"] = f"{hf_prefix}.mlp.linear_fc2.weight"
+      mapping[f"{prefix}-mlp_out-bias"] = f"{hf_prefix}.mlp.linear_fc2.bias"
+
+    # Vision projector (final merger)
+    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-ln_q-scale"] = "model.visual.merger.norm.weight"
+    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-ln_q-bias"] = "model.visual.merger.norm.bias"
+    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = (
+        "model.visual.merger.linear_fc1.weight"
+    )
+    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-bias"] = "model.visual.merger.linear_fc1.bias"
+    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = (
+        "model.visual.merger.linear_fc2.weight"
+    )
+    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-bias"] = "model.visual.merger.linear_fc2.bias"
+
   return mapping
 
 
@@ -1214,6 +1272,92 @@ def concat_ba_and_transpose(input_tensor, target_shape=None):
     hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = process_wi_0_wi_1
     hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert
 
+  # Vision hooks for Qwen3.5
+  vision_config = config.get("vision_config", None)
+  if vision_config and maxtext_config.use_multimodal:
+    n_vision_layers = vision_config["depth"]
+    hidden_size = vision_config["hidden_size"]
+
+    def reshape_kernel_vision(input_tensor, target_shape):
+      if saving_to_hf:
+        flipped_target_shape = np.flip(np.array(target_shape))
+        return input_tensor.reshape(flipped_target_shape).T
+      else:
+        return input_tensor.T.reshape(target_shape)
+
+    def reshape_conv3d_patch_embed(input_tensor, target_shape):
+      if saving_to_hf:
+        return input_tensor.transpose(4, 3, 0, 1, 2)
+      else:
+        return input_tensor.transpose(2, 3, 4, 1, 0)
+
+    def split_qkv_query(input_tensor, target_shape):
+      if saving_to_hf:
+        raise NotImplementedError("Use fusion hook for MaxText->HF")
+      else:
+        q_weight = input_tensor[:hidden_size, :]
+        return q_weight.T.reshape(target_shape)
+
+    def split_qkv_key(input_tensor, target_shape):
+      if saving_to_hf:
+        raise NotImplementedError("Use fusion hook for MaxText->HF")
+      else:
+        k_weight = input_tensor[hidden_size : 2 * hidden_size, :]
+        return k_weight.T.reshape(target_shape)
+
+    def split_qkv_value(input_tensor, target_shape):
+      if saving_to_hf:
+        raise NotImplementedError("Use fusion hook for MaxText->HF")
+      else:
+        v_weight = input_tensor[2 * hidden_size :, :]
+        return v_weight.T.reshape(target_shape)
+
+    def split_qkv_bias_query(input_tensor, target_shape):
+      if saving_to_hf:
+        raise NotImplementedError("Use fusion hook for MaxText->HF")
+      else:
+        q_bias = input_tensor[:hidden_size]
+        return q_bias.reshape(target_shape)
+
+    def split_qkv_bias_key(input_tensor, target_shape):
+      if saving_to_hf:
+        raise NotImplementedError("Use fusion hook for MaxText->HF")
+      else:
+        k_bias = input_tensor[hidden_size : 2 * hidden_size]
+        return k_bias.reshape(target_shape)
+
+    def split_qkv_bias_value(input_tensor, target_shape):
+      if saving_to_hf:
+        raise NotImplementedError("Use fusion hook for MaxText->HF")
+      else:
+        v_bias = input_tensor[2 * hidden_size :]
+        return v_bias.reshape(target_shape)
+
+    def reshape_vision_attn_out(input_tensor, target_shape):
+      if saving_to_hf:
+        return input_tensor.reshape(hidden_size, hidden_size).T
+      else:
+        return input_tensor.T.reshape(target_shape)
+
+    # Apply vision hooks
+    hooks["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-kernel"] = reshape_conv3d_patch_embed
+
+    for i in range(n_vision_layers):
+      prefix = f"params-vision_encoder-Qwen3_5MoeVisionEncoder_0-blocks_{i}"
+      hooks[f"{prefix}-attn-attn-query-kernel"] = split_qkv_query
+      hooks[f"{prefix}-attn-attn-query-bias"] = split_qkv_bias_query
+      hooks[f"{prefix}-attn-attn-key-kernel"] = split_qkv_key
+      hooks[f"{prefix}-attn-attn-key-bias"] = split_qkv_bias_key
+      hooks[f"{prefix}-attn-attn-value-kernel"] = split_qkv_value
+      hooks[f"{prefix}-attn-attn-value-bias"] = split_qkv_bias_value
+      hooks[f"{prefix}-attn-attn-out-kernel"] = reshape_vision_attn_out
+      hooks[f"{prefix}-mlp-kernel"] = reshape_kernel_vision
+      hooks[f"{prefix}-mlp_out-kernel"] = reshape_kernel_vision
+
+    # Vision projector
+    hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = reshape_kernel_vision
+    hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = reshape_kernel_vision
+
   return hooks
 
 

src/maxtext/configs/models/qwen3.5-35b-a3b.yml

@@ -48,3 +48,23 @@ partial_rotary_factor: 0.25
 
 # General Model Settings
 enable_dropout: False
+
+# Vision Encoder Configuration (need to set use_multimodal=true)
+# Based on Qwen3.5 MoE Vision Model Config
+image_size_for_vit: 768
+hidden_size_for_vit: 1152
+intermediate_size_for_vit: 4304
+num_attention_heads_for_vit: 16
+num_hidden_layers_for_vit: 27
+num_channels_for_vit: 3
+patch_size_for_vit: 16
+temporal_patch_size_for_vit: 2
+spatial_merge_size_for_vit: 2
+out_hidden_size_for_vit: 2048  # Projects to decoder emb_dim (2048)
+num_position_embeddings_for_vit: 2304
+deepstack_visual_indexes_for_vit: []  # No deepstack for Qwen3.5 VL
+rope_theta_for_vit: 10000
+
+# MRoPE Settings (Multi-dimensional RoPE for multimodal)
+use_mrope: true
+mrope_section: [11, 11, 10]

src/maxtext/configs/types.py

@@ -262,8 +262,8 @@ class ProfilerType(str, Enum):
     "qwen3-next-80b-a3b",
     "qwen3-omni-30b-a3b",
     "qwen3-custom-30b-a3b",
-    "qwen3.5-397b-a17b",
     "qwen3.5-35b-a3b",
+    "qwen3.5-397b-a17b",
     "gpt3-175b",
     "gpt3-22b",
     "gpt3-6b",
@@ -2927,6 +2927,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "llama4-17b-16e",
           "llama4-17b-128e",
           "qwen3-omni-30b-a3b",
+          "qwen3.5-35b-a3b",
           "qwen3.5-397b-a17b",
       )
       if self.model_name not in valid_mm_models and self.model_name != "default":

src/maxtext/inference/decode.py

@@ -154,6 +154,7 @@ def main(argv: Sequence[str]) -> None:
           second_per_grids=processor_outputs.video_second_per_grid,  # pytype: disable=attribute-error
           spatial_merge_size=config.spatial_merge_size_for_vit,  # pytype: disable=attribute-error
           position_id_per_seconds=config.position_id_per_seconds,
+          config=config,
       )
 
   assert (

src/maxtext/layers/decoders.py

@@ -660,6 +660,7 @@ def _apply_embedding(
             "llama4-17b-16e",
             "llama4-17b-128e",
             "qwen3-omni-30b-a3b",
+            "qwen3.5-35b-a3b",
             "qwen3.5-397b-a17b",
         ]:
           y = mm_utils.merge_mm_embeddings(
@@ -673,7 +674,7 @@ def _apply_embedding(
           raise ValueError(f"Unsupported model_name for multimodal: {cfg.model_name}")
 
       if video_embeddings is not None and cfg.use_multimodal:
-        if cfg.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-397b-a17b"]:
+        if cfg.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
           y = mm_utils.merge_mm_embeddings(
               text_embeddings=y,
               multimodal_embeddings=video_embeddings,

src/maxtext/layers/encoders.py

@@ -70,7 +70,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-397b-a17b"]:
+    elif self.config.model_name in ["qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
       from maxtext.models import qwen3_5_vision  # pylint: disable=import-outside-toplevel
 
       encoder_name = "Qwen3_5MoeVisionEncoder_0"

src/maxtext/models/qwen3.py

@@ -734,6 +734,8 @@ def __init__(
         use_qk_norm=cfg.use_qk_norm,
         query_pre_attn_scalar=scaling_factor,
         model_mode=model_mode,
+        use_mrope=cfg.use_mrope,
+        mrope_section=cfg.mrope_section,
         rngs=rngs,
     )
 

src/maxtext/multimodal/processor.py

@@ -44,7 +44,7 @@ def preprocess_mm_data(config):
 
     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.5-397b-a17b"]:
+  elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
     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)
@@ -68,7 +68,7 @@ def preprocess_image_for_training(image, model_name):
     from maxtext.multimodal.processor_llama4 import preprocess_mm_data_llama4  # pylint: disable=import-outside-toplevel
 
     return preprocess_mm_data_llama4(image)
-  elif model_name in ["qwen3-omni-30b-a3b", "qwen3.5-397b-a17b"]:
+  elif model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
     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)
@@ -90,7 +90,7 @@ def get_image_offsets(config, processor_output: mm_utils.PreprocessorOutput | No
     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.5-397b-a17b"]:
+  elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
     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)
@@ -112,7 +112,7 @@ def reformat_prompt(prompt, image_placeholder, model_name, num_images, video_pla
     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.5-397b-a17b"]:
+  elif model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
     from maxtext.multimodal.processor_qwen3_omni import reformat_prompt_qwen3_omni  # pylint: disable=import-outside-toplevel
 
     return reformat_prompt_qwen3_omni(
@@ -137,7 +137,7 @@ def reformat_response(response, model_name):
   elif model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]:
     formatted_response = f"{response}<end_of_turn>"
     return formatted_response
-  elif model_name in ["qwen3-omni-30b-a3b", "qwen3.5-397b-a17b"]:
+  elif model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
     formatted_response = f"{response}<|im_end|>"
     return formatted_response
   else:
@@ -158,7 +158,7 @@ def prepare_text_for_image_fusion(tokens, config, processor_output=None):
     from maxtext.multimodal.processor_llama4 import add_extra_tokens_for_images_llama4  # pylint: disable=import-outside-toplevel
 
     return add_extra_tokens_for_images_llama4(tokens, processor_output)
-  elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-397b-a17b"]:
+  elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
     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)
@@ -181,7 +181,7 @@ def get_dummy_image_shape_for_init(model_name, batch_size=1, num_image_per_seque
     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-30b-a3b") or model_name.startswith("qwen3.5-397b-a17b"):
+  elif model_name.startswith("qwen3-omni-30b-a3b") or model_name.startswith("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)
@@ -222,22 +222,26 @@ def get_bidirectional_mask_vision(config, decoder_input_tokens, is_video: bool =
     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.5-397b-a17b"]:
-    from maxtext.multimodal.processor_qwen3_omni import QWEN3_OMNI_IMAGE_TOKEN, QWEN3_OMNI_VIDEO_TOKEN  # pylint: disable=import-outside-toplevel
+  elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
+    from maxtext.multimodal.processor_qwen3_omni import QwenTokens  # pylint: disable=import-outside-toplevel
+
+    tokens = QwenTokens(config)
 
     if is_video:
-      bidirectional_mask_vision = decoder_input_tokens == QWEN3_OMNI_VIDEO_TOKEN
+      bidirectional_mask_vision = decoder_input_tokens == tokens.video_pad
     else:
-      bidirectional_mask_vision = decoder_input_tokens == QWEN3_OMNI_IMAGE_TOKEN
+      bidirectional_mask_vision = decoder_input_tokens == tokens.image_pad
   return bidirectional_mask_vision
 
 
 def get_bidirectional_mask_audio(config, decoder_input_tokens):
   """Get the bidirectional mask for specific models."""
   bidirectional_mask_audio = None
   if config.model_name in ["qwen3-omni-30b-a3b"]:
-    from maxtext.multimodal.processor_qwen3_omni import QWEN3_OMNI_AUDIO_TOKEN  # pylint: disable=import-outside-toplevel
+    from maxtext.multimodal.processor_qwen3_omni import QwenTokens  # pylint: disable=import-outside-toplevel
+
+    tokens = QwenTokens(config)
 
     # Create bidirectional_mask for audio token merging
-    bidirectional_mask_audio = decoder_input_tokens == QWEN3_OMNI_AUDIO_TOKEN
+    bidirectional_mask_audio = decoder_input_tokens == tokens.audio_pad
   return bidirectional_mask_audio