[SymmetricMemory] use the python version of empty() and rendezvous() for tests and library ops (#142154)

Yifu Wang · pytorchmergebot · commit 5513e2ec3521 · 2024-12-05T22:09:36.000Z
Pull Request resolved: #142154 Approved by: https://github.com/weifengpy
diff --git a/test/distributed/test_symmetric_memory.py b/test/distributed/test_symmetric_memory.py
@@ -221,7 +221,7 @@ def test_get_signal_pad(self) -> None:
 
         # Sanity check that writes to buffer doesn't corrupt signal_pad
         t = symm_mem.empty(0, device="cuda")
-        symm_mem_hdl = symm_mem.rendezvous(t)
+        symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD)
         signal_pad = symm_mem_hdl.get_signal_pad(self.rank)
         signal_pad.fill_(42)
         t.fill_(0)
@@ -235,7 +235,7 @@ def test_barrier_timeout(self) -> None:
         self._init_process()
 
         t = symm_mem.empty(1, device="cuda")
-        symm_mem_hdl = _SymmetricMemory.rendezvous(t, group=dist.group.WORLD)
+        symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD)
 
         if self.rank == 0:
             with self.assertRaises(RuntimeError):
@@ -256,7 +256,7 @@ def test_put_signal_timeout(self) -> None:
         self._init_process()
 
         t = symm_mem.empty(1, device="cuda")
-        symm_mem_hdl = _SymmetricMemory.rendezvous(t, group=dist.group.WORLD)
+        symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD)
 
         if self.rank == 0:
             with self.assertRaises(RuntimeError):
@@ -280,7 +280,7 @@ def test_wait_signal_timeout(self) -> None:
         self._init_process()
 
         t = symm_mem.empty(1, device="cuda")
-        symm_mem_hdl = _SymmetricMemory.rendezvous(t, group=dist.group.WORLD)
+        symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD)
 
         if self.rank == 0:
             with self.assertRaises(RuntimeError):
@@ -306,23 +306,14 @@ def test_allow_overlapping_devices(self) -> None:
             rank=self.rank,
             store=store,
         )
-        group_name = dist.group.WORLD.group_name
-        enable_symm_mem_for_group(group_name)
-
-        t = _SymmetricMemory.empty_strided_p2p(
-            (64,),
-            (1,),
-            dtype=torch.uint32,
-            device=torch.device("cuda:0"),
-            group_name=group_name,
-        ).fill_(0)
-        symm_mem_handle = _SymmetricMemory.rendezvous(t)
+        t = symm_mem.empty(64, device="cuda:0")
+        symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD)
 
-        self.assertEqual(symm_mem_handle.rank, self.rank)
-        self.assertEqual(symm_mem_handle.world_size, self.world_size)
+        self.assertEqual(symm_mem_hdl.rank, self.rank)
+        self.assertEqual(symm_mem_hdl.world_size, self.world_size)
 
         for rank in range(self.world_size):
-            buf = symm_mem_handle.get_buffer(rank, (64,), torch.float32)
+            buf = symm_mem_hdl.get_buffer(rank, (64,), torch.float32)
             if rank == self.rank:
                 self.assertEqual(buf.data_ptr(), t.data_ptr())
             else:
@@ -754,7 +745,6 @@ def _init_process(self):
             rank=self.rank,
             store=store,
         )
-        enable_symm_mem_for_group(dist.group.WORLD.group_name)
         torch.manual_seed(42 + self.rank)
 
     @skip_if_lt_x_gpu(4)
