modify init pg

Author: kip-cxj

checkpoint_engine/distributed/base.py

@@ -106,18 +106,17 @@ def init_process_group(
         rank: int,
         world_size: int,
         store: torch.distributed.TCPStore,
-        timeout: timedelta,
         **kwargs,
     ):
         backend = kwargs.get("backend", "nccl")
-        store_counter = kwargs.get("store_counter", "1")
-        sub_store = torch.distributed.PrefixStore(f"prefix-{store_counter}", store)
+        timeout = kwargs.get("timeout", timedelta(minutes=10))
+
         torch.distributed.init_process_group(
             backend=backend,
             world_size=world_size,
             rank=rank,
             timeout=timeout,
-            store=sub_store,
+            store=store,
         )
 
     def destroy_process_group(self, group: DistributedProcessGroup | None = None):
@@ -244,10 +243,9 @@ def init_process_group(
     rank: int,
     world_size: int,
     store: torch.distributed.TCPStore,
-    timeout: timedelta = timedelta(seconds=300),
     **kwargs,
 ):
-    _BACKEND_INSTANCE.init_process_group(rank, world_size, store, timeout, **kwargs)
+    _BACKEND_INSTANCE.init_process_group(rank, world_size, store, **kwargs)
 
 
 def destroy_process_group(group: DistributedProcessGroup | None = None):

checkpoint_engine/distributed/hccl.py

@@ -1,6 +1,5 @@
 import ctypes
 from contextlib import contextmanager
-from datetime import timedelta
 from typing import Any, ClassVar
 
 import torch
@@ -233,20 +232,15 @@ def init_process_group(
         rank: int,
         world_size: int,
         store: torch.distributed.TCPStore,
-        timeout: timedelta = timedelta(seconds=300),
         **kwargs,
     ):
         assert not self.initialized, "already initialized"
 
-        self.host = store.host
-        self.port = store.port + 1
         self.rank = rank
         self.world_size = world_size
         self.device = torch.device("npu", torch.npu.current_device())
 
-        self.pg = StatelessProcessGroup.create(
-            self.host, self.port, rank, world_size, store_timeout=int(timeout.total_seconds())
-        )
+        self.pg = StatelessProcessGroup(rank=rank, world_size=world_size, store=store, socket=None)
         self.pyhccl = PyHcclCommunicatorEx(group=self.pg, device=self.device)
         self.comm = self.pyhccl.comm
         self.initialized = True

checkpoint_engine/distributed/nccl.py

@@ -1,6 +1,5 @@
 import ctypes
 from contextlib import contextmanager
-from datetime import timedelta
 from typing import Any, ClassVar
 
 import torch
@@ -136,21 +135,15 @@ def init_process_group(
         rank: int,
         world_size: int,
         store: torch.distributed.TCPStore,
-        timeout: timedelta = timedelta(seconds=300),
         **kwargs,
     ):
         assert not self.initialized, "already initialized"
 
-        self.host = store.host
-        self.port = store.port + 1
         self.rank = rank
         self.world_size = world_size
         self.device = torch.device("cuda", torch.cuda.current_device())
 
-        self.pg = StatelessProcessGroup.create(
-            self.host, self.port, rank, world_size, store_timeout=int(timeout.total_seconds())
-        )
-
+        self.pg = StatelessProcessGroup(rank=rank, world_size=world_size, store=store, socket=None)
         self.pynccl = PyNcclCommunicatorEx(group=self.pg, device=self.device)
         self.comm = self.pynccl.comm
         self.initialized = True

checkpoint_engine/ps.py

@@ -511,13 +511,13 @@ def init_process_group(
             timeout: The timeout of the process group.
         """
         self._store_counter += 1
+        sub_store = torch.distributed.PrefixStore(f"prefix-{self._store_counter}", self._store)
         dist.init_process_group(
-            rank=self._rank,
+            backend=self.device_manager.backend,
             world_size=self._world_size,
-            store=self._store,
+            rank=self._rank,
             timeout=timeout,
-            backend=self.device_manager.backend,
-            store_counter=self._store_counter,
+            store=sub_store,
         )
         logger.info(f"[rank{self._rank}] init process group successfully.")