diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py index b4328a47d5..05bc1fa193 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -1401,44 +1401,120 @@ def __call__( height: int, width: int, token_mask: Array | None = None, - valid_grid: tuple[int, int, int] | None = None, + valid_grid: Array | None = None, ) -> Array: """Apply rotary position embeddings directly to inputs (Q or K tensors). Args: - inputs: Input tensor of shape [B, T*H*W, N, head_dim] (batch, sequence, heads, head_dim) - where T=num_frames, H=height, W=width (all static) + inputs: Input tensor of shape [B, S, N, head_dim] (batch, sequence, heads, head_dim) + where sequence length S = num_frames * height * width (static maximum). num_frames: Number of temporal frames (static) height: Height in patches (static) width: Width in patches (static) token_mask: Optional mask identifying valid tokens in the padded sequence. - valid_grid: Optional unpadded `(frames, height, width)` grid used to compute - RoPE for valid tokens. + valid_grid: Optional actual video grid with shape (batch, 3) (dynamic) Returns: - Rotated inputs with same shape [B, T*H*W, N, head_dim] + Rotated inputs with same shape [B, S, N, head_dim] """ - cos_emb, sin_emb = self.compute_cos_sin(num_frames, height, width) - if token_mask is not None and valid_grid is not None: - valid_cos, valid_sin = self.compute_cos_sin(*valid_grid) - valid_len = math.prod(valid_grid) - valid_indices = jnp.nonzero(token_mask[0], size=valid_len)[0] - cos_emb = jnp.ones_like(cos_emb).at[valid_indices].set(valid_cos) - sin_emb = jnp.zeros_like(sin_emb).at[valid_indices].set(valid_sin) - - if len(inputs.shape) == 4: - cos_emb = cos_emb[None, :, None, :] # [1, S, 1, H] - sin_emb = sin_emb[None, :, None, :] - elif len(inputs.shape) == 3: - # For [S, N, H] case - cos_emb = cos_emb[:, None, :] # [S, 1, H] - sin_emb = sin_emb[:, None, :] + is_3d = inputs.ndim == 3 + if is_3d: + inputs = inputs[None, :, :, :] + + batch_size = inputs.shape[0] + max_patches = num_frames * height * width + + if valid_grid is not None: + if isinstance(valid_grid, (tuple, list)): + valid_grid = jnp.array([valid_grid], dtype=jnp.int32) + elif valid_grid.ndim == 1: + valid_grid = valid_grid[None, :] + else: + valid_grid = jnp.tile( + jnp.array([[num_frames, height, width]], dtype=jnp.int32), + (batch_size, 1), + ) + + # Generate coordinates in block-based order + row, col = generate_block_coords(valid_grid, max_patches, self.spatial_merge_size) + + # Compute frequencies dynamically using coordinates + max_hw = max(height, width) + freq_table = self._compute_freq_table(max_hw) # [max_hw, head_dim//4] + + row_freqs = freq_table[row] # [B, max_patches, head_dim//4] + col_freqs = freq_table[col] # [B, max_patches, head_dim//4] + + # Concatenate row and column frequencies + embeddings = jnp.concatenate([row_freqs, col_freqs], axis=-1) # [B, max_patches, head_dim//2] + + # Double the embeddings to match head_dim + embeddings = jnp.concatenate([embeddings, embeddings], axis=-1) # [B, max_patches, head_dim] + + cos_emb = jnp.cos(embeddings) + sin_emb = jnp.sin(embeddings) + + # Mask out invalid (padded) tokens by setting cos to 1 and sin to 0 + if token_mask is not None: + is_valid = token_mask[:, :, None] # [B, max_patches, 1] + cos_emb = jnp.where(is_valid, cos_emb, 1.0) + sin_emb = jnp.where(is_valid, sin_emb, 0.0) + + if self.cast_as_fprop_dtype: + cos_emb = cos_emb.astype(self.fprop_dtype) + sin_emb = sin_emb.astype(self.fprop_dtype) + + # Reshape for broadcasting to inputs [B, S, N, H] + cos_emb = cos_emb[:, :, None, :] # [B, S, 1, H] + sin_emb = sin_emb[:, :, None, :] rotated = inputs * cos_emb + self._rotate_half(inputs) * sin_emb + if is_3d: + rotated = rotated[0] + return rotated +def generate_block_coords(video_grid_thw: Array, max_patches: int, merge_size: int) -> tuple[Array, Array]: + """Generate row and col coordinates in block-based order for padded video. + + Args: + video_grid_thw: Actual video grid with shape (batch, 3), in Qwen grid units. + max_patches: Maximum number of patches (static). + merge_size: Spatial merge block size (static). + + Returns: + Tuple of (row, col) each of shape [batch, max_patches] + """ + V_T = video_grid_thw[:, 0:1] # [B, 1] + V_H = video_grid_thw[:, 1:2] + V_W = video_grid_thw[:, 2:3] + + merged_w = V_W // merge_size + V_len = V_T * V_H * V_W + + idx = jnp.arange(max_patches, dtype=jnp.int32)[None, :] # [1, max_patches] + is_valid = idx < V_len # [B, max_patches] + + stride_lh = V_H * V_W + safe_stride_lh = jnp.maximum(stride_lh, 1) + s_idx = idx % safe_stride_lh + + intra_col = s_idx % merge_size + intra_row = (s_idx // merge_size) % merge_size + + safe_merged_w = jnp.maximum(merged_w, 1) + block_elements = merge_size * merge_size + block_col = (s_idx // block_elements) % safe_merged_w + block_row = s_idx // (safe_merged_w * block_elements) + + row = jnp.where(is_valid, block_row * merge_size + intra_row, 0) + col = jnp.where(is_valid, block_col * merge_size + intra_col, 0) + + return row, col + + def qwen3omnimoe_vision_pos_embed_interpolate_as_linen( *, num_position_embeddings: int, @@ -1596,46 +1672,87 @@ def _interpolate_single(self, t: int, h: int, w: int) -> tuple[Array, Array]: return indices, weights - def __call__(self, num_frames: int, height: int, width: int) -> Array: + def __call__( + self, + num_frames: int, + height: int, + width: int, + video_grid_thw: Array | None = None, + attention_mask: Array | None = None, + ) -> Array: """Interpolate positional embeddings for given static grid dimensions. Args: num_frames: Number of temporal frames (static) height: Height in patches (static) width: Width in patches (static) + video_grid_thw: Optional actual video grid with shape (batch, 3) (dynamic) + attention_mask: Optional attention mask with shape (batch, max_patches) (dynamic) Returns: - Interpolated positional embeddings of shape [num_frames * height * width, hidden_size] + Interpolated positional embeddings of shape [batch, num_frames * height * width, hidden_size] """ - # Get interpolation indices and weights - indices, weights = self._interpolate_single(num_frames, height, width) # [4, h*w], [4, h*w] + if video_grid_thw is not None: + if isinstance(video_grid_thw, (tuple, list)): + video_grid_thw = jnp.array([video_grid_thw], dtype=jnp.int32) + elif video_grid_thw.ndim == 1: + video_grid_thw = video_grid_thw[None, :] + batch_size = video_grid_thw.shape[0] + elif attention_mask is not None: + batch_size = attention_mask.shape[0] + else: + batch_size = 1 - # Lookup embeddings for all 4 corners - corner_embeds = self.pos_embed.value[indices] # [4, h*w, hidden_size] + max_patches = num_frames * height * width - # Apply bilinear weights and sum - weighted_embeds = corner_embeds * weights[:, :, None] # [4, h*w, hidden_size] - interpolated = jnp.sum(weighted_embeds, axis=0) # [h*w, hidden_size] + if video_grid_thw is None: + video_grid_thw = jnp.tile( + jnp.array([[num_frames, height, width]], dtype=jnp.int32), + (batch_size, 1), + ) + if attention_mask is None: + attention_mask = jnp.ones((batch_size, max_patches), dtype=jnp.int32) - # Repeat for temporal frames - if num_frames > 1: - interpolated = jnp.tile(interpolated, (num_frames, 1)) # [t*h*w, hidden_size] + # Generate coordinates in block-based order + row, col = generate_block_coords(video_grid_thw, max_patches, self.spatial_merge_size) - # Apply spatial merge permutation - # Reshape to [t, h, w, hidden_size] then permute for block-based processing - merge_size = self.spatial_merge_size - merged_h = height // merge_size - merged_w = width // merge_size + # Normalize coordinates to [0, 1] range based on valid grid sizes + V_H = video_grid_thw[:, 1:2] + V_W = video_grid_thw[:, 2:3] + row_norm = row / jnp.maximum(V_H - 1, 1) + col_norm = col / jnp.maximum(V_W - 1, 1) + + # Interpolate from N x N grid + N = self.num_grid_per_side + table = self.pos_embed.value.reshape(N, N, self.hidden_size) - # Reshape: [t*h*w, hidden_size] -> [t, h, w, hidden_size] - interpolated = interpolated.reshape(num_frames, height, width, self.hidden_size) + y = row_norm * (N - 1) + x = col_norm * (N - 1) + + y0 = jnp.floor(y).astype(jnp.int32) + x0 = jnp.floor(x).astype(jnp.int32) + y1 = jnp.minimum(y0 + 1, N - 1) + x1 = jnp.minimum(x0 + 1, N - 1) + + dy = y - y0 + dx = x - x0 + + # Gather 4 corners + embed_00 = table[y0, x0] + embed_01 = table[y0, x1] + embed_10 = table[y1, x0] + embed_11 = table[y1, x1] + + # Apply bilinear weights + dy = dy[:, :, None] + dx = dx[:, :, None] + + interpolated = ( + (1.0 - dy) * (1.0 - dx) * embed_00 + (1.0 - dy) * dx * embed_01 + dy * (1.0 - dx) * embed_10 + dy * dx * embed_11 + ) - # Permute for spatial merging: [t, merged_h, merge_size, merged_w, merge_size, hidden_size] - interpolated = interpolated.reshape(num_frames, merged_h, merge_size, merged_w, merge_size, self.hidden_size) - # -> [t, merged_h, merged_w, merge_size, merge_size, hidden_size] - interpolated = jnp.transpose(interpolated, (0, 1, 3, 2, 4, 5)) - # Flatten back to [t*merged_h*merged_w*merge_size*merge_size, hidden_size] - interpolated = interpolated.reshape(-1, self.hidden_size) + # Mask out invalid (padded) tokens by setting their embeddings to zero + interpolated = interpolated * attention_mask[:, :, None] if self.cast_as_fprop_dtype: interpolated = interpolated.astype(self.fprop_dtype) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index e1685f72e4..6920565cdb 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -2136,22 +2136,17 @@ def __call__( x, _ = self.patch_embed(hidden_states) x = x.reshape(batch_size, -1, self.config.hidden_size_for_vit) - valid_grid = None if attention_mask is not None and video_grid_thw is None: raise ValueError("video_grid_thw is required when video_mask is provided.") - if attention_mask is not None and batch_size != 1: - raise ValueError("Padded Qwen3-Omni vision encoding currently supports batch size one.") - if video_grid_thw is not None: - grid = video_grid_thw[0] if getattr(video_grid_thw, "ndim", 1) == 2 else video_grid_thw - valid_grid = tuple(int(dim) for dim in grid) - pos = self.pos_embed_interpolate(num_frames, height, width) - if attention_mask is not None and valid_grid is not None: - valid_pos = self.pos_embed_interpolate(*valid_grid) - valid_indices = jnp.nonzero(attention_mask[0], size=math.prod(valid_grid))[0] - pos = jnp.zeros_like(pos).at[valid_indices].set(valid_pos) - - pos = pos[jnp.newaxis, :, :] + pos = self.pos_embed_interpolate( + num_frames, + height, + width, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + ) x = x + pos + valid_grid = video_grid_thw h_traj = [] for i in range(self.depth): diff --git a/tests/unit/qwen3_omni_layers_test.py b/tests/unit/qwen3_omni_layers_test.py index 04dd2d191f..753ba30251 100644 --- a/tests/unit/qwen3_omni_layers_test.py +++ b/tests/unit/qwen3_omni_layers_test.py @@ -37,7 +37,7 @@ Qwen3OmniMoeVisionRotaryEmbedding as JaxQwen3OmniMoeVisionRotaryEmbedding, ) from maxtext.layers.decoders import deepstack_process -from maxtext.layers.encoders import AudioEncoder +from maxtext.layers.encoders import AudioEncoder, VisionEncoder from maxtext.multimodal.processor_qwen3_omni import ( maybe_pad_video_values_to_max_grid, preprocess_video, @@ -53,6 +53,7 @@ Qwen3OmniMoeVisionProjector as JaxQwen3OmniMoeVisionProjector, ) from maxtext.multimodal import processor as mm_processor +from maxtext.multimodal import utils as mm_utils from tests.utils.multimodal_test_utils import ( assert_all_close_jax_torch, copy_attention_weights_to_maxtext, @@ -645,6 +646,9 @@ def _copy_weights_and_test(self, num_frames, height, width): pos_embed_jax = self.jax_model(num_frames, height, width) pos_embed_torch = self.torch_encoder.fast_pos_embed_interpolate(grid_thw_torch) + if pos_embed_jax.ndim == 3 and pos_embed_torch.ndim == 2: + pos_embed_jax = pos_embed_jax[0] + assert_all_close_jax_torch(pos_embed_jax, pos_embed_torch, rtol=1e-2, atol=1e-2) def test_pos_embed_interpolate_matches_torch(self): @@ -762,7 +766,7 @@ def test_padded_video_valid_outputs_match_unpadded(self): encoder = JaxQwen3OmniMoeVisionEncoder(config=config, mesh=self.mesh, rngs=nnx.Rngs(42)) raw_output, _ = encoder(raw_video) - padded_output, _ = encoder(padded_video, video_mask=video_mask, video_grid_thw=raw_grid) + padded_output, _ = encoder(padded_video, video_mask=video_mask, video_grid_thw=jnp.array([raw_grid], dtype=jnp.int32)) patch_mask = np.asarray(video_mask).reshape(1, -1, temporal_patch_size * patch_size * patch_size).max(-1).astype(bool) np.testing.assert_allclose( @@ -772,6 +776,161 @@ def test_padded_video_valid_outputs_match_unpadded(self): atol=1e-4, ) + def test_padded_video_projected_outputs_match_unpadded_batched(self): + """Padded tokens do not change valid projected outputs for batched inputs.""" + raw_grid_1 = (2, 2, 2) + raw_grid_2 = (1, 4, 4) + max_grid = (3, 4, 4) + config = pyconfig.initialize( + ["", base_config_path], + model_name="qwen3-omni-30b-a3b", + attention="dot_product", + attention_type="full", + dtype="float32", + dtype_mm="float32", + weight_dtype="float32", + override_model_config=True, + attention_for_vit="dot_product", + hidden_size_for_vit=16, + num_attention_heads_for_vit=2, + intermediate_size_for_vit=32, + num_hidden_layers_for_vit=1, + deepstack_visual_indexes_for_vit=[], + video_max_grid_t=max_grid[0], + video_max_grid_h=max_grid[1], + video_max_grid_w=max_grid[2], + out_hidden_size_for_vit=24, + vision_encoder_block=common_types.VisionEncoderBlockType.QWEN3_OMNI, + ) + patch_size = config.patch_size_for_vit + temporal_patch_size = config.temporal_patch_size_for_vit + + # Generate raw videos + raw_shape_1 = ( + 1, + config.num_channels_for_vit, + raw_grid_1[0] * temporal_patch_size, + raw_grid_1[1] * patch_size, + raw_grid_1[2] * patch_size, + ) + raw_shape_2 = ( + 1, + config.num_channels_for_vit, + raw_grid_2[0] * temporal_patch_size, + raw_grid_2[1] * patch_size, + raw_grid_2[2] * patch_size, + ) + raw_video_1, _ = create_random_jax_torch(*raw_shape_1) + raw_video_2, _ = create_random_jax_torch(*raw_shape_2) + + # Pad them individually + padded_video_1, _, video_mask_1 = maybe_pad_video_values_to_max_grid( + np.asarray(raw_video_1), np.asarray([raw_grid_1]), config + ) + padded_video_2, _, video_mask_2 = maybe_pad_video_values_to_max_grid( + np.asarray(raw_video_2), np.asarray([raw_grid_2]), config + ) + + # Batch them + batched_padded_video = jnp.concatenate([padded_video_1, padded_video_2], axis=0) + batched_video_mask = jnp.concatenate([video_mask_1, video_mask_2], axis=0) + batched_grid_thw = jnp.array([raw_grid_1, raw_grid_2], dtype=jnp.int32) + + # Instantiate VisionEncoder (wraps encoder + projector) + vision_encoder = VisionEncoder(config=config, mesh=self.mesh, rngs=nnx.Rngs(42)) + + # Run individual unpadded + raw_output_1, _ = vision_encoder(raw_video_1) + raw_output_2, _ = vision_encoder(raw_video_2) + + # Run batched padded + batched_padded_output, _ = vision_encoder( + batched_padded_video, + input_masks=batched_video_mask, + video_grid_thw=batched_grid_thw, + ) + + # Downsample masks to projected tokens + projected_mask_1 = mm_processor.downsample_video_mask_to_tokens(video_mask_1, config) + projected_mask_2 = mm_processor.downsample_video_mask_to_tokens(video_mask_2, config) + + # Verify equivalence for batch item 0 + np.testing.assert_allclose( + np.asarray(raw_output_1).reshape(-1, config.out_hidden_size_for_vit), + np.asarray(batched_padded_output[0])[np.asarray(projected_mask_1[0]).astype(bool)], + rtol=1e-4, + atol=1e-4, + ) + + # Verify equivalence for batch item 1 + np.testing.assert_allclose( + np.asarray(raw_output_2).reshape(-1, config.out_hidden_size_for_vit), + np.asarray(batched_padded_output[1])[np.asarray(projected_mask_2[0]).astype(bool)], + rtol=1e-4, + atol=1e-4, + ) + + # --- Test Merging into Text Sequence --- + emb_dim = config.out_hidden_size_for_vit + + # Example 1: 2 placeholders at indices 5, 6. + text_embeddings_1 = jnp.arange(1 * 20 * emb_dim, dtype=jnp.float32).reshape(1, 20, emb_dim) + placeholder_mask_1 = jnp.zeros((1, 20), dtype=jnp.int32) + placeholder_mask_1 = placeholder_mask_1.at[:, 5:7].set(1) + + # Example 2: 4 placeholders at indices 5, 6, 7, 8. + text_embeddings_2 = jnp.arange(1 * 20 * emb_dim, dtype=jnp.float32).reshape(1, 20, emb_dim) + 1000.0 + placeholder_mask_2 = jnp.zeros((1, 20), dtype=jnp.int32) + placeholder_mask_2 = placeholder_mask_2.at[:, 5:9].set(1) + + # Batch text inputs + batched_text_embeddings = jnp.concatenate([text_embeddings_1, text_embeddings_2], axis=0) + batched_placeholder_mask = jnp.concatenate([placeholder_mask_1, placeholder_mask_2], axis=0) + batched_projected_mask = jnp.concatenate([projected_mask_1, projected_mask_2], axis=0) + + # Run batched merge + batched_merged = mm_utils.merge_mm_embeddings( + text_embeddings=batched_text_embeddings, + multimodal_embeddings=batched_padded_output, + mask=batched_placeholder_mask, + token_masks=batched_projected_mask, + ) + + # Run individual unpadded merges + token_mask_1 = jnp.ones((1, 2), dtype=jnp.int32) + merged_1 = mm_utils.merge_mm_embeddings( + text_embeddings=text_embeddings_1, + multimodal_embeddings=raw_output_1, + mask=placeholder_mask_1, + token_masks=token_mask_1, + ) + + # Example 2 has 4 valid blocks + token_mask_2 = jnp.ones((1, 4), dtype=jnp.int32) + merged_2 = mm_utils.merge_mm_embeddings( + text_embeddings=text_embeddings_2, + multimodal_embeddings=raw_output_2, + mask=placeholder_mask_2, + token_masks=token_mask_2, + ) + + # Verify that batched merge matches individual unpadded merges + # Batch item 0 + np.testing.assert_allclose( + np.asarray(batched_merged[0]), + np.asarray(merged_1[0]), + rtol=1e-4, + atol=1e-4, + ) + + # Batch item 1 + np.testing.assert_allclose( + np.asarray(batched_merged[1]), + np.asarray(merged_2[0]), + rtol=1e-4, + atol=1e-4, + ) + class TestDeepstackProcess(unittest.TestCase): """Tests for deepstack_process.