[4/N] Test NaN checker against broadcast (#134701)

kwen2501 · pytorchmergebot · commit d9d95dc55e3e · 2024-08-29T17:00:07.000Z
Pull Request resolved: #134701 Approved by: https://github.com/wconstab ghstack dependencies: #134345, #134357
diff --git a/test/distributed/test_c10d_nccl.py b/test/distributed/test_c10d_nccl.py
@@ -369,7 +369,7 @@ def test_nan_assert(self, type):
 
     @requires_nccl()
     @skip_if_lt_x_gpu(2)
-    def test_nan_p2p(self):
+    def test_nan_rank_filter(self):
         # Putting NaN at recv buffer, program should not fail as NaN checker
         # should not check on receive buffer
         os.environ["TORCH_NCCL_NAN_CHECK"] = "1"
@@ -379,11 +379,15 @@ def test_nan_p2p(self):
             backend="nccl", store=store, rank=self.rank, world_size=self.world_size
         )
         t = torch.ones(3, 4, dtype=torch.bfloat16, device=device)
+        if self.rank != 0:
+            # Putting NaN at recv buffer
+            t[1, 1] = float("nan")
+        # Against broadcast
+        c10d.broadcast(t, 0)
+        # Against P2P
         if self.rank == 0:
             c10d.send(t, 1)
         elif self.rank == 1:
-            # Putting NaN at recv buffer
-            t[1, 1] = float("nan")
             c10d.recv(t, 0)
         c10d.destroy_process_group()
         # reset env
