fix(qwen3-vl): per-segment mRoPE + vision under CP + THD packing

miles/backends/megatron_utils/model_provider.py

@@ -98,6 +98,8 @@ def wrapped_model_provider(
         provider.expert_tensor_parallel_size = args.expert_tensor_parallel_size
         provider.sequence_parallel = args.sequence_parallel
         provider.context_parallel_size = args.context_parallel_size
+        # CP>1 VL models assert this; bridge configs skip core_transformer_config_from_args.
+        provider.calculate_per_token_loss = args.calculate_per_token_loss
         provider.attention_softmax_in_fp32 = args.attention_softmax_in_fp32
         provider.variable_seq_lengths = args.variable_seq_lengths
         if hasattr(args, "moe_token_dispatcher_type"):

miles_plugins/models/qwen3_vl.py

@@ -11,6 +11,7 @@
 import importlib
 import logging
 import threading
+from typing import NamedTuple
 
 import torch
 
@@ -23,6 +24,31 @@
 def install_qwen3_vl_packed_mrope_patch() -> None:
     _patch_rotary_signature()
     _patch_model_forward_and_rope_index()
+    _patch_allgather_vision_embeddings_kwarg()
+
+
+def _patch_allgather_vision_embeddings_kwarg() -> None:
+    """megatron-bridge 0.5.0 calls AllGatherVisionEmbeddings.apply(..., cp_group=...) in the
+    Qwen3-VL vision_dp_when_cp path, but torch.autograd.Function.apply rejects keyword args
+    (TypeError: apply() takes no keyword arguments). Replace the symbol with a shim whose
+    .apply accepts cp_group as a kwarg and forwards it positionally.
+    """
+    try:
+        model_mod = importlib.import_module("megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model")
+    except ImportError:
+        return
+    orig = getattr(model_mod, "AllGatherVisionEmbeddings", None)
+    if orig is None or getattr(orig, "_miles_kwarg_shim", False):
+        return
+
+    class _AllGatherVisionEmbeddingsKwargShim:
+        _miles_kwarg_shim = True
+
+        @staticmethod
+        def apply(input, seqlens_on_cp_ranks, cp_group=None):
+            return orig.apply(input, seqlens_on_cp_ranks, cp_group)
+
+    model_mod.AllGatherVisionEmbeddings = _AllGatherVisionEmbeddingsKwargShim
 
 
 def _patch_rotary_signature() -> None:
@@ -65,30 +91,78 @@ def patched_get_rope_index(*args, **kwargs):
 
     model_mod.get_rope_index = patched_get_rope_index
 
+    # Under CP, miles pre-shards the THD row (slice_with_cp), but the bridge forward re-shards
+    # internally via preprocess_packed_seqs (it expects the FULL input). When miles already
+    # sharded, make that internal call an identity that returns miles' packed_seq_params (so CP
+    # attention still sees the full cu_seqlens) instead of re-splitting the already-local data.
+    _patch_preprocess_packed_seqs_identity(model_mod)
+
     Qwen3VLModel = getattr(model_mod, "Qwen3VLModel", None)
+    # The bridge selects CP-local vision embeds natively; warn when running an old one.
+    if Qwen3VLModel is not None and not hasattr(Qwen3VLModel, "_cp_local_vision_embed_indices"):
+        logger.warning(
+            "megatron-bridge Qwen3VLModel lacks native CP-local vision-embed selection; "
+            "CP runs with vision tokens will mis-place vision embeddings. "
+            "Apply the matching Megatron-Bridge patch (radixark/Megatron-Bridge PR #9)."
+        )
+
     if Qwen3VLModel is None or Qwen3VLModel.__dict__.get(_PATCHED, False):
         setattr(model_mod, _PATCHED, True)
         return
 
     orig_forward = Qwen3VLModel.forward
 
     def patched_forward(self, *args, **kwargs):
-        packed = _build_packed_positions(self, args, kwargs, orig_get_rope_index)
-        if packed is None:
-            return orig_forward(self, *args, **kwargs)
-        _tls.packed_positions = packed
+        parsed = _parse_packed_thd(args, kwargs)
+        packed = _build_packed_positions(self, parsed, kwargs, orig_get_rope_index)
+        ctx = _prepare_cp_local_context(parsed)
+        if packed is not None:
+            _tls.packed_positions = packed
+        if ctx is not None:
+            _tls.cp_local = ctx
         try:
             return orig_forward(self, *args, **kwargs)
         finally:
             _tls.packed_positions = None
+            _tls.cp_local = None
 
     Qwen3VLModel.forward = patched_forward
     setattr(Qwen3VLModel, _PATCHED, True)
     setattr(model_mod, _PATCHED, True)
 
 
-def _build_packed_positions(self, args, kwargs, orig_get_rope_index):
-    # Per-segment MRoPE positions for a THD single-row packed batch; else None (run normally).
+def _patch_preprocess_packed_seqs_identity(model_mod) -> None:
+    orig = getattr(model_mod, "preprocess_packed_seqs", None)
+    if orig is None or getattr(orig, "_miles_identity_wrapped", False):
+        return
+
+    def wrapped(input_ids, attention_mask, *args, **kwargs):
+        ctx = getattr(_tls, "cp_local", None)
+        if ctx is not None:
+            # Already CP-local: skip re-sharding and hand back miles' full-cu packed_seq_params.
+            return input_ids, ctx["psp"]
+        return orig(input_ids, attention_mask, *args, **kwargs)
+
+    wrapped._miles_identity_wrapped = True
+    model_mod.preprocess_packed_seqs = wrapped
+
+
+class _PackedTHD(NamedTuple):
+    """Parsed view of a THD single-row packed batch (see ``_parse_packed_thd``)."""
+
+    psp: object
+    flat: torch.Tensor
+    cu: list  # cu_seqlens_q as a host-side list (one GPU->host copy, shared by both paths)
+    cu_t: torch.Tensor
+    local_len: int
+    cp_size: int
+    cp_rank: int
+
+
+def _parse_packed_thd(args, kwargs):
+    """Extract the preamble shared by the mRoPE-position and CP-context paths, doing the one
+    GPU->host copy of cu_seqlens_q a single time. Returns None when this is not a THD packed
+    ``[1, T]`` batch, in which case both callers run the dense / unchanged path."""
     input_ids = kwargs.get("input_ids")
     if input_ids is None and args:
         input_ids = args[0]
@@ -101,21 +175,97 @@ def _build_packed_positions(self, args, kwargs, orig_get_rope_index):
     if cu_t is None or cu_t.numel() < 2:
         return None
     flat = input_ids.reshape(-1)
-    cu = cu_t.detach().cpu().tolist()
-    if cu[0] != 0 or cu[-1] != flat.numel():
-        # cu_seqlens_q doesn't describe this local row. Under context parallelism the THD
-        # row is laid out against cu_seqlens_q_padded (load-balanced CP chunks), which this
-        # path does not yet reconstruct, so we fall back to the dense get_rope_index.
-        # TODO(follow-up): per-segment positions for the CP + padded layout.
-        logger.debug(
-            "qwen3_vl packed mRoPE: cu_seqlens_q (%d) != local len (%d); using dense path", cu[-1], flat.numel()
-        )
+    cp_size, cp_rank = _cp_size_rank()
+    return _PackedTHD(psp, flat, cu_t.detach().cpu().tolist(), cu_t, flat.numel(), cp_size, cp_rank)
+
+
+def _prepare_cp_local_context(parsed):
+    """When CP has already pre-sharded this THD row, capture miles' packed_seq_params so the
+    preprocess_packed_seqs identity wrapper can hand them back unchanged (the bridge does
+    CP-local vision-embed selection natively). Returns None for the non-CP / full-input cases
+    (bridge runs unchanged)."""
+    if parsed is None or parsed.cp_size <= 1 or parsed.cu[0] != 0:
         return None
+    if parsed.cu[-1] != parsed.cp_size * parsed.local_len:
+        return None
+    return {"psp": parsed.psp}
+
+
+def _cp_size_rank():
+    """Context-parallel (size, rank); (1, 0) when CP is unavailable."""
+    try:
+        from megatron.core import parallel_state as _ps
+
+        return _ps.get_context_parallel_world_size(), _ps.get_context_parallel_rank()
+    except Exception:
+        return 1, 0
+
+
+def _natural_to_zigzag_slice(t, cp_size, cp_rank, dim):
+    """Slice a full-length tensor into this rank's zigzag (load-balanced ring-attn) CP chunks.
+
+    Mirrors miles.backends.training_utils.cp_utils.slice_with_cp / natural_to_zigzag_slice:
+    rank r owns chunks [r, 2*cp_size-1-r] of the 2*cp_size equal partitions along ``dim``.
+    """
+    total = t.shape[dim]
+    num_chunks = 2 * cp_size
+    chunk = total // num_chunks
+    idxs = [cp_rank, 2 * cp_size - 1 - cp_rank]
+    return torch.cat([t.narrow(dim, i * chunk, chunk) for i in idxs], dim=dim)
+
+
+def _cp_allgather_unzigzag(flat, cu, cp_size):
+    """Reconstruct the full THD packed row from this rank's zigzag chunks.
+
+    Under CP, ``flat`` holds only chunks [cp_rank, 2*cp-1-cp_rank] of every segment, and
+    ``cu`` (== psp.cu_seqlens_q) gives the FULL padded per-segment boundaries. All-gather the
+    per-rank rows over the CP group and de-interleave each segment back to natural order.
+    Returns None (caller falls back to dense) if a segment is not divisible by 2*cp.
+    """
+    import torch.distributed as dist
+    from megatron.core import parallel_state as _ps
+
+    group = _ps.get_context_parallel_group()
+    gathered = [torch.empty_like(flat) for _ in range(cp_size)]
+    dist.all_gather(gathered, flat.contiguous(), group=group)
+    return _reassemble_full_row(gathered, cu, cp_size)
 
+
+def _reassemble_full_row(gathered, cu, cp_size):
+    """De-interleave per-rank zigzag rows back into the full natural-order packed row.
+
+    ``gathered[r]`` is rank r's local row; ``cu`` (full padded per-segment boundaries, i.e.
+    miles' cu_seqlens * cp) locates each segment. For segment ``i`` of full length ``L`` (a
+    multiple of 2*cp), rank r contributed chunk r and chunk 2*cp-1-r, each of size L/(2*cp),
+    at local offset cu[i]//cp. Pure (no collectives) so it is unit-testable. Returns None if a
+    segment is not divisible by 2*cp (caller falls back to the dense path).
+    """
+    full = torch.zeros(cu[-1], dtype=gathered[0].dtype, device=gathered[0].device)
+    for i in range(len(cu) - 1):
+        seg_full = cu[i + 1] - cu[i]
+        if seg_full <= 0:
+            continue
+        if seg_full % (2 * cp_size) != 0:
+            return None
+        c = seg_full // (2 * cp_size)
+        local_off = cu[i] // cp_size  # this segment's offset within a per-rank (local) row
+        for r in range(cp_size):
+            mir = 2 * cp_size - 1 - r
+            full[cu[i] + r * c : cu[i] + (r + 1) * c] = gathered[r][local_off : local_off + c]
+            full[cu[i] + mir * c : cu[i] + (mir + 1) * c] = gathered[r][local_off + c : local_off + 2 * c]
+    return full
+
+
+def _segment_positions(model, flat, cu, cu_t, kwargs, orig_get_rope_index):
+    """Per-segment MRoPE positions for a full (unsharded) packed row `flat` with boundaries `cu`.
+
+    Returns a list of [3, 1, seg_len] tensors (one per non-empty segment), text segments get
+    a linear 0..L range, media segments call get_rope_index with the matching grid slice.
+    """
     image_grid_thw = kwargs.get("image_grid_thw")
     video_grid_thw = kwargs.get("video_grid_thw")
-    merge = self.config.spatial_merge_size
-    img_id, vid_id, vstart = self.image_token_id, self.video_token_id, self.vision_start_token_id
+    merge = model.config.spatial_merge_size
+    img_id, vid_id, vstart = model.image_token_id, model.video_token_id, model.vision_start_token_id
 
     # Vectorized media count per segment (one GPU->host copy total, no per-segment .item()).
     num_segments = len(cu) - 1
@@ -153,7 +303,45 @@ def _build_packed_positions(self, args, kwargs, orig_get_rope_index):
         img_off += ic
         vid_off += vc
         segments.append(pos)
-    return torch.cat(segments, dim=2).contiguous() if segments else None
+    return segments
+
+
+def _build_packed_positions(model, parsed, kwargs, orig_get_rope_index):
+    """Per-segment MRoPE positions for a THD packed batch; None to run the dense path."""
+    if parsed is None or parsed.cu[0] != 0:
+        return None
+    flat, cu, cu_t = parsed.flat, parsed.cu, parsed.cu_t
+    local_len, cp_size, cp_rank = parsed.local_len, parsed.cp_size, parsed.cp_rank
+
+    # Non-CP (or single chunk): cu_seqlens_q already describes this row exactly.
+    if cu[-1] == local_len:
+        segments = _segment_positions(model, flat, cu, cu_t, kwargs, orig_get_rope_index)
+        return torch.cat(segments, dim=2).contiguous() if segments else None
+
+    # CP + THD packing: cu_seqlens_q gives the FULL padded per-segment boundaries (miles
+    # builds cu_seqlens * cp_size), while this row holds only this rank's zigzag chunks
+    # (full_len / cp). Reconstruct the full row across the CP group, build full per-segment
+    # MRoPE positions, then re-slice each segment into this rank's zigzag layout so the
+    # positions line up with the tokens that slice_with_cp produced.
+    if cp_size > 1 and cu[-1] == cp_size * local_len:
+        full_flat = _cp_allgather_unzigzag(flat, cu, cp_size)
+        if full_flat is None:
+            logger.debug("qwen3_vl packed mRoPE: CP segment not divisible by 2*cp; dense path")
+            return None
+        segments = _segment_positions(model, full_flat, cu, cu_t, kwargs, orig_get_rope_index)
+        if not segments:
+            return None
+        local_segments = [_natural_to_zigzag_slice(p, cp_size, cp_rank, dim=2) for p in segments]
+        return torch.cat(local_segments, dim=2).contiguous()
+
+    # Unrecognized layout -> let the dense get_rope_index run.
+    logger.debug(
+        "qwen3_vl packed mRoPE: cu_seqlens_q (%d) vs local len (%d), cp=%d; dense path",
+        cu[-1],
+        local_len,
+        cp_size,
+    )
+    return None
 
 
 def _slice(grid, offset, count):

tests/fast/test_qwen3_vl_cp_mrope.py

@@ -0,0 +1,80 @@
+"""CPU unit test for the Qwen3-VL CP+THD packed mRoPE reconstruction (issue #1296).
+
+Under context parallelism each rank's THD row holds only its zigzag chunks of every packed
+segment. `_reassemble_full_row` de-interleaves the all-gathered per-rank rows back to the
+full natural-order row so per-segment MRoPE positions can be rebuilt and re-sliced. This
+test checks that reconstruction is the exact inverse of `slice_with_cp` (the function miles
+uses to shard the tokens), and that re-slicing with `_natural_to_zigzag_slice` round-trips.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from miles_plugins.models.qwen3_vl import _natural_to_zigzag_slice, _reassemble_full_row
+
+
+def _slice_with_cp(tokens, cp_size, cp_rank, pad_value=0):
+    """Reference copy of cp_utils.slice_with_cp's THD zigzag slicing (per sample)."""
+    token_len = len(tokens)
+    chunk = (token_len + 2 * cp_size - 1) // (2 * cp_size)
+    pad = 2 * cp_size * chunk - token_len
+    if pad:
+        tokens = F.pad(tokens, (0, pad), value=pad_value)
+    s1, e1 = chunk * cp_rank, chunk * (cp_rank + 1)
+    s2, e2 = chunk * (2 * cp_size - cp_rank - 1), chunk * (2 * cp_size - cp_rank)
+    return torch.cat([tokens[s1:e1], tokens[s2:e2]])
+
+
+def _build_like_get_batch(sample_lens, cp_size, pad_size=8):
+    """Mimic miles get_batch THD+CP packing: per-sample zigzag slice, concat, pad, cu*cp."""
+    samples = []
+    base = 1
+    for L in sample_lens:
+        samples.append(torch.arange(base, base + L))  # unique nonzero ids
+        base += L
+    per_rank = []
+    for r in range(cp_size):
+        row = torch.cat([_slice_with_cp(t, cp_size, r) for t in samples])
+        per_rank.append(row)
+    cu = [0]
+    for t in samples:
+        cu.append(cu[-1] + _slice_with_cp(t, cp_size, 0).size(0))
+    final_pad = (pad_size - per_rank[0].size(0) % pad_size) % pad_size
+    if final_pad:
+        per_rank = [F.pad(row, (0, final_pad), value=0) for row in per_rank]
+        cu.append(cu[-1] + final_pad)
+    cu = [x * cp_size for x in cu]
+    return samples, per_rank, cu
+
+
+@pytest.mark.parametrize(
+    "cp_size,sample_lens",
+    [(2, [10, 7, 13]), (2, [16, 16]), (4, [20, 9, 30, 5]), (2, [3]), (4, [40, 17])],
+)
+def test_reassemble_is_inverse_of_slice_with_cp(cp_size, sample_lens):
+    samples, per_rank, cu = _build_like_get_batch(sample_lens, cp_size)
+    local_len = per_rank[0].size(0)
+    assert cu[-1] == cp_size * local_len
+
+    full = _reassemble_full_row(per_rank, cu, cp_size)
+    assert full is not None and full.numel() == cu[-1]
+
+    # Each real sample's tokens reappear (in order) at the start of its segment.
+    for i, t in enumerate(samples):
+        seg = full[cu[i] : cu[i + 1]]
+        assert torch.equal(seg[: t.numel()], t)
+
+    # Re-slicing the full row per segment recovers exactly each rank's local chunks.
+    for r in range(cp_size):
+        recon = []
+        for i in range(len(cu) - 1):
+            recon.append(_natural_to_zigzag_slice(full[cu[i] : cu[i + 1]], cp_size, r, dim=0))
+        assert torch.equal(torch.cat(recon), per_rank[r])
+
+
+def test_reassemble_bails_on_indivisible_segment():
+    # A segment length not divisible by 2*cp -> None (caller falls back to dense path).
+    cu = [0, 6]  # 6 not divisible by 2*cp=4
+    gathered = [torch.zeros(3, dtype=torch.long), torch.zeros(3, dtype=torch.long)]
+    assert _reassemble_full_row(gathered, cu, cp_size=2) is None