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

@@ -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

@@ -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"
 

checkpoint_engine/distributed/nccl.py

@@ -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"
 

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,