@@ -768,13 +758,8 @@ def test_multimem_all_reduce(
         self._init_process()
         group_name = dist.group.WORLD.group_name
 
-        t = _SymmetricMemory.empty_strided_p2p(
-            size=(16384,),
-            stride=(1,),
-            dtype=dtype,
-            device=self.device,
-            group_name=group_name,
-        ).fill_(0)
+        t = symm_mem.empty((16384), dtype=dtype, device=self.device)
+        symm_mem.rendezvous(t, group=dist.group.WORLD)
 
         self.assertTrue(t.data_ptr() % 16 == 0)
         self.assertTrue(align_bytes % t.element_size() == 0)
@@ -806,13 +791,10 @@ def test_multimem_one_shot_all_reduce(
         self._init_process()
         group_name = dist.group.WORLD.group_name
 
-        inp = _SymmetricMemory.empty_strided_p2p(
-            size=(size_bytes,),
-            stride=(1,),
-            dtype=dtype,
-            device=self.device,
-            group_name=group_name,
+        inp = symm_mem.empty(
+            size_bytes // dtype.itemsize, dtype=dtype, device=self.device
         ).normal_()
+        symm_mem.rendezvous(inp, group=group_name)
 
         res = torch.ops.symm_mem.multimem_one_shot_all_reduce(inp, "sum", group_name)
 
@@ -835,13 +817,10 @@ def test_one_shot_all_reduce(
         self._init_process()
         group_name = dist.group.WORLD.group_name
 
-        inp = _SymmetricMemory.empty_strided_p2p(
-            size=(size_bytes,),
-            stride=(1,),
-            dtype=dtype,
-            device=self.device,
-            group_name=group_name,
+        inp = symm_mem.empty(
+            size_bytes // dtype.itemsize, dtype=dtype, device=self.device
         ).normal_()
+        symm_mem.rendezvous(inp, group=group_name)
 
         res = torch.ops.symm_mem.one_shot_all_reduce(inp, "sum", group_name)
         self._verify_all_reduce_result(inp, res)
@@ -858,13 +837,8 @@ def test_two_shot_all_reduce(
         self._init_process()
         group_name = dist.group.WORLD.group_name
 
-        t = _SymmetricMemory.empty_strided_p2p(
-            size=(16384,),
-            stride=(1,),
-            dtype=dtype,
-            device=self.device,
-            group_name=group_name,
-        ).fill_(0)
+        t = symm_mem.empty(16384, dtype=dtype, device=self.device).fill_(0)
+        symm_mem.rendezvous(t, group=group_name)
 
         self.assertTrue(t.data_ptr() % 16 == 0)
         self.assertTrue(align_bytes % t.element_size() == 0)
diff --git a/torch/csrc/distributed/c10d/CUDASymmetricMemoryOps.cu b/torch/csrc/distributed/c10d/CUDASymmetricMemoryOps.cu
@@ -154,7 +154,7 @@ at::Tensor multimem_all_reduce_(
       reduce_op == "sum",
       "multimem_all_reduce_: only sum is supported for now.");
 
-  auto symm_mem = c10d::symmetric_memory::rendezvous(input);
+  auto symm_mem = c10d::symmetric_memory::rendezvous(input, group_name);
   TORCH_CHECK(
       symm_mem != nullptr,
       "multimem_all_reduce_: input must be allocated with empty_strided_p2p().");
@@ -240,7 +240,7 @@ at::Tensor multimem_one_shot_all_reduce_out(
       reduce_op == "sum",
       "multimem_one_shot_all_reduce: only sum is supported for now.");
 
-  auto symm_mem = c10d::symmetric_memory::rendezvous(input);
+  auto symm_mem = c10d::symmetric_memory::rendezvous(input, group_name);
   TORCH_CHECK(
       symm_mem != nullptr,
       "multimem_one_shot_all_reduce: input must be allocated with empty_strided_p2p().");
@@ -343,7 +343,7 @@ at::Tensor one_shot_all_reduce_out(
       reduce_op == "sum",
       "one_shot_all_reduce: only sum is supported for now.");
 
-  auto symm_mem = c10d::symmetric_memory::rendezvous(input);
+  auto symm_mem = c10d::symmetric_memory::rendezvous(input, group_name);
   TORCH_CHECK(
       symm_mem != nullptr,
       "one_shot_all_reduce: input must be allocated with empty_strided_p2p().");
@@ -453,7 +453,7 @@ at::Tensor two_shot_all_reduce_(
       reduce_op == "sum",
       "two_shot_all_reduce: only sum is supported for now.");
 
-  auto symm_mem = c10d::symmetric_memory::rendezvous(input);
+  auto symm_mem = c10d::symmetric_memory::rendezvous(input, group_name);
   TORCH_CHECK(
       symm_mem != nullptr,
       "two_shot_all_reduce: input must be allocated with empty_strided_p2p().");
diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py
@@ -667,7 +667,7 @@ def _fused_all_gather_matmul_native(
     B: torch.Tensor,
     group_name: str,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    symm_mem = _SymmetricMemory.rendezvous(A_shard)
+    symm_mem = rendezvous(A_shard, group_name)
     if symm_mem is None:
         symm_mem = get_symm_mem_workspace(
             group_name, A_shard.numel() * A_shard.element_size()
@@ -1235,7 +1235,7 @@ def _low_contention_all_gather(
           symmetric memory workspace.
         - Symmetric memory workspace size requirement: the size of `tensor`.
     """
-    symm_mem = _SymmetricMemory.rendezvous(tensor)
+    symm_mem = rendezvous(tensor, group_name)
     if symm_mem is not None:
         input_is_symm_mem = True
     else:
@@ -1376,7 +1376,7 @@ def _low_contention_reduce_scatter(
     TODO(yifu): the SM-based copy can be avoided with a list-based reduction
     kernel.
     """
-    symm_mem = _SymmetricMemory.rendezvous(tensor)
+    symm_mem = rendezvous(tensor, group_name)
     if symm_mem is not None:
         return _low_contention_reduce_scatter_with_symm_mem_input(
             tensor, reduce_op, symm_mem
