Fix MTP CP-boundary roll to communicate on the context-parallel process group

megatron/core/transformer/multi_token_prediction.py

@@ -278,28 +278,42 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non
     next_rank = global_ranks[(local_rank + 1) % len(global_ranks)]
     prev_rank = global_ranks[(local_rank - 1) % len(global_ranks)]
 
-    # Start send and recv ops
+    # Keep each boundary exchange on its CP communicator. Using the default WORLD
+    # communicator couples otherwise independent DP lanes and can deadlock when
+    # packed batches issue different P2P sequences.
     ops = []
     if local_rank != 0:
-        req_send_first_part = torch.distributed.isend(tensor=tensor_send_list[0], dst=prev_rank)
-        ops.append(req_send_first_part)
-        req_recv_second_part = torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank)
-        ops.append(req_recv_second_part)
+        ops.append(
+            torch.distributed.P2POp(
+                torch.distributed.isend, tensor_send_list[0], prev_rank, group=cp_group
+            )
+        )
+        ops.append(
+            torch.distributed.P2POp(
+                torch.distributed.irecv, tensor_recv_list[1], prev_rank, group=cp_group
+            )
+        )
     else:
         # Inserted elements are set to be 0.0.
         tensor_recv_list[1] = 0
     if local_rank != len(global_ranks) - 1:
-        req_recv_first_part = torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank)
-        ops.append(req_recv_first_part)
-        req_send_second_part = torch.distributed.isend(tensor=tensor_send_list[1], dst=next_rank)
-        ops.append(req_send_second_part)
+        ops.append(
+            torch.distributed.P2POp(
+                torch.distributed.irecv, tensor_recv_list[0], next_rank, group=cp_group
+            )
+        )
+        ops.append(
+            torch.distributed.P2POp(
+                torch.distributed.isend, tensor_send_list[1], next_rank, group=cp_group
+            )
+        )
     else:
         # For the last CP rank, the removed elements of second part go into the first part
         tensor_recv_list[0] = tensor_send_list[1]
 
     # Wait for all communication operations to complete
-    for op in ops:
-        op.wait()
+    for request in torch.distributed.batch_isend_irecv(ops):
+        request.wait()
 
     # Splicing: Replace boundary elements with received elements from adjacent ranks
     # This ensures proper sequence continuity across CP boundaries
@@ -398,19 +412,35 @@ def _roll_tensor_packed_seq(
 
         ops = []
         if local_rank != 0:
-            ops.append(torch.distributed.isend(tensor=tensor_send_list[0], dst=prev_rank))
-            ops.append(torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank))
+            ops.append(
+                torch.distributed.P2POp(
+                    torch.distributed.isend, tensor_send_list[0], prev_rank, group=cp_group
+                )
+            )
+            ops.append(
+                torch.distributed.P2POp(
+                    torch.distributed.irecv, tensor_recv_list[1], prev_rank, group=cp_group
+                )
+            )
         else:
             tensor_recv_list[1].zero_()
 
         if local_rank != cp_size - 1:
-            ops.append(torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank))
-            ops.append(torch.distributed.isend(tensor=tensor_send_list[1], dst=next_rank))
+            ops.append(
+                torch.distributed.P2POp(
+                    torch.distributed.irecv, tensor_recv_list[0], next_rank, group=cp_group
+                )
+            )
+            ops.append(
+                torch.distributed.P2POp(
+                    torch.distributed.isend, tensor_send_list[1], next_rank, group=cp_group
+                )
+            )
         else:
             tensor_recv_list[0].copy_(tensor_send_list[1])
 
-        for op in ops:
-            op.wait()
+        for request in torch.distributed.batch_isend_irecv(ops):
+            request.wait()
 
         index = [slice(None)] * rolled_chunks[0].dim()
         index[dims] = shifts

tests/unit_tests/transformer/test_multi_token_prediction.py

@@ -3931,6 +3931,90 @@ def test_attention_mask_validation_mamba(self):
                 raise
 
 
