add torch.distributed + rpc orchestrator

Signed-off-by: Superjomn <328693+Superjomn@users.noreply.github.com>

Author: Superjomn

tensorrt_llm/_torch/distributed/communicator.py

@@ -476,14 +476,22 @@ def _get_cluster_info(self):
         if self.cluster_info is not None:
             return self.cluster_info
 
-        if ray.is_initialized():
+        is_ray_initialized = False
+        try:
+            if ray.is_initialized():
+                is_ray_initialized = True
+        except Exception:
+            pass
+
+        if is_ray_initialized:
             node_ip = ray.util.get_node_ip_address()
+            gpu_index = [int(id) for id in ray.get_gpu_ids()]
+            assert len(gpu_index) == 1
+            gpu_id = gpu_index[0]
         else:
-            raise RuntimeError("Ray is not initialized")
-
-        gpu_index = [int(id) for id in ray.get_gpu_ids()]
-
-        assert len(gpu_index) == 1
+            import socket
+            node_ip = socket.gethostbyname(socket.gethostname())
+            gpu_id = torch.cuda.current_device()
 
         # Gather node ip
         node_list = [None] * torch.distributed.get_world_size()
@@ -492,7 +500,7 @@ def _get_cluster_info(self):
 
         # Gather gpu index
         gpu_list = [None] * torch.distributed.get_world_size()
-        torch.distributed.all_gather_object(gpu_list, gpu_index[0])
+        torch.distributed.all_gather_object(gpu_list, gpu_id)
 
         # Gather rank
         rank_list = [None] * torch.distributed.get_world_size()
@@ -639,8 +647,15 @@ def allreduce(self,
                   obj: int | float | torch.Tensor,
                   op=torch.distributed.ReduceOp.SUM):
         is_base_type = isinstance(obj, int) or isinstance(obj, float)
+        device = torch.device(
+            "cuda") if dist.get_backend() == "nccl" else torch.device("cpu")
+
         if is_base_type:
-            obj = torch.tensor(obj)
+            obj = torch.tensor(obj, device=device)
+        elif isinstance(obj, torch.Tensor):
+            # Ensure tensor is on the correct device
+            if obj.device != device:
+                obj = obj.to(device)
 
         dist.all_reduce(obj, op=op)
 

tensorrt_llm/executor/rpc_torch_dist_executor.py

