Merge pull request AI-Hypercomputer#4251 from AI-Hypercomputer:hengtaoguo-resize

PiperOrigin-RevId: 937714775

Author: Google-ML-Automation

src/maxtext/configs/base.yml

@@ -1167,6 +1167,9 @@ image_placeholder: "<|image|>"
 video_placeholder: "<|video|>"
 audio_placeholder: "<|audio|>"
 use_audio_in_video: false
+video_max_grid_t: null
+video_max_grid_h: null
+video_max_grid_w: null
 posemb_type_for_vit: "learn"
 # max_num_images_per_example only applies for training when your image column is a list of images.
 # -1 means no limit, and will pad to the max possible number of images determined by sequence length.

src/maxtext/configs/types.py

@@ -1955,6 +1955,25 @@ class MultimodalGeneral(BaseModel):
   use_mrope: bool = Field(False, description="Enable Multi-dimensional RoPE for Qwen3-Omni models.")
   mrope_section: list[int] = Field([24, 20, 20], description="Dimensions for temporal, height, width in MRoPE.")
   position_id_per_seconds: int = Field(25, description="Temporal granularity for MRoPE (tokens per second).")
+  video_max_grid_t: Optional[int] = Field(
+      None,
+      description="Maximum Qwen3-Omni/3.5 video grid size in temporal patches, before spatial merge.",
+  )
+  video_max_grid_h: Optional[int] = Field(
+      None,
+      description="Maximum Qwen3-Omni/3.5 video grid size in height patches, before spatial merge.",
+  )
+  video_max_grid_w: Optional[int] = Field(
+      None,
+      description="Maximum Qwen3-Omni/3.5 video grid size in width patches, before spatial merge.",
+  )
+
+  @model_validator(mode="after")
+  def validate_video_max_grid(self) -> "MultimodalGeneral":
+    max_grid = (self.video_max_grid_t, self.video_max_grid_h, self.video_max_grid_w)
+    if any(dim is None for dim in max_grid) and not all(dim is None for dim in max_grid):
+      raise ValueError("video_max_grid_t, video_max_grid_h, and video_max_grid_w must be set together.")
+    return self
 
 
 class VisionTower(BaseModel):

src/maxtext/models/qwen3.py

@@ -1866,19 +1866,29 @@ def __init__(
         rngs=rngs,
     )
 
-  def __call__(self, hidden_states: Array) -> Array:
+  def __call__(self, hidden_states: Array, video_mask: Array | None = None) -> tuple[Array, Array | None]:
     """
     Args:
         hidden_states: Input tensor of shape (batch, in_channels, temporal*patch_size, height*patch_size, width*patch_size)
+        video_mask: Optional pixel-level mask with shape
+          (batch, 1, temporal*patch_size, height*patch_size, width*patch_size).
     Returns:
-        Output tensor of shape (batch, T*H*W, embed_dim) where T, H, W are the number of patches
+        Tuple of:
+        - Output tensor of shape (batch, T*H*W, embed_dim) where T, H, W are the number of patches
+        - Attention mask of shape (batch, T*H*W), or None when video_mask is not provided
     """
     hidden_states = jnp.transpose(hidden_states, (0, 2, 3, 4, 1))
     hidden_states = self.proj(hidden_states)
     batch_size = hidden_states.shape[0]
     seq_len = hidden_states.shape[1] * hidden_states.shape[2] * hidden_states.shape[3]
     hidden_states = hidden_states.reshape(batch_size, seq_len, self.embed_dim)
-    return hidden_states
+
+    attention_mask = None
+    if video_mask is not None:
+      patch_mask = video_mask[:, 0, :: self.temporal_patch_size, :: self.patch_size, :: self.patch_size]
+      attention_mask = patch_mask.reshape(video_mask.shape[0], -1).astype(jnp.int32)
+
+    return hidden_states, attention_mask
 
 
 class Qwen3OmniMoeVisionAttention(nnx.Module):
@@ -2102,7 +2112,7 @@ def __call__(
         self.config.patch_size_for_vit,
     )
 
-    x = self.patch_embed(hidden_states)
+    x, _ = self.patch_embed(hidden_states)
     x = x.reshape(batch_size, -1, self.config.hidden_size_for_vit)
     pos = self.pos_embed_interpolate(num_frames, height, width)
 

