fix(model): correct Qwen3-VL vision_dp_when_cp CP gradient and avoid 0-image-rank hang

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,28 +517,29 @@ 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:
                         deepstack_feature_lists.append(
                             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(

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,

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]])