@@ -0,0 +1,158 @@
+import multiprocessing
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+from tensorrt_llm._utils import get_free_port
+from tensorrt_llm.bindings import executor as tllm
+from tensorrt_llm.builder import Engine
+from tensorrt_llm.executor.executor import GenerationExecutor
+from tensorrt_llm.executor.postproc_worker import PostprocWorkerConfig
+from tensorrt_llm.executor.rpc_proxy_mixin import RpcExecutorMixin
+from tensorrt_llm.executor.rpc_torch_dist_worker import RpcTorchDistWorker
+from tensorrt_llm.llmapi.llm_args import BaseLlmArgs
+from tensorrt_llm.llmapi.tokenizer import TokenizerBase
+from tensorrt_llm.logger import logger
+from tensorrt_llm.sampling_params import BatchedLogitsProcessor
+
+
+class RpcTorchDistExecutor(RpcExecutorMixin, GenerationExecutor):
+    def __init__(
+        self,
+        worker_kwargs: Dict,
+        model_world_size: int,
+        postproc_worker_config: PostprocWorkerConfig,
+        is_llm_executor: bool,
+    ):
+        # Initialize GenerationExecutor
+        super().__init__(
+            num_postprocess_workers=postproc_worker_config.num_postprocess_workers,
+            postprocess_tokenizer_dir=postproc_worker_config.postprocess_tokenizer_dir,
+            is_llm_executor=is_llm_executor,
+        )
+
+        self.world_size = model_world_size
+        self.processes: List[multiprocessing.Process] = []
+
+        # Setup RPC
+        self.init_rpc_executor()
+
+        # Determine Master Addr/Port for torch.distributed
+        self.master_addr = "127.0.0.1"
+        self.master_port = str(get_free_port())
+
+        logger.info(
+            f"RpcTorchDistExecutor starting with {model_world_size} workers."
+            f"Master: {self.master_addr}:{self.master_port}"
+        )
+
+        # Spawn workers
+        self.start_workers(worker_kwargs)
+
+        # Setup engine (remote)
+        # This will trigger setup_engine on rank 0 via RPC, which broadcasts to other ranks
+        try:
+            logger.info("Setting up remote engine...")
+            self.setup_engine_remote()
+        except Exception as e:
+            logger.error(f"Failed to setup remote engine: {e}")
+            self.shutdown()
+            raise e
+
+        # Setup main loop for receiving results from RPC
+        self.setup_mainloop()
+
+    def start_workers(self, worker_kwargs: Dict):
+        ctx = multiprocessing.get_context("spawn")
+
+        for rank in range(self.world_size):
+            p = ctx.Process(
+                target=RpcTorchDistWorker.worker_main,
+                args=(
+                    rank,
+                    self.world_size,
+                    self.master_addr,
+                    self.master_port,
+                    self.rpc_addr,  # Passed to all, but only used by rank 0
+                    worker_kwargs,
+                ),
+                name=f"RpcTorchDistWorker-{rank}",
+            )
+            p.start()
+            self.processes.append(p)
+
+    def setup_engine_remote(self):
+        # Call setup_engine on Rank 0 via RPC
+        # We wait for the result to ensure everything is initialized
+        self.rpc_client.setup_engine().remote()
+
+    def shutdown(self):
+        if self.doing_shutdown:
+            return
+        self.doing_shutdown = True
+
+        logger.info("Shutting down RpcTorchDistExecutor...")
+
+        # RPC shutdown to Rank 0
+        try:
+            if hasattr(self, "rpc_client") and self.rpc_client:
+                # This tells Rank 0 to shutdown, which broadcasts shutdown to others
+                self.rpc_client.shutdown().remote(need_response=False)
+        except Exception as e:
+            logger.warning(f"Error during RPC shutdown: {e}")
+
+        # Cleanup RPC client
+        if hasattr(self, "rpc_client") and self.rpc_client:
+            self.rpc_client.close()
+
+        # Join processes
+        for p in self.processes:
+            if p.is_alive():
+                p.join(timeout=5)
+                if p.is_alive():
+                    logger.warning(f"Process {p.name} did not exit, terminating...")
+                    p.terminate()
+
+        super().shutdown()
+
+    @classmethod
+    def create(
+        cls,
+        engine: Union[Path, Engine],
+        executor_config: Optional[tllm.ExecutorConfig] = None,
+        batched_logits_processor: Optional[BatchedLogitsProcessor] = None,
+        model_world_size: int = 1,
+        mpi_session: Optional[Any] = None,
+        reuse_mpi_comm: bool = False,
+        return_logits: bool = False,
+        postproc_worker_config: Optional[PostprocWorkerConfig] = None,
+        is_llm_executor: Optional[bool] = None,
+        hf_model_dir: Optional[Path] = None,
+        tokenizer: Optional[TokenizerBase] = None,
+        llm_args: Optional[BaseLlmArgs] = None,
+        **kwargs,
+    ):
+        postproc_worker_config = postproc_worker_config or PostprocWorkerConfig()
+
+        worker_kwargs = {
+            "engine": engine,
+            "executor_config": executor_config,
+            "batched_logits_processor": batched_logits_processor,
+            "hf_model_dir": hf_model_dir,
+            "tokenizer": tokenizer,
+            "llm_args": llm_args,
+        }
+
+        return cls(
+            worker_kwargs=worker_kwargs,
+            model_world_size=model_world_size,
+            postproc_worker_config=postproc_worker_config,
+            is_llm_executor=is_llm_executor or False,
+        )
+
+    # Implement abstract methods from GenerationExecutor
+    def submit(self, request):
+        return super().submit(request)  # RpcExecutorMixin.submit
+
+    def abort_request(self, request_id: int):
+        # Forward to Rank 0
+        self.rpc_client.abort_request(request_id).remote(need_response=False)

tensorrt_llm/executor/rpc_torch_dist_worker.py

