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 8a2876fb30e..49f90eedcc6 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 @@ -40,6 +40,7 @@ AllGatherVisionEmbeddings, PatchMergerSubmodules, collapse_thw, + ensure_requires_grad_for_cp_collective, get_dist_train_vision_dp_data, get_vision_cp_data, pack_dist_train_vision_module_output, @@ -516,7 +517,7 @@ def forward( vision_embeds = torch.zeros( (0, self.config.hidden_size), device=vision_data.device, - dtype=torch.bfloat16, + dtype=self.config.params_dtype, ) deepstack_feature_lists = [] for _ in self.vision_transformer_config.deepstack_visual_indexes: @@ -524,20 +525,21 @@ def forward( torch.zeros( (0, self.language_model.config.hidden_size), device=vision_data.device, - dtype=torch.bfloat16, + dtype=self.config.params_dtype, ) ) if cp_size > 1 and self.config.vision_dp_when_cp: + ensure_requires_grad_for_cp_collective((vision_embeds, *deepstack_feature_lists)) vision_embeds = AllGatherVisionEmbeddings.apply( vision_embeds, seqlen_on_cp_ranks, - cp_group=self.pg_collection.cp, + self.pg_collection.cp, ) for i in range(len(deepstack_feature_lists)): deepstack_feature_lists[i] = AllGatherVisionEmbeddings.apply( deepstack_feature_lists[i], seqlen_on_cp_ranks, - cp_group=self.pg_collection.cp, + self.pg_collection.cp, ) combined_embeddings = self.language_model.embedding( diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py index be12d17a798..4078b147cb9 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py @@ -663,6 +663,7 @@ def forward(ctx, input, seqlens_on_cp_ranks, cp_group: torch.distributed.Process outputs.append(o) torch.distributed.all_gather(outputs, input, group=cp_group) ctx.cp_rank = torch.distributed.get_rank(group=cp_group) + ctx.cp_group = cp_group ctx.save_for_backward(*seqlens_on_cp_ranks) output = torch.cat(outputs, dim=0) @@ -677,10 +678,28 @@ def backward(ctx, grad_output): seqlens_on_cp_ranks = ctx.saved_tensors start_idx = torch.cat(seqlens_on_cp_ranks[:cp_rank]).sum() if cp_rank != 0 else 0 end_idx = start_idx + seqlens_on_cp_ranks[cp_rank].sum() + grad_output = grad_output.contiguous() + torch.distributed.all_reduce(grad_output, group=ctx.cp_group) grad_output = grad_output[start_idx:end_idx] return grad_output, None, None +def ensure_requires_grad_for_cp_collective(tensors: Sequence[torch.Tensor]) -> None: + """Make every tensor entering a CP collective participate in autograd. + + ``AllGatherVisionEmbeddings.backward`` issues a cp_group all_reduce, and autograd + only invokes it for tensors that require grad. Tensors that arrive without grad + (0-image placeholders, or the outputs of a fully frozen vision tower) would skip + the collective on their rank and deadlock the remaining ranks, so force + ``requires_grad`` here. No-op under ``torch.no_grad`` (no rank records backward). + """ + if not torch.is_grad_enabled(): + return + for tensor in tensors: + if not tensor.requires_grad: + tensor.requires_grad_(True) + + def preprocess_packed_seqs( input_ids: torch.Tensor, attention_mask: torch.Tensor, diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py index 03e79b8a7e6..933e9e3e127 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py @@ -41,6 +41,7 @@ Qwen3VLVisionPatchMerger, Qwen3VLVisionRotaryEmbedding, collapse_thw, + ensure_requires_grad_for_cp_collective, expand_thw, get_dist_train_vision_dp_data, get_vision_cp_data, @@ -147,6 +148,58 @@ def test_get_rope_index_video_tokens_use_video_grid(): assert deltas.shape == (1, 1) +def test_ensure_requires_grad_for_cp_collective(): + """Grad mode: no-grad tensors are forced; already-requiring tensors untouched. + Under torch.no_grad the helper is a no-op (no rank records a backward).""" + plain = torch.zeros(2, 3) + empty = torch.zeros(0, 3) + already = torch.zeros(2, 3, requires_grad=True) + ensure_requires_grad_for_cp_collective((plain, empty, already)) + assert plain.requires_grad and empty.requires_grad and already.requires_grad + + with torch.no_grad(): + untouched = torch.zeros(2, 3) + ensure_requires_grad_for_cp_collective((untouched,)) + assert not untouched.requires_grad + + +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_allgather_vision_embeddings_backward_mocked_collectives(cp_rank, monkeypatch): + """Verify backward = reduce-scatter (sum across CP ranks, then slice this rank's + range) without a process group, by stubbing the collectives. The injected peer + gradient varies per row, so a wrong slice offset changes the result, and dropping + the all_reduce loses the peer contribution entirely. + """ + hidden_size = 4 + seqlens_on_cp_ranks = [torch.tensor([2]), torch.tensor([3])] + total = int(sum(s.sum() for s in seqlens_on_cp_ranks)) + peer_grad = torch.arange(total, dtype=torch.float).unsqueeze(1).expand(total, hidden_size) + fake_group = object() + + def fake_all_gather(outputs, inp, group=None): + assert group is fake_group + outputs[cp_rank].copy_(inp) + outputs[1 - cp_rank].fill_(7.0) + + def fake_all_reduce(tensor, group=None): + assert group is fake_group + tensor.add_(peer_grad) + + monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather) + monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: cp_rank) + + local_seqlen = int(seqlens_on_cp_ranks[cp_rank].sum()) + input_ = torch.randn(local_seqlen, hidden_size, requires_grad=True) + output = AllGatherVisionEmbeddings.apply(input_, seqlens_on_cp_ranks, fake_group) + assert output.shape == (total, hidden_size) + output.sum().backward() + + start = int(sum(s.sum() for s in seqlens_on_cp_ranks[:cp_rank])) + expected = (torch.ones(total, hidden_size) + peer_grad)[start : start + local_seqlen] + torch.testing.assert_close(input_.grad, expected) + + class TestQwen3VLUtils: """Test suite for Qwen3VL utility functions.""" @@ -480,6 +533,133 @@ def test_allgather_vision_embeddings(self): self.destroy_parallel_state() + @pytest.mark.skipif( + not torch.cuda.is_available() or int(os.environ.get("WORLD_SIZE", "1")) < 2, + reason="Requires at least 2 GPUs", + ) + def test_allgather_vision_embeddings_backward_reduce_scatter(self): + """Backward must reduce-scatter: grad w.r.t. input is the CP-summed gradient sliced + to this rank's range, not just this rank's local slice. + + Each rank drives a *rank-dependent* upstream gradient on the shared gathered output + (weight = cp_rank + 1). The correct gradient for every rank is therefore the same + row-weighted tensor scaled by S = sum_r (r + 1), sliced to the rank's own range. The + buggy local-slice-only backward would instead yield the per-rank weight, which differs + across ranks and misses the cross-rank contributions. + """ + cp_size = 2 + self._setup_parallel_state(tp_size=1, cp_size=cp_size) + cp_group = parallel_state.get_context_parallel_group() + cp_rank = parallel_state.get_context_parallel_rank() + + local_seqlen = 3 + hidden_size = 8 + device = torch.cuda.current_device() + seqlens_on_cp_ranks = [torch.tensor([local_seqlen], dtype=torch.long, device=device) for _ in range(cp_size)] + total = local_seqlen * cp_size + + input_ = torch.randn(local_seqlen, hidden_size, dtype=torch.float, device=device, requires_grad=True) + output = AllGatherVisionEmbeddings.apply(input_, seqlens_on_cp_ranks, cp_group) + + # Row weight varies along the gathered sequence so the test also exercises the slice offset. + row = (torch.arange(total, device=device, dtype=torch.float) + 1.0).unsqueeze(1) # (total, 1) + rank_weight = float(cp_rank + 1) + loss = (output * (row * rank_weight)).sum() + loss.backward() + + s = sum(r + 1 for r in range(cp_size)) # 3 for cp_size=2 + start = cp_rank * local_seqlen + expected = (row[start : start + local_seqlen] * s).expand(local_seqlen, hidden_size) + assert input_.grad is not None + torch.testing.assert_close(input_.grad, expected) + + self.destroy_parallel_state() + + @pytest.mark.skipif( + not torch.cuda.is_available() or int(os.environ.get("WORLD_SIZE", "1")) < 2, + reason="Requires at least 2 GPUs", + ) + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + def test_allgather_vision_embeddings_empty_rank(self, dtype): + """A CP rank with 0 tokens (num_images < cp_size) must still participate in the + backward all_reduce so the collective stays symmetric and does not hang. + + The empty rank supplies a 0-length input with requires_grad=True (the model forces + this before the gather), so autograd invokes the backward -- and thus the cp_group + all_reduce -- on every rank. The rank that owns the tokens must receive the CP-summed + gradient. + """ + cp_size = 2 + self._setup_parallel_state(tp_size=1, cp_size=cp_size) + cp_group = parallel_state.get_context_parallel_group() + cp_rank = parallel_state.get_context_parallel_rank() + + owner_seqlen = 4 + hidden_size = 8 + device = torch.cuda.current_device() + # Only rank 0 owns tokens; rank 1 is the empty (0-image) rank. + seqlens_on_cp_ranks = [ + torch.tensor([owner_seqlen], dtype=torch.long, device=device), + torch.tensor([0], dtype=torch.long, device=device), + ] + local_seqlen = owner_seqlen if cp_rank == 0 else 0 + + input_ = torch.zeros(local_seqlen, hidden_size, dtype=dtype, device=device) + ensure_requires_grad_for_cp_collective((input_,)) + assert input_.requires_grad + output = AllGatherVisionEmbeddings.apply(input_, seqlens_on_cp_ranks, cp_group) + assert output.shape == (owner_seqlen, hidden_size) + assert output.dtype == dtype + + # Rank-dependent upstream gradient again, so the owner rank must see the CP-summed grad. + rank_weight = float(cp_rank + 1) + (output * rank_weight).sum().backward() + + assert input_.grad is not None + assert input_.grad.shape == (local_seqlen, hidden_size) + if cp_rank == 0: + s = sum(r + 1 for r in range(cp_size)) # 3 + torch.testing.assert_close( + input_.grad, torch.full((owner_seqlen, hidden_size), float(s), device=device, dtype=dtype) + ) + + self.destroy_parallel_state() + + @pytest.mark.skipif( + not torch.cuda.is_available() or int(os.environ.get("WORLD_SIZE", "1")) < 2, + reason="Requires at least 2 GPUs", + ) + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + def test_allgather_vision_embeddings_frozen_input(self, dtype): + """A fully frozen vision tower produces embeddings that do not require grad; the + model forces requires_grad=True before the gather so every CP rank still enters the + backward collective. A forced no-grad-fn leaf must receive the CP-summed gradient. + """ + cp_size = 2 + self._setup_parallel_state(tp_size=1, cp_size=cp_size) + cp_group = parallel_state.get_context_parallel_group() + cp_rank = parallel_state.get_context_parallel_rank() + + local_seqlen = 3 + hidden_size = 8 + device = torch.cuda.current_device() + seqlens_on_cp_ranks = [torch.tensor([local_seqlen], dtype=torch.long, device=device) for _ in range(cp_size)] + + # Frozen tower output: a plain tensor with no grad_fn and requires_grad=False. + input_ = torch.randn(local_seqlen, hidden_size, dtype=dtype, device=device) + assert not input_.requires_grad + ensure_requires_grad_for_cp_collective((input_,)) # what the model does right before the gather + assert input_.requires_grad + + output = AllGatherVisionEmbeddings.apply(input_, seqlens_on_cp_ranks, cp_group) + rank_weight = float(cp_rank + 1) + (output * rank_weight).sum().backward() + + s = float(sum(r + 1 for r in range(cp_size))) # 3 for cp_size=2 + torch.testing.assert_close(input_.grad, torch.full((local_seqlen, hidden_size), s, device=device, dtype=dtype)) + + self.destroy_parallel_state() + def test_split_deepstack_embs_no_tp(self): """Test split_deepstack_embs with tp_size=1.""" visual_pos_masks = torch.tensor([[True, False, True], [False, True, False]])