feat: Fix Qwen3-VL THD packed mRoPE positions

[Zhichenzzz]

Problem

Under THD sequence packing, several variable-length sequences are concatenated into one
row input_ids = [1, total], with packed_seq_params.cu_seqlens_q marking the segment
boundaries. Megatron-Bridge's Qwen3VLModel.forward resets position_ids = None and then
recomputes MRoPE positions via the module-level get_rope_index over the whole packed
row, so positions run monotonically across segment boundaries instead of restarting per
sequence.

This is harmless for text (RoPE is relative — a constant per-segment offset cancels), but
wrong for multimodal: image/video (t, h, w) grid positions do not commute with a
whole-row offset, so the positions become bogus for packed VL training.

Why not just inject position_ids

Passing position_ids into forward from the outside does not work: forward discards it
at the position_ids = None reset. Also, the current Bridge Qwen3VLModel has no
get_rope_index method (it calls the module-level free function), so wiring that assumes a
get_rope_index method does not attach.

Fix (miles-side, no Megatron-Bridge edit)

Hijack the internal call instead: wrap Qwen3VLModel.forward to build correct per-segment
positions (slice the [1, total] row by cu_seqlens_q; text-only segment → linear 0..L,
media segment → per-segment get_rope_index with the matching grid slice; concat to
[3, 1, total]) and stash them on a thread-local; patch the module-level get_rope_index
to return the stash. The reset is then harmless — the following get_rope_index call yields
our positions, which flow through preprocess_packed_seqs (THD/CP) unchanged. Only the
position values change; tensor shapes and the downstream flow are identical.

Guarded to THD single-row packed input; otherwise Bridge runs normally (e.g. CP-sharded or
[batch, seq] falls back). Replaces the prior 1-line rotary-only monkey-patch (the rotary
signature shim is kept).

Validation

• Per-segment position check: packed multimodal positions equal running each segment
  unpacked exactly (e.g. seg B t-channel [0,1,2,...] vs the bug's [43,44,45,...]).
• End-to-end Qwen3-VL RL (geo3k, Qwen3-VL-2B-Instruct, megatron + THD, 8×H200): job
  succeeds, 8 training steps with a weight update each, gradients flow, and the
  train↔rollout logprob abs-diff stays small/stable (~0.011–0.016) throughout.

[gemini-code-assist]

Code Review

This pull request introduces a new module qwen_vl_packed_mrope.py to fix Qwen3-VL THD packed mRoPE positions by patching the rotary embedding signature and model forward pass. Feedback on the changes highlights two main improvement opportunities: first, vectorizing the media token counting across segments to avoid performance-degrading GPU-to-host synchronizations inside a loop; second, removing a redundant list comprehension when converting an integer tensor to a list.

[gemini-code-assist]

The current implementation of _count_media performs multiple GPU-to-host synchronizations (.item()) inside a loop over all segments. In a training run with many segments, this will cause significant performance degradation due to device-to-host latency.

We can completely vectorize the media token counting across all segments using torch.bucketize and torch.bincount on the GPU, and then copy the counts to the host in a single step. This reduces the number of synchronizations to a constant (independent of the number of segments) and avoids calling _count_media in a loop.

    # Vectorized counting of media tokens per segment to avoid per-segment GPU-CPU syncs.
    # Using right=True ensures left-closed, right-open intervals [start, end) map correctly.
    num_segments = len(cu) - 1
    img_counts = [0] * num_segments
    vid_counts = [0] * num_segments

    starts = torch.nonzero(flat == vstart, as_tuple=False).flatten()
    starts = starts[starts + 1 < flat.numel()]
    if starts.numel() > 0:
        toks = flat[starts + 1]
        is_img = (toks == img_id)
        is_vid = (toks == vid_id)
        
        segment_indices = torch.bucketize(starts, psp.cu_seqlens_q, right=True) - 1
        
        img_seg_indices = segment_indices[is_img]
        vid_seg_indices = segment_indices[is_vid]
        
        img_counts_tensor = torch.bincount(img_seg_indices, minlength=num_segments)
        vid_counts_tensor = torch.bincount(vid_seg_indices, minlength=num_segments)
        
        img_counts = img_counts_tensor.cpu().tolist()
        vid_counts = vid_counts_tensor.cpu().tolist()

    img_off = vid_off = 0
    segments = []
    for i, (start, end) in enumerate(zip(cu[:-1], cu[1:])):
        if end <= start:
            continue
        seg = flat[start:end]
        ic = img_counts[i]
        vc = vid_counts[i]
        if ic == 0 and vc == 0:
            pos = torch.arange(seg.numel(), dtype=seg.dtype, device=seg.device).view(1, 1, -1).expand(3, 1, -1)
        else:
            pos, _ = orig_get_rope_index(
                merge, img_id, vid_id, vstart, seg.unsqueeze(0),
                image_grid_thw=_slice(image_grid_thw, img_off, ic),
                video_grid_thw=_slice(video_grid_thw, vid_off, vc),
                attention_mask=None,
            )
            pos = pos[:, :, : seg.numel()]
        img_off += ic
        vid_off += vc
        segments.append(pos)
    return torch.cat(segments, dim=2).contiguous() if segments else None

[gemini-code-assist]

Since cu is an integer tensor, calling .tolist() already returns a list of Python integers. The list comprehension [int(x) for x in ...] is redundant and adds unnecessary overhead.

Suggested change

	flat = input_ids.reshape(-1)
	cu = [int(x) for x in cu.detach().cpu().tolist()]
	if cu[0] != 0 or cu[-1] != flat.numel(): # e.g. CP-sharded: doesn't match local input
	return None
	flat = input_ids.reshape(-1)
	cu = cu.detach().cpu().tolist()
	if cu[0] != 0 or cu[-1] != flat.numel(): # e.g. CP-sharded: doesn't match local input
	return None

[TSunny007]

_build_packed_positions uses cu_seqlens_q and bails when cu[-1] != flat.numel(). With CP + THD packing the input_ids are laid out against cu_seqlens_q_padded (see bridge's thd_to_bshd at modelling_qwen3_vl/model.py:271-275), so this no-ops in that case and CP+packing falls back to the broken cross-sample positions.

Plan to handle that here, or as a follow-up?

[TSunny007]

Our rewards are looking good with this change on sequence packed examples

[Zhichenzzz]

_build_packed_positions uses cu_seqlens_q and bails when cu[-1] != flat.numel(). With CP + THD packing the input_ids are laid out against cu_seqlens_q_padded (see bridge's thd_to_bshd at modelling_qwen3_vl/model.py:271-275), so this no-ops in that case and CP+packing falls back to the broken cross-sample positions.

Plan to handle that here, or as a follow-up?

Thanks for raising the issue here. I will fix this in later VL PRs.

[yueming-yuan]

I would prefer to move this to /miles_plugins/models/qwen3-vl (or similar place), given we don't want to put patches in miles core logic..

[yueming-yuan]

approved (and better to clean up the patches in the future)