diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 116675e6118..03e924637b7 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/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"): diff --git a/miles_plugins/models/qwen3_vl.py b/miles_plugins/models/qwen3_vl.py index bca4872269e..6dc29bc02c7 100644 --- a/miles_plugins/models/qwen3_vl.py +++ b/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,7 +91,21 @@ 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 @@ -73,22 +113,56 @@ def patched_get_rope_index(*args, **kwargs): 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): diff --git a/tests/fast/test_qwen3_vl_cp_mrope.py b/tests/fast/test_qwen3_vl_cp_mrope.py new file mode 100644 index 00000000000..a6aaa8908f0 --- /dev/null +++ b/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