Update on "[PGNCCL] Use non-blocking mode by default in eager init"

kwen2501 · kwen2501 · commit 821fc923f336 · 2024-10-25T16:59:20.000-07:00
### Why use non-blocking mode in eager init? For overlapping comm init and model init, etc. ![image](https://github.com/user-attachments/assets/9b0bf7a9-be26-4d16-827b-dbe861f083cd) ### Why can we set non-blocking as default? If the setting is dangling -- i.e. not passed in by user nor set via env -- `ProcessGroupNCCL` can have some preferred logic. And torch-level API semantics does not change whether the NCCL comm is blocking or non-blocking (handled within `ProcessGroupNCCL`). ### Why not make non-blocking default for lazy mode as well? PR #137544 tried it. Two reasons why that's not preferred today: 1. It is hard -- too big a blast. 2. There is no gain by doing lazy init in non-blocking mode, because the right next CPU call is a collective, and we will block there waiting for comm to be ready, so same effect as blocked init, no "opening" compared to eager mode. cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]
diff --git a/test/distributed/test_c10d_nccl.py b/test/distributed/test_c10d_nccl.py
@@ -865,15 +865,11 @@ def test_comm_eager_init_subgroup(self):
         pg = self._create_process_group_nccl(store, self.opts())
         backend = pg._get_backend(torch.device(device))
         self.assertEqual(backend._is_initialized(), False)
-
-        tensor = torch.full((1,), self.rank).cuda(device)
+        # create a subgroup eagerly
         new_group = c10d.new_group([0, 1], device_id=device)
-        self.assertEqual(backend.comm_split_count(), 0)
-
-        new_backend = new_group._get_backend(torch.device(device))
-        self.assertEqual(new_backend._is_initialized(), True)
+        tensor = torch.full((1,), self.rank).cuda(device)
         dist.broadcast(tensor, 0, group=new_group)
-        self.assertEqual(new_backend.comm_split_count(), 0)
+        # the default group should stay lazy
         self.assertEqual(backend._is_initialized(), False)
         torch.cuda.synchronize()
         dist.destroy_process_group()