+class TestRollTensorWithCPSubgroup:
+    """Roll-tensor CP paths on a CP group that is a strict subgroup of WORLD.
+
+    Regression coverage for the CP-boundary exchange being issued on the
+    default (WORLD) process group instead of ``cp_group``: with
+    WORLD == CP group the two are indistinguishable, so these tests run
+    TP2 x CP2 to make the CP group a strict subgroup.
+    """
+
+    def setup_method(self, method):
+        Utils.initialize_model_parallel(tensor_model_parallel_size=2, context_parallel_size=2)
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+
+    @staticmethod
+    def _cp_shard(tensor, cp_rank, cp_size, dim=-1):
+        """Shard along ``dim`` in the load-balanced two-chunk CP layout."""
+        chunks = tensor.chunk(2 * cp_size, dim=dim)
+        return torch.cat([chunks[cp_rank], chunks[2 * cp_size - 1 - cp_rank]], dim=dim)
+
+    def test_roll_tensor_cp_boundary_exchange(self, monkeypatch):
+        cp_group = get_context_parallel_group()
+        cp_size = cp_group.size()
+        cp_rank = torch.distributed.get_rank(group=cp_group)
+
+        # Spy on the batched P2P launch: every boundary-exchange op must be
+        # addressed to the CP communicator, never the default (WORLD) group.
+        seen_groups = []
+        real_batch_isend_irecv = torch.distributed.batch_isend_irecv
+
+        def spy_batch_isend_irecv(ops):
+            seen_groups.extend(op.group for op in ops)
+            return real_batch_isend_irecv(ops)
+
+        monkeypatch.setattr(torch.distributed, "batch_isend_irecv", spy_batch_isend_irecv)
+
+        full = torch.arange(2 * 16, dtype=torch.float32, device="cuda").reshape(2, 16)
+        # Oracle: the documented single-rank semantics on the full sequence.
+        expected_full, _ = roll_tensor(full.clone(), shifts=-1, dims=-1)
+        expected_local = self._cp_shard(expected_full, cp_rank, cp_size)
+
+        local = self._cp_shard(full, cp_rank, cp_size).contiguous()
+        rolled, rolled_sum = roll_tensor(local, shifts=-1, dims=-1, cp_group=cp_group)
+
+        assert torch.equal(rolled, expected_local)
+        assert rolled_sum == expected_local.sum()
+        assert seen_groups, "CP boundary exchange must go through batch_isend_irecv"
+        assert all(group is cp_group for group in seen_groups)
+
+    def test_roll_tensor_packed_seq_cp_boundary_exchange(self):
+        cp_group = get_context_parallel_group()
+        cp_size = cp_group.size()
+        cp_rank = torch.distributed.get_rank(group=cp_group)
+
+        seq_lengths = [8, 16]  # each divisible by 2 * cp_size
+        cu_seqlens = torch.tensor([0, 8, 24], dtype=torch.int32, device="cuda")
+        full = torch.arange(1 * 24, dtype=torch.float32, device="cuda").reshape(1, 24)
+
+        packed_seq_params = PackedSeqParams(cu_seqlens_q=cu_seqlens)
+        # Oracle: single-rank packed semantics (per-sequence roll) on the full tensor.
+        expected_full, _ = roll_tensor(
+            full.clone(), shifts=-1, dims=-1, packed_seq_params=packed_seq_params
+        )
+
+        def shard_packed(tensor):
+            pieces = []
+            start = 0
+            for length in seq_lengths:
+                seq = tensor[..., start : start + length]
+                pieces.append(self._cp_shard(seq, cp_rank, cp_size))
+                start += length
+            return torch.cat(pieces, dim=-1)
+
+        expected_local = shard_packed(expected_full)
+        local = shard_packed(full).contiguous()
+        rolled, rolled_sum = roll_tensor(
+            local, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
+        )
+
+        assert torch.equal(rolled, expected_local)
+        assert rolled_sum == expected_local.sum()
+
+
 class TestLearnedOutputContract:
     """Tests for the learned n-stream to one-stream mHC contraction."""