src/maxtext/multimodal/processor_qwen3_omni.py

@@ -129,13 +129,75 @@ class Qwen3OmniPreprocessorOutput(mm_utils.PreprocessorOutput):
   video_values: None | np.ndarray = None
   video_grid_thw: None | np.ndarray = None
   video_second_per_grid: None | np.ndarray = None
+  video_mask: None | np.ndarray = None
   # Audio attributes.
   num_audios: int = 0
   audio_values: None | np.ndarray = None
   audio_mask: None | np.ndarray = None
   audio_lengths: None | np.ndarray = None
 
 
+def maybe_pad_video_values_to_max_grid(
+    video_values: np.ndarray,
+    video_grid_thw: np.ndarray,
+    config,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
+  """Pad Qwen3-Omni video pixels to configured static grid limits when enabled.
+
+  Args:
+    video_values: Video pixels of shape (batch, channels, T*tps, H*patch, W*patch).
+    video_grid_thw: Actual video grid with shape (1, 3), in Qwen grid units.
+    config: Config carrying video_max_grid_t/h/w and ViT patch sizes.
+
+  Returns:
+    Tuple of:
+    - padded video pixels, or the input when no max grid is configured
+    - input grid_thw
+    - pixel-level mask of shape (batch, 1, max_T*tps, max_H*patch, max_W*patch), or None
+  """
+  max_grid = (
+      getattr(config, "video_max_grid_t", None),
+      getattr(config, "video_max_grid_h", None),
+      getattr(config, "video_max_grid_w", None),
+  )
+  if all(dim is None for dim in max_grid):
+    return video_values, video_grid_thw, None
+  if any(dim is None for dim in max_grid):
+    raise ValueError("video_max_grid_t, video_max_grid_h, and video_max_grid_w must be set together.")
+  if video_values.ndim != 5:
+    raise ValueError(f"video_values must have shape (batch, channels, time, height, width), got {video_values.shape}.")
+
+  max_t, max_h, max_w = (int(dim) for dim in max_grid)
+  actual_t, actual_h, actual_w = (int(dim) for dim in video_grid_thw[0])
+  if actual_t > max_t or actual_h > max_h or actual_w > max_w:
+    raise ValueError(
+        f"video grid {video_grid_thw[0].tolist()} exceeds max grid {(max_t, max_h, max_w)}. "
+        "Scale or resize the video before padding."
+    )
+
+  temporal_patch_size = config.temporal_patch_size_for_vit
+  patch_size = config.patch_size_for_vit
+  valid_t_px = actual_t * temporal_patch_size
+  valid_h_px = actual_h * patch_size
+  valid_w_px = actual_w * patch_size
+  max_t_px = max_t * temporal_patch_size
+  max_h_px = max_h * patch_size
+  max_w_px = max_w * patch_size
+
+  padded_video_values = np.zeros(
+      (video_values.shape[0], video_values.shape[1], max_t_px, max_h_px, max_w_px),
+      dtype=video_values.dtype,
+  )
+  padded_video_values[:, :, :valid_t_px, :valid_h_px, :valid_w_px] = video_values[
+      :, :, :valid_t_px, :valid_h_px, :valid_w_px
+  ]
+
+  video_mask = np.zeros((video_values.shape[0], 1, max_t_px, max_h_px, max_w_px), dtype=np.int32)
+  video_mask[:, :, :valid_t_px, :valid_h_px, :valid_w_px] = 1
+
+  return padded_video_values, video_grid_thw, video_mask
+
+
 def smart_resize(
     height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280
 ):
@@ -589,8 +651,10 @@ def preprocess_mm_data_qwen3_omni(config):
             config.patch_size_for_vit * video_grid_thw[0, 2],
         ),
     )
+    video_values, video_grid_thw, video_mask = maybe_pad_video_values_to_max_grid(video_values, video_grid_thw, config)
     processor_outputs.video_values = video_values
     processor_outputs.video_grid_thw = video_grid_thw
+    processor_outputs.video_mask = video_mask
     processor_outputs.video_second_per_grid = np.asarray([config.temporal_patch_size_for_vit], dtype=np.float32)
     processor_outputs.num_videos = 1  # Only one video for now.
 

