fix: add defensive fallback for Qwen VL mrope get_rope_index shape mismatch

[Ottohere-Mourn]

Description

When training Qwen VL models with mixed multimodal data, get_rope_index can raise:

RuntimeError: "value tensor of shape [3, A] cannot be broadcast to indexing result of shape [3, B]"

This is not a frame-count issue — it is caused by missing media files triggering inconsistent fallback paths.

Root cause chain

1. Some samples reference npy/video files that are missing from the dataset directory
2. The custom video loader catches the load error and produces a fallback volume (e.g. a fixed-shape all-black array)
3. Tokenize stage processes this fallback through custom grid metadata logic, producing a video_grid_thw with specific placeholder token count (e.g. 462 or 504)
4. Forward stage re-loads the same missing file, also fallback, but computes video_grid_thw via transformers' native video_processor — a different code path → token count diverges
5. Placeholder count in input_ids ≠ actual visual feature count → shape mismatch in get_rope_index

In a 9.16M-sample run, only 19 unique IDs triggered this (105 rows, 0.0011%), all referencing missing files in M3D_CT/MedMNIST.

Fix

1. try/except around get_rope_func (qwen vl branch), catching only the known shape mismatch
2. _fallback_rope_position_ids: flat 3-axis positional encoding, fully vectorized
3. _log_rope_mismatch: one-line diagnostic for downstream data cleaning

Validation

• 9.16M-sample Qwen3.5-27B run: 0 false positives, step 279+ stable
• 105 affected samples identified via diagnostic log and cleaned at data level

[gemini-code-assist]

Code Review

This pull request introduces a defensive fallback mechanism in the data collator for Qwen VL models to handle mrope shape mismatch errors, preventing training runs from crashing. The review feedback highlights a performance bottleneck in the fallback method due to GPU-CPU synchronization and loop overhead, and provides a vectorized PyTorch implementation to resolve it.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Performance Bottleneck: GPU-CPU Synchronization & Loop Overhead

In the current implementation of _fallback_rope_position_ids, there are two major performance issues that can slow down training throughput:

1. GPU-CPU Synchronization: Calling .item() on a GPU tensor (inside attention_mask[batch_idx].bool().sum().item()) forces the CPU to block and wait for the GPU to finish its computation. Since this is executed inside a loop for every batch in the data collator, it can severely bottleneck the training pipeline and lower GPU utilization.
2. Redundant Tensor Creation & Loop Overhead: Re-creating the positions tensor bsz times inside a Python loop is inefficient.

Solution

We can completely vectorize this method to run loop-free and avoid any GPU-CPU synchronization by leveraging PyTorch's native tensor operations (expand, torch.where, and sum along dimensions).

    @staticmethod
    def _fallback_rope_position_ids(input_ids: "torch.Tensor", attention_mask: "torch.Tensor"):
        r"""Flat positional fallback: every non-pad token gets a monotonic position; all 3 mrope axes identical."""
        bsz, seq_len = input_ids.shape
        positions = torch.arange(seq_len, device=input_ids.device, dtype=input_ids.dtype)
        position_ids = positions.view(1, 1, seq_len).expand(3, bsz, seq_len).contiguous()
        valid_len = attention_mask.bool().sum(dim=-1)
        rope_deltas = torch.where(valid_len > 0, seq_len - valid_len, torch.zeros_like(valid_len))
        return position_ids, rope_deltas.to(dtype=input_ids.dtype)

[Ottohere-Mourn]

Thanks, applied the vectorized version in 2e67cbb. The Python loop and .item() call have been removed.