add dist.use_backend

Author: kip-cxj

checkpoint_engine/distributed/base.py

@@ -223,32 +223,32 @@ 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),
 ):
-    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)
 
 

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,8 +496,6 @@ 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,
         )
         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,