diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py index 0fc881d9d6..7d1e1b651d 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py @@ -51,7 +51,10 @@ split_deepstack_embs, ) from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.vision_model import Qwen3VLVisionModel -from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_cp_partition_indices +from megatron.bridge.training.utils.packed_seq_utils import ( + get_packed_seq_cp_partition_indices, + get_packed_seq_q_cu_seqlens, +) def _is_mrope_position_ids(position_ids: torch.Tensor | None) -> bool: @@ -85,6 +88,123 @@ def _split_if_full_sequence( return split_data_cp_rank(val, cp_size, seq_dim, cp_rank), True +def _is_packed_input_pre_sharded( + input_ids: torch.Tensor | None, + packed_seq_params: PackedSeqParams | None, + *, + cp_size: int, +) -> bool: + """Return whether a THD input already uses MCore's zigzag CP layout.""" + if ( + input_ids is None + or packed_seq_params is None + or cp_size <= 1 + or input_ids.dim() != 2 + or input_ids.size(0) != 1 + or packed_seq_params.qkv_format != "thd" + ): + return False + + _, physical_cu_seqlens = get_packed_seq_q_cu_seqlens(packed_seq_params) + if ( + not isinstance(physical_cu_seqlens, torch.Tensor) + or physical_cu_seqlens.dim() != 1 + or physical_cu_seqlens.numel() < 2 + ): + return False + + full_token_count = int(physical_cu_seqlens[-1].item()) + if full_token_count != cp_size * input_ids.numel(): + return False + if int(physical_cu_seqlens[0].item()) != 0: + raise ValueError("Pre-sharded packed CP metadata must start at token offset 0.") + + chunk_count = 2 * cp_size + segment_lengths = physical_cu_seqlens[1:] - physical_cu_seqlens[:-1] + if bool(torch.any(segment_lengths % chunk_count != 0).item()): + raise ValueError( + "Pre-sharded packed CP inputs require every physical segment length to be divisible by 2 * cp_size." + ) + return True + + +def _get_cp_local_vision_embed_indices( + vision_mask: torch.Tensor, + packed_seq_params: PackedSeqParams, + *, + vision_embed_count: int, + cp_group: torch.distributed.ProcessGroup, + embed_device: torch.device, +) -> torch.Tensor: + """Map full-sequence vision embeddings to one pre-sharded CP rank. + + The local THD row contains two chunks from each packed segment: chunk + ``cp_rank`` and chunk ``2 * cp_size - 1 - cp_rank``. Gathering only the + per-chunk vision-token counts reconstructs offsets into the full, natural + vision-embedding order without communicating token IDs. + """ + _, physical_cu_seqlens = get_packed_seq_q_cu_seqlens(packed_seq_params) + if not isinstance(physical_cu_seqlens, torch.Tensor): + raise ValueError("Pre-sharded packed CP vision selection requires physical cu_seqlens metadata.") + + cp_size = cp_group.size() + cp_rank = cp_group.rank() + chunk_count = 2 * cp_size + cu_seqlens = physical_cu_seqlens.tolist() + local_vision_mask = vision_mask.reshape(-1) + expected_local_tokens = cu_seqlens[-1] // cp_size + if local_vision_mask.numel() != expected_local_tokens: + raise ValueError( + "Pre-sharded packed CP vision mask length does not match the global packed-sequence metadata." + ) + + segment_count = len(cu_seqlens) - 1 + local_counts = torch.zeros(segment_count, 2, dtype=torch.long, device=local_vision_mask.device) + for segment_idx in range(segment_count): + segment_length = cu_seqlens[segment_idx + 1] - cu_seqlens[segment_idx] + chunk_length = segment_length // chunk_count + local_start = cu_seqlens[segment_idx] // cp_size + local_counts[segment_idx, 0] = local_vision_mask[local_start : local_start + chunk_length].sum() + local_counts[segment_idx, 1] = local_vision_mask[ + local_start + chunk_length : local_start + 2 * chunk_length + ].sum() + + gathered_counts = [torch.empty_like(local_counts) for _ in range(cp_size)] + torch.distributed.all_gather(gathered_counts, local_counts, group=cp_group) + + full_counts = torch.zeros( + segment_count, + chunk_count, + dtype=torch.long, + device=local_vision_mask.device, + ) + for rank, rank_counts in enumerate(gathered_counts): + full_counts[:, rank] = rank_counts[:, 0] + full_counts[:, chunk_count - 1 - rank] = rank_counts[:, 1] + + gathered_vision_count = int(full_counts.sum().item()) + if gathered_vision_count != vision_embed_count: + raise ValueError( + f"Packed CP ranks contain {gathered_vision_count} vision tokens, but the vision encoder produced " + f"{vision_embed_count} embeddings." + ) + + flattened_counts = full_counts.reshape(-1) + offsets = (torch.cumsum(flattened_counts, dim=0) - flattened_counts).reshape(segment_count, chunk_count) + index_parts = [] + for segment_idx in range(segment_count): + for local_chunk_idx, full_chunk_idx in enumerate((cp_rank, chunk_count - 1 - cp_rank)): + count = int(local_counts[segment_idx, local_chunk_idx].item()) + if count == 0: + continue + start = int(offsets[segment_idx, full_chunk_idx].item()) + index_parts.append(torch.arange(start, start + count, dtype=torch.long, device=embed_device)) + + if not index_parts: + return torch.empty(0, dtype=torch.long, device=embed_device) + return torch.cat(index_parts) + + class Qwen3VLModel(MegatronModule): """Qwen3VL multi-modal model. @@ -397,6 +517,7 @@ def forward( input_ids (torch.Tensor): input text ids [batch, text_seq_len]. position_ids (torch.Tensor): Optional explicit Qwen MRoPE position ids [3, batch, text_seq_len]. Ordinary 2D text position ids are ignored and MRoPE is computed from ``input_ids`` and visual grids. + Inputs that are already sharded across context-parallel ranks must provide rank-local MRoPE IDs. attention_mask (torch.Tensor): attention mask for the language model [batch, 1, combined_seq_len, combined_seq_len]. labels (torch.Tensor): Optional target text labels [batch, combined_seq_len]. @@ -429,6 +550,13 @@ def forward( legacy_packed_bshd = ( packed_seq_params is not None and input_ids is not None and input_ids.dim() == 2 and input_ids.size(0) > 1 ) + packed_input_pre_sharded = _is_packed_input_pre_sharded( + input_ids, + packed_seq_params, + cp_size=cp_size, + ) + if packed_input_pre_sharded and position_ids is None: + raise ValueError("Pre-sharded packed CP inputs require explicit rank-local 3D MRoPE position_ids.") packed_cp_index = ( get_packed_seq_cp_partition_indices( packed_seq_params, @@ -438,7 +566,11 @@ def forward( device=input_ids.device, cp_group=self.pg_collection.cp, ) - if packed_seq_params is not None and cp_size > 1 and input_ids is not None and not legacy_packed_bshd + if packed_seq_params is not None + and cp_size > 1 + and input_ids is not None + and not legacy_packed_bshd + and not packed_input_pre_sharded else None ) @@ -550,6 +682,22 @@ def forward( if vision_embeds is not None: combined_embeddings = combined_embeddings.transpose(0, 1).contiguous() + if packed_input_pre_sharded: + if vision_mask is None: + raise ValueError("Pre-sharded packed CP vision inputs require a rank-local vision mask.") + vision_embed_indices = _get_cp_local_vision_embed_indices( + vision_mask, + packed_seq_params, + vision_embed_count=vision_embeds.size(0), + cp_group=self.pg_collection.cp, + embed_device=vision_embeds.device, + ) + vision_embeds = vision_embeds.index_select(0, vision_embed_indices) + if deepstack_feature_lists is not None: + deepstack_feature_lists = [ + deepstack_embeds.index_select(0, vision_embed_indices) + for deepstack_embeds in deepstack_feature_lists + ] combined_embeddings[vision_mask] = vision_embeds combined_embeddings = combined_embeddings.transpose(0, 1).contiguous() @@ -601,7 +749,7 @@ def forward( .transpose(0, 1) .contiguous() ) - elif cp_size > 1 and packed_cp_index is None: + elif cp_size > 1 and packed_cp_index is None and not packed_input_pre_sharded: raise ValueError("Qwen3VLModel requires input_ids for packed CP slicing") elif packed_cp_index is not None: lm_input_ids = _select_sequence(input_ids, packed_cp_index, seq_dim=1) diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py index 5bbb8a53fd..344f5446dc 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py @@ -35,11 +35,96 @@ from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from transformers import Qwen3VLMoeConfig -from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel +from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import ( + Qwen3VLModel, + _get_cp_local_vision_embed_indices, + _is_packed_input_pre_sharded, +) from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.transformer_config import Qwen3VLTransformerConfig from megatron.bridge.models.qwen_vl.qwen3_vl_provider import DistTrainConfig +def _make_packed_seq_params(cu_seqlens: list[int]) -> PackedSeqParams: + cu_seqlens_tensor = torch.tensor(cu_seqlens, dtype=torch.int32) + max_seqlen = max(end - start for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:])) + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_tensor, + cu_seqlens_kv=cu_seqlens_tensor, + cu_seqlens_q_padded=cu_seqlens_tensor, + cu_seqlens_kv_padded=cu_seqlens_tensor, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + ) + + +def test_is_packed_input_pre_sharded_uses_global_physical_length(): + packed_seq_params = _make_packed_seq_params([0, 8, 16]) + + assert _is_packed_input_pre_sharded(torch.zeros((1, 8), dtype=torch.long), packed_seq_params, cp_size=2) + assert not _is_packed_input_pre_sharded(torch.zeros((1, 16), dtype=torch.long), packed_seq_params, cp_size=2) + + +@pytest.mark.parametrize( + ("cp_rank", "local_vision_mask", "expected_indices"), + [ + (0, [1, 0, 1, 1, 0, 0, 1, 0], [0, 3, 4, 8]), + (1, [1, 1, 0, 0, 1, 0, 1, 1], [1, 2, 5, 6, 7]), + ], +) +def test_get_cp_local_vision_embed_indices_multiple_segments( + cp_rank, + local_vision_mask, + expected_indices, + monkeypatch, +): + packed_seq_params = _make_packed_seq_params([0, 8, 16]) + counts_by_rank = ( + torch.tensor([[1, 2], [0, 1]], dtype=torch.long), + torch.tensor([[2, 0], [1, 2]], dtype=torch.long), + ) + cp_group = SimpleNamespace(size=lambda: 2, rank=lambda: cp_rank) + + def fake_all_gather(outputs, _local_counts, group): + assert group is cp_group + for output, counts in zip(outputs, counts_by_rank): + output.copy_(counts) + + monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather) + + indices = _get_cp_local_vision_embed_indices( + torch.tensor(local_vision_mask, dtype=torch.bool).reshape(1, -1), + packed_seq_params, + vision_embed_count=9, + cp_group=cp_group, + embed_device=torch.device("cpu"), + ) + + assert indices.tolist() == expected_indices + + +def test_get_cp_local_vision_embed_indices_allows_zero_vision_tokens(monkeypatch): + packed_seq_params = _make_packed_seq_params([0, 8]) + cp_group = SimpleNamespace(size=lambda: 2, rank=lambda: 0) + + def fake_all_gather(outputs, _local_counts, group): + assert group is cp_group + outputs[0].zero_() + outputs[1].copy_(torch.tensor([[1, 1]], dtype=torch.long)) + + monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather) + + indices = _get_cp_local_vision_embed_indices( + torch.zeros((1, 4), dtype=torch.bool), + packed_seq_params, + vision_embed_count=2, + cp_group=cp_group, + embed_device=torch.device("cpu"), + ) + + assert indices.shape == (0,) + + @pytest.fixture(scope="module") def hf_config(): """Create a local HuggingFace config once for all tests.""" @@ -717,6 +802,105 @@ def __call__(self, **kwargs): assert language_model.last_kwargs["loss_mask"] is loss_mask assert language_model.last_kwargs["packed_seq_params"] is packed_seq_params + def test_forward_preserves_pre_sharded_packed_cp_layout_and_selects_vision_embeds(self, monkeypatch): + """Pre-sharded CP inputs stay local and select matching vision and deepstack rows.""" + monkeypatch.setattr( + "megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model.torch.cuda.nvtx.range_push", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + "megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model.torch.cuda.nvtx.range_pop", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + "megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model.get_packed_seq_cp_partition_indices", + lambda *_args, **_kwargs: pytest.fail("pre-sharded input must not be partitioned again"), + ) + + local_vision_mask = torch.tensor([[True, False, True, False]]) + monkeypatch.setattr( + "megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model.reorganize_inputs", + lambda **_kwargs: ( + torch.ones((1, 2)), + torch.tensor([[1, 1, 1]], dtype=torch.long), + local_vision_mask, + ), + ) + + cp_group = SimpleNamespace(size=lambda: 2, rank=lambda: 0) + + def fake_all_gather(outputs, _local_counts, group): + assert group is cp_group + outputs[0].copy_(torch.tensor([[1, 1]], dtype=torch.long)) + outputs[1].copy_(torch.tensor([[1, 1]], dtype=torch.long)) + + monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather) + + full_vision_embeds = torch.tensor([[10.0, 10.0], [20.0, 20.0], [30.0, 30.0], [40.0, 40.0]]) + full_deepstack_embeds = full_vision_embeds + 100.0 + + class DummyVisionModel: + def __call__(self, **_kwargs): + return full_vision_embeds, [full_deepstack_embeds] + + class DummyLanguageModel: + def __init__(self): + self.rotary_pos_emb = SimpleNamespace(is_thd_format=False) + self.last_kwargs = None + + def embedding(self, input_ids, position_ids=None): + del position_ids + return torch.zeros((input_ids.size(1), input_ids.size(0), 2)) + + def __call__(self, **kwargs): + self.last_kwargs = kwargs + return torch.ones(1) + + language_model = DummyLanguageModel() + model = SimpleNamespace( + pre_process=True, + square_merge_size=1, + config=SimpleNamespace( + vision_dp_when_cp=False, + sequence_parallel=False, + spatial_merge_size=1, + ), + pg_collection=SimpleNamespace( + cp=cp_group, + tp=SimpleNamespace(rank=lambda: 0, size=lambda: 1), + pp=object(), + ), + language_model=language_model, + vision_model=DummyVisionModel(), + image_token_id=1, + video_token_id=2, + vision_start_token_id=3, + use_dist_train=False, + ) + input_ids = torch.tensor([[1, 11, 1, 12]], dtype=torch.long) + position_ids = torch.arange(4).reshape(1, 1, 4).expand(3, -1, -1).clone() + packed_seq_params = _make_packed_seq_params([0, 8]) + + output = Qwen3VLModel.forward( + model, + input_ids=input_ids, + position_ids=position_ids, + packed_seq_params=packed_seq_params, + pixel_values=torch.ones(1), + image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long), + ) + + assert torch.equal(output, torch.ones(1)) + assert language_model.last_kwargs is not None + assert language_model.last_kwargs["input_ids"] is input_ids + assert language_model.last_kwargs["position_ids"] is position_ids + decoder_input = language_model.last_kwargs["decoder_input"].transpose(0, 1) + torch.testing.assert_close(decoder_input[local_vision_mask], full_vision_embeds[[0, 3]]) + assert torch.equal(language_model.last_kwargs["visual_pos_masks"], local_vision_mask) + deepstack_embeds = language_model.last_kwargs["deepstack_visual_embeds"] + assert deepstack_embeds is not None + torch.testing.assert_close(deepstack_embeds[0], full_deepstack_embeds[[0, 3]]) + def test_forward_applies_one_partition_index_to_packed_cp_tensors(self, monkeypatch): """Packed CP slices every sequence-aligned tensor with the same index.""" monkeypatch.setattr(