tests/unit/qwen3_omni_layers_test.py

@@ -37,6 +37,7 @@
 )
 from maxtext.layers.decoders import deepstack_process
 from maxtext.layers.encoders import AudioEncoder
+from maxtext.multimodal.processor_qwen3_omni import maybe_pad_video_values_to_max_grid
 from maxtext.models.qwen3 import (
     Qwen3OmniAudioEncoder,
     Qwen3OmniAudioEncoderLayer,
@@ -360,13 +361,80 @@ def test_patch_embed_output_matches_torch(self):
     )
 
     torch_output = self.torch_model(torch_hidden_states)
-    jax_output = self.jax_model(jax_hidden_states)
+    jax_output, attention_mask = self.jax_model(jax_hidden_states)
+    self.assertIsNone(attention_mask)
 
     torch_output_squeezed = torch_output.squeeze(1)
     jax_output_squeezed = jax_output.squeeze(1)
 
     assert_all_close_jax_torch(jax_output_squeezed, torch_output_squeezed, rtol=1e-3, atol=5e-3)
 
+  def test_patch_embed_padded_video_valid_outputs_match_unpadded(self):
+    """Test that valid padded video patches match embeddings from unpadded input."""
+    config_video_pad = pyconfig.initialize(
+        ["", base_config_path],
+        model_name="qwen3-omni-30b-a3b",
+        attention="dot_product",
+        attention_type="full",
+        matmul_precision="highest",
+        dropout_rate=0.0,
+        dtype="float32",
+        dtype_mm="float32",
+        weight_dtype="float32",
+        float32_logits=True,
+        float32_qk_product=True,
+        video_max_grid_t=3,
+        video_max_grid_h=4,
+        video_max_grid_w=4,
+    )
+    batch_size = 1
+    raw_grid = (2, 2, 3)  # Padding to (3, 4, 4)
+    max_grid = (
+        config_video_pad.video_max_grid_t,
+        config_video_pad.video_max_grid_h,
+        config_video_pad.video_max_grid_w,
+    )
+    temporal_patch_size = self.config.temporal_patch_size_for_vit
+    patch_size = self.config.patch_size_for_vit
+    in_channels = self.config.num_channels_for_vit
+
+    raw_shape = (
+        batch_size,
+        in_channels,
+        raw_grid[0] * temporal_patch_size,
+        raw_grid[1] * patch_size,
+        raw_grid[2] * patch_size,
+    )
+    raw_hidden_states, _ = create_random_jax_torch(*raw_shape)
+    raw_grid_thw = np.asarray([raw_grid], dtype=np.int32)
+
+    padded_hidden_states, padded_grid_thw, video_mask = maybe_pad_video_values_to_max_grid(
+        np.asarray(raw_hidden_states),
+        raw_grid_thw,
+        config_video_pad,
+    )
+
+    raw_output, raw_attention_mask = self.jax_model(raw_hidden_states)
+    padded_output, attention_mask = self.jax_model(
+        jnp.asarray(padded_hidden_states),
+        video_mask=jnp.asarray(video_mask),
+    )
+
+    self.assertIsNone(raw_attention_mask)
+    np.testing.assert_array_equal(padded_grid_thw, raw_grid_thw)
+    self.assertEqual(raw_output.shape, (batch_size, math.prod(raw_grid), self.config.hidden_size_for_vit))
+    self.assertEqual(padded_output.shape, (batch_size, math.prod(max_grid), self.config.hidden_size_for_vit))
+    self.assertEqual(attention_mask.shape, (batch_size, math.prod(max_grid)))
+    self.assertEqual(int(jnp.sum(attention_mask)), math.prod(raw_grid))
+
+    padded_valid_output = np.array(padded_output)[np.array(attention_mask, dtype=bool)]
+    np.testing.assert_allclose(
+        np.array(raw_output).reshape(-1, self.config.hidden_size_for_vit),
+        padded_valid_output,
+        rtol=1e-3,
+        atol=5e-3,
+    )
+
   def test_patch_embed_is_jittable(self):
     """Test that patch embed is JIT-compilable."""