@@ -0,0 +1,149 @@
+import os
+from typing import Any, Dict, Optional
+
+import torch
+import torch.distributed as dist
+
+from tensorrt_llm.executor.base_worker import BaseWorker
+from tensorrt_llm.executor.rpc_worker_mixin import RpcWorkerMixin
+from tensorrt_llm.logger import logger
+
+
+class RpcTorchDistWorker(RpcWorkerMixin, BaseWorker):
+    def __init__(
+        self, rank: int, world_size: int, device_id: int, rpc_addr: Optional[str] = None, **kwargs
+    ):
+        # Initialize BaseWorker
+        super().__init__(**kwargs)
+
+        self.rank = rank
+        self.global_rank = rank
+        self.world_size = world_size
+        self.device_id = device_id
+
+        # Create control group for worker orchestration
+        # Use Gloo for control messages as it doesn't require GPU
+        # and is robust.
+        self.control_group = dist.new_group(backend="gloo")
+
+        if self.rank == 0:
+            if rpc_addr is None:
+                raise ValueError("rpc_addr must be provided for rank 0")
+            self.init_rpc_worker(self.rank, rpc_addr)
+            self.start_rpc_server()
+
+    def setup_engine(self):
+        # Broadcast command if rank 0
+        if self.rank == 0:
+            self._broadcast_command("setup_engine")
+
+        # Ensure we are synchronized before setting up engine if needed
+        if dist.is_initialized():
+            dist.barrier()
+
+        super().setup_engine()
+
+    def start(self):
+        pass
+
+    def shutdown(self):
+        if self.doing_shutdown:
+            return
+
+        # Broadcast command if rank 0
+        if self.rank == 0:
+            try:
+                self._broadcast_command("shutdown")
+            except Exception as e:
+                logger.warning(f"Failed to broadcast shutdown command: {e}")
+
+        super().shutdown()
+
+        if self.rank == 0 and hasattr(self, "rpc_server") and self.rpc_server:
+            self.rpc_server.shutdown()
+
+    def _broadcast_command(self, command: str, args: Any = None):
+        if not dist.is_initialized():
+            return
+        cmd_list = [command, args]
+        try:
+            dist.broadcast_object_list(cmd_list, src=0, group=self.control_group)
+        except Exception as e:
+            logger.error(f"Broadcast error: {e}")
+
+    @classmethod
+    def worker_main(
+        cls,
+        rank: int,
+        world_size: int,
+        master_addr: str,
+        master_port: str,
+        rpc_addr: Optional[str],
+        worker_kwargs: Dict,
+    ):
+        # Setup environment
+        os.environ["MASTER_ADDR"] = master_addr
+        os.environ["MASTER_PORT"] = master_port
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+        os.environ["TLLM_DISABLE_MPI"] = "1"
+
+        # Setup device
+        if torch.cuda.is_available():
+            device_id = rank % torch.cuda.device_count()
+            torch.cuda.set_device(device_id)
+        else:
+            device_id = 0
+
+        # Initialize process group
+        # Use nccl for GPU, gloo for CPU
+        backend = "nccl" if torch.cuda.is_available() else "gloo"
+        dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
+
+        logger.info(f"Worker {rank}/{world_size} initialized with backend {backend}")
+
+        try:
+            worker = cls(
+                rank=rank,
+                world_size=world_size,
+                device_id=device_id,
+                rpc_addr=rpc_addr,
+                **worker_kwargs,
+            )
+
+            if rank == 0:
+                # Rank 0 waits for RPCs.
+                # The RPC server runs in a background thread started by start_rpc_server.
+                # We wait on the shutdown event which is set by shutdown() method (called via RPC).
+                worker.shutdown_event.wait()
+            else:
+                # Rank > 0 command loop
+                while True:
+                    cmd_list = [None, None]
+                    try:
+                        dist.broadcast_object_list(cmd_list, src=0, group=worker.control_group)
+                    except Exception as e:
+                        # If broadcast fails (e.g. rank 0 died), we should exit
+                        logger.error(f"Rank {rank} broadcast receive error: {e}")
+                        break
+
+                    cmd, args = cmd_list
+                    # logger.debug(f"Rank {rank} received command: {cmd}")
+
+                    if cmd == "setup_engine":
+                        worker.setup_engine()
+                    elif cmd == "shutdown":
+                        worker.shutdown()
+                        break
+                    elif cmd == "report_device_id":
+                        # Optional: handle other commands if needed
+                        pass
+                    else:
+                        logger.warning(f"Rank {rank} received unknown command: {cmd}")
+
+        except Exception as e:
+            logger.error(f"Worker {rank} failed with error: {e}")
+            raise e
+        finally:
+            if dist.is_initialized():
+                dist.destroy_process_group()

tensorrt_llm/llmapi/llm_args.py

@@ -1852,12 +1852,13 @@ class BaseLlmArgs(StrictBaseModel):
                                       description="Return perf metrics.",
                                       status="prototype")
 
-    orchestrator_type: Optional[Literal["rpc", "ray"]] = Field(
-        default=None,
-        description=
-        "The orchestrator type to use. Defaults to None, which uses MPI.",
-        status="prototype",
-    )
+    orchestrator_type: Optional[Literal[
+        "rpc", "ray", "rpc_torch_dist"]] = Field(
+            default=None,
+            description=
+            "The orchestrator type to use. Defaults to None, which uses MPI.",
+            status="prototype",
+        )
 
     _parallel_config: Optional[_ParallelConfig] = PrivateAttr(default=None)
     _model_format: Optional[_ModelFormatKind] = PrivateAttr(default=None)