revert to store_based_barrier

Author: kip-cxj

checkpoint_engine/distributed/nccl.py

@@ -223,6 +223,7 @@ def new_group(self, ranks: list[int], **kwargs) -> CommGroup:
         else:
             ranks.sort()
 
+        group: CommGroup = None
         newcomm = self.pynccl.create_newcomm(ranks)
         if newcomm:
             group = CommGroup(newcomm.value, ranks)

checkpoint_engine/ps.py

@@ -553,14 +553,26 @@ def update(
         try:
             master_addr = os.getenv("MASTER_ADDR") or master_addr
             assert master_addr, "master_addr is required"
-            if self._auto_pg and not dist.is_initialized():
-                self.init_process_group(
-                    timeout=timeout, master_addr=master_addr, master_port=master_port
+            if self._auto_pg:
+                if not dist.is_initialized():
+                    self.init_process_group(
+                        timeout=timeout, master_addr=master_addr, master_port=master_port
+                    )
+                manager_store = torch.distributed.distributed_c10d._get_default_store()
+            else:
+                # HACK: MASTER_PORT+2 for barrier store if master_port is not provided, _get_master_port() returns MASTER_PORT+1
+                # If master_port is provided, use master_port+1 for barrier store
+                manager_store = torch.distributed.TCPStore(
+                    master_addr,
+                    _get_master_port(master_port) + 1,
+                    self._world_size,
+                    timeout=timeout,
+                    is_master=self._rank == 0,
                 )
             # if ranks is None or [], it will use fully broadcast to update to all ranks
             ranks_group = dist.new_group(ranks) if ranks else None
             self._update_per_bucket(checkpoint_name, req_func, ranks_group, ranks)
-            dist.barrier()
+            self.store_based_barrier(manager_store)
         except Exception as e:
             logger.exception(
                 f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} error {e}"