add dist.use_backend

Author: kip-cxj

checkpoint_engine/distributed/__init__.py

@@ -8,6 +8,7 @@
     init_process_group,
     is_initialized,
     new_group,
+    use_backend,
 )
 
 
@@ -21,4 +22,5 @@
     "init_process_group",
     "is_initialized",
     "new_group",
+    "use_backend",
 ]

checkpoint_engine/distributed/base.py

@@ -16,15 +16,15 @@ def all_gather(self, *args: Any, **kwargs: Any) -> torch.Tensor: ...
 class CommGroup:
     def __init__(self, comm_handle: int, ranks: list[int]):
         self._comm = comm_handle
-        self.ranks = ranks
+        self._ranks = ranks
 
     @property
     def handle(self) -> int:
         return self._comm
 
     @property
     def ranks(self) -> list[int]:
-        return self.ranks
+        return self._ranks
 
 
 DistributedProcessGroup = torch_dist.ProcessGroup | CommGroup
@@ -39,6 +39,7 @@ def init_process_group(
         rank: int,
         world_size: int,
         timeout: timedelta,
+        **kwargs,
     ):
         raise NotImplementedError
 
@@ -100,22 +101,21 @@ def new_group(
 
 
 class TorchBackend(Distributed):
-    def __init__(self, backend_type: str):
-        self.backend_type = backend_type
-
     def init_process_group(
         self,
         host: str,
         port: int,
         rank: int,
         world_size: int,
         timeout: timedelta,
+        **kwargs,
     ):
+        backend = kwargs.get("backend", "nccl")
         store = torch.distributed.TCPStore(
             host, port, world_size, timeout=timeout, is_master=(rank == 0)
         )
         torch.distributed.init_process_group(
-            backend=self.backend_type,
+            backend=backend,
             world_size=world_size,
             rank=rank,
             timeout=timeout,
@@ -159,7 +159,7 @@ def new_group(self, ranks: list[int], **kwargs) -> DistributedProcessGroup | Non
 
 
 # specific device instance
-_BACKEND_INSTANCE: Distributed = TorchBackend(backend_type="nccl")
+_BACKEND_INSTANCE: Distributed = TorchBackend()
 
 _pickler = pickle.Pickler
 _unpickler = pickle.Unpickler
@@ -223,33 +223,34 @@ def _common_all_gather_object(
         object_list[i] = _tensor_to_object(tensor, tensor_size)
 
 
+def use_backend(backend: str | None):
+    global _BACKEND_INSTANCE
+
+    if not backend:
+        return
+
+    mapping = {
+        "vllm_nccl": ".nccl.DistributedNccl",
+        "vllm_hccl": ".hccl.DistributedHccl",
+    }
+    if backend not in mapping:
+        raise ValueError(f"Unsupported custom backend: {backend}")
+
+    module_path, class_name = mapping[backend].rsplit(".", 1)
+    module = importlib.import_module(module_path, "checkpoint_engine.distributed")
+    backend_class = getattr(module, class_name)
+    _BACKEND_INSTANCE = backend_class()
+
+
 def init_process_group(
     host: str,
     port: int,
     rank: int,
     world_size: int,
-    custom_dist: bool,
-    backend: str,
     timeout: timedelta = timedelta(seconds=300),
+    **kwargs,
 ):
-    global _BACKEND_INSTANCE
-
-    if not custom_dist:
-        _BACKEND_INSTANCE = TorchBackend(backend_type=backend)
-    else:
-        mapping = {
-            "nccl": ".nccl.DistributedNccl",
-            "hccl": ".hccl.DistributedHccl",
-        }
-        if backend not in mapping:
-            raise ValueError(f"Unsupported custom backend: {backend}")
-
-        module_path, class_name = mapping[backend].rsplit(".", 1)
-        module = importlib.import_module(module_path, "checkpoint_engine.distributed")
-        backend_class = getattr(module, class_name)
-        _BACKEND_INSTANCE = backend_class()
-
-    _BACKEND_INSTANCE.init_process_group(host, port, rank, world_size, timeout)
+    _BACKEND_INSTANCE.init_process_group(host, port, rank, world_size, timeout, **kwargs)
 
 
 def destroy_process_group(group: DistributedProcessGroup | None = None):

checkpoint_engine/distributed/hccl.py

@@ -208,25 +208,25 @@ def __init__(self):
 
     @contextmanager
     def _use_group(self, group: CommGroup | None, src: int | None = None):
+        active_src = src
         if group:
-            assert group.handle() in self.sub_groups, "invalid sub_group"
-            newcomm = ctypes.c_void_p(group.handle())
-            self.pynccl.comm = newcomm
-            active_src = src
+            assert group.handle in self.sub_groups, "invalid sub_group"
+            newcomm = ctypes.c_void_p(group.handle)
+            self.pyhccl.comm = newcomm
 
             if src is not None:
-                assert src in group.ranks(), "src rank not in group"
+                assert src in group.ranks, "src rank not in group"
                 # convert src rank id in default world to newcomm
-                active_src = group.ranks().index(src)
-                self.pynccl.rank = group.ranks().index(self.rank)
+                active_src = group.ranks.index(src)
+                self.pyhccl.rank = group.ranks.index(self.rank)
 
         try:
             yield active_src
         finally:
             if group:
-                self.pynccl.comm = self.comm
+                self.pyhccl.comm = self.comm
                 if src is not None:
-                    self.pynccl.rank = self.rank
+                    self.pyhccl.rank = self.rank
 
     def init_process_group(
         self,
@@ -235,6 +235,7 @@ def init_process_group(
         rank: int,
         world_size: int,
         timeout: timedelta = timedelta(seconds=300),
+        **kwargs,
     ):
         assert not self.initialized, "already initialized"
 
@@ -257,10 +258,10 @@ def destroy_process_group(
     ):
         assert self.initialized, "not initialized"
 
-        if group in self.sub_groups:
-            subcomm = ctypes.c_void_p(group)
+        if group and group.handle in self.sub_groups:
+            subcomm = ctypes.c_void_p(group.handle)
             self.pyhccl.destroy_comm(subcomm)
-            del self.sub_groups[group]
+            del self.sub_groups[group.handle]
             return
 
         self.pyhccl.destroy_comm()
@@ -297,7 +298,7 @@ def broadcast(
     ):
         assert self.initialized, "not initialized"
 
-        with self._use_group(group) as local_rank:
+        with self._use_group(group, src) as local_rank:
             self.pyhccl.broadcast(tensor, local_rank)
             current_stream().synchronize()
 
@@ -318,8 +319,11 @@ def new_group(self, ranks: list[int], **kwargs) -> CommGroup:
         else:
             ranks.sort()
 
-        newcomm = self.pynccl.create_newcomm(ranks)
-        if newcomm:
-            group = CommGroup(newcomm.value, ranks)
-            self.sub_groups[newcomm.value] = group
+        if self.rank not in ranks:
+            return
+
+        subcomm = self.pyhccl.create_subcomm(ranks)
+        if subcomm:
+            group = CommGroup(subcomm.value, ranks)
+            self.sub_groups[subcomm.value] = group
         return group

checkpoint_engine/distributed/nccl.py

@@ -111,17 +111,17 @@ def __init__(self):
 
     @contextmanager
     def _use_group(self, group: CommGroup | None, src: int | None = None):
+        active_src = src
         if group:
-            assert group.handle() in self.sub_groups, "invalid sub_group"
-            newcomm = ctypes.c_void_p(group.handle())
+            assert group.handle in self.sub_groups, "invalid sub_group"
+            newcomm = ctypes.c_void_p(group.handle)
             self.pynccl.comm = newcomm
-            active_src = src
 
             if src is not None:
-                assert src in group.ranks(), "src rank not in group"
+                assert src in group.ranks, "src rank not in group"
                 # convert src rank id in default world to newcomm
-                active_src = group.ranks().index(src)
-                self.pynccl.rank = group.ranks().index(self.rank)
+                active_src = group.ranks.index(src)
+                self.pynccl.rank = group.ranks.index(self.rank)
 
         try:
             yield active_src
@@ -138,6 +138,7 @@ def init_process_group(
         rank: int,
         world_size: int,
         timeout: timedelta = timedelta(seconds=300),
+        **kwargs,
     ):
         assert not self.initialized, "already initialized"
 
@@ -161,10 +162,10 @@ def destroy_process_group(
     ):
         assert self.initialized, "not initialized"
 
-        if group.handle() in self.sub_groups:
-            newcomm = ctypes.c_void_p(group.handle())
+        if group and group.handle in self.sub_groups:
+            newcomm = ctypes.c_void_p(group.handle)
             self.pynccl.destroy_comm(newcomm)
-            del self.sub_groups[group.handle()]
+            del self.sub_groups[group.handle]
             return
 
         self.pynccl.destroy_comm()

checkpoint_engine/ps.py

@@ -176,7 +176,6 @@ def __init__(
         auto_pg: bool = True,
         gpu_count: int | None = None,
         mem_fraction: float | None = None,
-        custom_dist: bool = False,
     ):
         """
         Initialize the parameter server. env RANK, WORLD_SIZE and MASTER_ADDR must be set.
@@ -197,7 +196,6 @@ def __init__(
         self._local_rdma_devices: dict[str, set[int]] = defaultdict(set)
         self._remote_rdma_devices: dict[str, set[int]] = defaultdict(set)
         self._mem_fraction = mem_fraction or float(os.getenv("PS_MEM_FRACTION", "0.9"))
-        self._custom_dist = custom_dist
 
         assert self._rank is not None and self._rank >= 0, self._rank
         assert self._world_size and self._world_size > 0, self._world_size
@@ -498,9 +496,8 @@ def init_process_group(
             port=_get_master_port(master_port),
             rank=self._rank,
             world_size=self._world_size,
-            custom_dist=self._custom_dist,
-            backend=self.device_manager.backend,
             timeout=timeout,
+            backend=self.device_manager.backend,
         )
         logger.info(f"[rank{self._rank}] init process group successfully.")
 

examples/update.py

@@ -159,13 +159,14 @@ def join(
     parser.add_argument("--checkpoint-name", type=str, default="my-checkpoint-iter-0")
     parser.add_argument("--update-method", type=str, default="broadcast")
     parser.add_argument("--uds", type=str, default=None)
-    parser.add_argument("--custom-dist", action="store_true")
+    parser.add_argument("--custom-dist", type=str, default=None)
     args = parser.parse_args()
     rank = int(os.getenv("RANK"))
     world_size = int(os.getenv("WORLD_SIZE"))
 
     req_func = req_inference(args.endpoint, args.inference_parallel_size, args.uds)
-    ps = ParameterServer(auto_pg=True, custom_dist=args.custom_dist)
+    dist.use_backend(args.custom_dist)
+    ps = ParameterServer(auto_pg=True)
     if args.load_metas_file:
         join(
             ps,