one kv trans process per tp.

Author: Weichao Luo

Diff for: lightllm/distributed/pynccl.py

@@ -248,51 +248,6 @@ def all_reduce(self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, strea
         )
         return out_tensor
 
-    def all_gather(self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None):
-        if self.disabled:
-            return
-        # nccl communicator created on a specific device
-        # will only work on tensors on the same device
-        # otherwise it will cause "illegal memory access"
-        assert input_tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {input_tensor.device}"
-        )
-        if stream is None:
-            stream = current_stream()
-        self.nccl.ncclAllGather(
-            buffer_type(input_tensor.data_ptr()),
-            buffer_type(output_tensor.data_ptr()),
-            input_tensor.numel(),
-            ncclDataTypeEnum.from_torch(input_tensor.dtype),
-            self.comm,
-            cudaStream_t(stream.cuda_stream),
-        )
-
-    def reduce_scatter(
-        self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None
-    ):
-        if self.disabled:
-            return
-        # nccl communicator created on a specific device
-        # will only work on tensors on the same device
-        # otherwise it will cause "illegal memory access"
-        assert input_tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {input_tensor.device}"
-        )
-        if stream is None:
-            stream = current_stream()
-        self.nccl.ncclReduceScatter(
-            buffer_type(input_tensor.data_ptr()),
-            buffer_type(output_tensor.data_ptr()),
-            output_tensor.numel(),
-            ncclDataTypeEnum.from_torch(input_tensor.dtype),
-            ncclRedOpTypeEnum.from_torch(op),
-            self.comm,
-            cudaStream_t(stream.cuda_stream),
-        )
-
     def send(self, tensor: torch.Tensor, dst: int, stream=None):
         if self.disabled:
             return
@@ -328,29 +283,3 @@ def recv(self, tensor: torch.Tensor, src: int, stream=None):
             self.comm,
             cudaStream_t(stream.cuda_stream),
         )
-
-    def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
-        if self.disabled:
-            return
-        assert tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {tensor.device}"
-        )
-        if stream is None:
-            stream = current_stream()
-        if src == self.rank:
-            sendbuff = buffer_type(tensor.data_ptr())
-            # NCCL requires the sender also to have a receive buffer
-            recvbuff = buffer_type(tensor.data_ptr())
-        else:
-            sendbuff = buffer_type()
-            recvbuff = buffer_type(tensor.data_ptr())
-        self.nccl.ncclBroadcast(
-            sendbuff,
-            recvbuff,
-            tensor.numel(),
-            ncclDataTypeEnum.from_torch(tensor.dtype),
-            src,
-            self.comm,
-            cudaStream_t(stream.cuda_stream),
-        )

Diff for: lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_kv_move_manager.py

@@ -51,8 +51,8 @@ def create(self, prefill_node_id: str, prefill_ip: str, prefill_port: int, manag
 
         device_index = manager.get_next_device_index()
         decode_node_id = manager.args.pd_node_id
-        task_in_queue = manager.kv_trans_task_in_queue
-        task_out_queue = manager.kv_trans_task_out_queue
+        task_in_queue = manager.kv_trans_task_in_queues[device_index]
+        task_out_queue = manager.kv_trans_task_out_queues[device_index]
 
         task_in_queue.put(
             PDTransJoinInfo(
@@ -136,7 +136,6 @@ def kv_move_loop(self):
                 self.manager.put_to_fail_release_task_queue(move_tasks)
 
         logger.error(f"{func_name} prefill id {self.prefill_node_id} device_index {self.device_index} thread quit")
-        self.task_in_queue.put(PDTransLeaveInfo(decode_id=self.decode_node_id, prefill_id=self.prefill_node_id))
         return
 
     def put_to_radix_loop(self):
@@ -217,6 +216,7 @@ def __del__(self):
         try:
             self.set_has_error()
             self.wait_thread_quit()
+            self.task_in_queue.put(PDTransLeaveInfo(decode_id=self.decode_node_id, prefill_id=self.prefill_node_id))
             if self.ready_to_move_queue is not None:
                 self.ready_to_move_queue.clear_tasks()
             if self.move_finished_queue is not None:
@@ -266,18 +266,31 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         # 需要每个卡有一个锁来规划每次只能有一个tran obj 操作对应显卡上的传输任务。
         self.device_locks = [threading.Lock() for _ in range(self.node_world_size)]
 
-        # start a single kv trans process
-        self.kv_trans_task_in_queue = mp.Queue()
-        self.kv_trans_task_out_queue = mp.Queue()
         from .decode_trans_process import start_decode_trans_process
 
-        self.kv_trans_process = start_decode_trans_process(
-            self.args, self.kv_trans_task_in_queue, self.kv_trans_task_out_queue, self.mem_queues
-        )
+        self.kv_trans_processes = []
+        self.kv_trans_task_in_queues = []
+        self.kv_trans_task_out_queues = []
+        self.kv_trans_process_alive = []
+
+        for device_index in range(self.node_world_size):
+            kv_trans_task_in_queue = mp.Queue()
+            kv_trans_task_out_queue = mp.Queue()
+            kv_trans_process = start_decode_trans_process(
+                self.args,
+                device_index,
+                kv_trans_task_in_queue,
+                kv_trans_task_out_queue,
+                self.mem_queues,
+            )
+            assert kv_trans_task_out_queue.get(timeout=30) == "proc_start"
+            self._put_mem_manager_to_mem_queue()
+            assert kv_trans_task_out_queue.get(timeout=60) == "get_mem_managers_ok"
 
-        assert self.kv_trans_task_out_queue.get(timeout=30) == "proc_start"
-        self._put_mem_manager_to_mem_queue()
-        assert self.kv_trans_task_out_queue.get(timeout=60) == "get_mem_managers_ok"
+            self.kv_trans_processes.append(kv_trans_process)
+            self.kv_trans_task_in_queues.append(kv_trans_task_in_queue)
+            self.kv_trans_task_out_queues.append(kv_trans_task_out_queue)
+            self.kv_trans_process_alive.append(True)
 
         return
 
@@ -462,7 +475,9 @@ def exposed_request_data_transfer(self, tasks: List[KVMoveTask]) -> List[Optiona
         return ans_list
 
     def get_next_device_index(self):
-        counts = [0 for _ in range(self.node_world_size)]
+        counts = [
+            0 if self.kv_trans_process_alive[device_id] else (1 << 20) for device_id in range(self.node_world_size)
+        ]
         for obj in self.node_id_to_trans_obj.values():
             counts[obj.device_index] += 1
         device_index = int(np.argmin(counts))
@@ -495,10 +510,22 @@ def remove_trans_obj(self, prefill_node_id):
         return
 
     def check_trans_process(self, raise_exception=True):
-        process = psutil.Process(self.kv_trans_process.pid)
-        if not (process.is_running() and process.status() != psutil.STATUS_ZOMBIE):
+        at_least_one_alive = False
+        for device_id in range(self.node_world_size):
+            if not self.kv_trans_process_alive[device_id]:
+                continue
+
+            process = psutil.Process(self.kv_trans_processes[device_id].pid)
+            if not (process.is_running() and process.status() != psutil.STATUS_ZOMBIE):
+                self.kv_trans_process_alive[device_id] = False
+                logger.error(f"kv trans process for device: {device_id} dead!!!")
+            else:
+                at_least_one_alive = True
+
+        if not at_least_one_alive:
             if raise_exception:
-                raise Exception(f"trans process: {self.kv_trans_process.pid} is dead")
+                raise Exception("All trans process are dead!!!")
+
         return
 
     def timer_loop(self):

Diff for: lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_trans_process.py

@@ -66,16 +66,17 @@ def _handle_prefill_join(
         logger.warning(f"error while connect to prefill node: {e}")
 
 
-def _init_env(args, task_in_queue: mp.Queue, task_out_queue: mp.Queue, mem_queues: List[mp.Queue]):
+def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, mem_queues: List[mp.Queue]):
 
     dp_size_in_node = max(1, args.dp // args.nnodes)
-    node_world_size = args.tp // args.nnodes
 
     try:
+        torch.cuda.set_device(device_id)
         graceful_registry(inspect.currentframe().f_code.co_name)
         task_out_queue.put("proc_start")
+
         mem_managers: List[MemoryManager] = [mem_queue.get(timeout=60) for mem_queue in mem_queues]
-        assert len(mem_managers) == node_world_size
+
         task_out_queue.put("get_mem_managers_ok")
         prefill_to_comm: Dict[int, PyNcclCommunicator] = {}
         while True:
@@ -97,12 +98,13 @@ def _init_env(args, task_in_queue: mp.Queue, task_out_queue: mp.Queue, mem_queue
 
 def start_decode_trans_process(
     args,
+    device_id: int,
     task_in_queue: mp.Queue,
     task_out_queue: mp.Queue,
     mem_queues: List[mp.Queue],
 ):
-    proc = mp.Process(target=_init_env, args=(args, task_in_queue, task_out_queue, mem_queues))
+    proc = mp.Process(target=_init_env, args=(args, device_id, task_in_queue, task_out_queue, mem_queues))
     proc.start()
     assert proc.is_alive()
-    logger.info("decode trans kv process start!")
+    logger.info(f"decode trans kv process for device: {device_id} start!")
     return proc

Diff for: lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_kv_move_manager.py

@@ -39,7 +39,7 @@ class TransProcessObj:
     rpyc_conn: object = None  # rpyc_con 的连接对象
     task_in_queue: mp.Queue = None
     task_out_queue: mp.Queue = None
-    device_index: str = None  # 使用的gpu序号
+    device_index: int = None  # 使用的gpu序号
     manager: "PrefillKVMoveManager" = None
     has_error: bool = False
     request_kv_trans_task_queue: TaskQueue = None
@@ -57,15 +57,15 @@ def create(
 
         device_index = manager.get_next_device_index()  # 分配 trans 进程使用的显卡
         prefill_node_id = manager.args.pd_node_id
-        task_in_queue = manager.kv_trans_task_in_queue
-        task_out_queue = manager.kv_trans_task_out_queue
+        task_in_queue = manager.kv_trans_task_in_queues[device_index]
+        task_out_queue = manager.kv_trans_task_out_queues[device_index]
 
         task_in_queue.put(
             PDTransJoinInfo(
                 prefill_id=prefill_node_id,
                 prefill_device_id=device_index,
                 prefill_ip=manager.host_ip,
-                prefill_port=manager.kv_trans_port,
+                prefill_port=manager.kv_trans_ports[device_index],
                 decode_id=decode_node_id,
                 decode_device_id=-1,
             )
@@ -74,7 +74,7 @@ def create(
         # 异步调用, 让decode节点建立与prefill节点进行nccl通信的进程
         max_kv_trans_token_num = obtain(
             con.root.build_trans_process(
-                prefill_node_id, manager.host_ip, manager.kv_trans_port, manager.args.max_total_token_num
+                prefill_node_id, manager.host_ip, manager.kv_trans_ports[device_index], manager.args.max_total_token_num
             )
         )
         self.max_kv_trans_token_num = max_kv_trans_token_num
@@ -237,7 +237,6 @@ def kv_trans_handle_loop(self):
                 self.manager.put_to_release_task_queue(move_tasks)
 
         logger.error(f"trans kv thread, decode id {self.decode_node_id} device_index {self.device_index} thread quit")
-        self.task_in_queue.put(PDTransLeaveInfo(decode_id=self.decode_node_id, prefill_id=self.prefill_node_id))
         return
 
     def wait_thread_quit(self):
@@ -282,6 +281,7 @@ def __del__(self):
         try:
             self.set_has_error()
             self.wait_thread_quit()
+            self.task_in_queue.put(PDTransLeaveInfo(decode_id=self.decode_node_id, prefill_id=self.prefill_node_id))
             if self.request_kv_trans_task_queue is not None:
                 self.request_kv_trans_task_queue.clear_tasks()
             if self.ready_kv_trans_task_queue is not None:
@@ -329,24 +329,37 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         self.release_tasks_thread.start()
 
         # start a single kv trans process
-        self.kv_trans_task_in_queue = mp.Queue()
-        self.kv_trans_task_out_queue = mp.Queue()
-        from .prefill_trans_process import start_prefill_trans_process
-
-        self.kv_trans_port = find_available_port(self.args.pd_p_allowed_port_min, self.args.pd_p_allowed_port_max)
-        self.kv_trans_process = start_prefill_trans_process(
-            self.args,
-            self.host_ip,
-            self.kv_trans_port,
-            self.kv_trans_task_in_queue,
-            self.kv_trans_task_out_queue,
-            self.mem_queues,
-        )
 
-        assert self.kv_trans_task_out_queue.get(timeout=30) == "proc_start"
-        self._put_mem_manager_to_mem_queue()
-        assert self.kv_trans_task_out_queue.get(timeout=60) == "get_mem_managers_ok"
+        from .prefill_trans_process import start_prefill_trans_process
 
+        self.kv_trans_ports = []
+        self.kv_trans_processes = []
+        self.kv_trans_task_in_queues = []
+        self.kv_trans_task_out_queues = []
+        self.kv_trans_process_alive = []
+
+        for device_id in range(self.node_world_size):
+            kv_trans_task_in_queue = mp.Queue()
+            kv_trans_task_out_queue = mp.Queue()
+            kv_trans_port = find_available_port(self.args.pd_p_allowed_port_min, self.args.pd_p_allowed_port_max)
+            kv_trans_process = start_prefill_trans_process(
+                self.args,
+                self.host_ip,
+                kv_trans_port,
+                device_id,
+                kv_trans_task_in_queue,
+                kv_trans_task_out_queue,
+                self.mem_queues,
+            )
+            assert kv_trans_task_out_queue.get(timeout=30) == "proc_start"
+            self._put_mem_manager_to_mem_queue()
+            assert kv_trans_task_out_queue.get(timeout=60) == "get_mem_managers_ok"
+
+            self.kv_trans_ports.append(kv_trans_port)
+            self.kv_trans_processes.append(kv_trans_process)
+            self.kv_trans_task_in_queues.append(kv_trans_task_in_queue)
+            self.kv_trans_task_out_queues.append(kv_trans_task_out_queue)
+            self.kv_trans_process_alive.append(True)
         return
 
     def put_to_release_task_queue(self, task: Union[KVMoveTask, List[KVMoveTask]]):
@@ -368,14 +381,28 @@ def handle_release_task_loop(self):
         return
 
     def check_trans_process(self, raise_exception=True):
-        process = psutil.Process(self.kv_trans_process.pid)
-        if not (process.is_running() and process.status() != psutil.STATUS_ZOMBIE):
+        at_least_one_alive = False
+        for device_id in range(self.node_world_size):
+            if not self.kv_trans_process_alive[device_id]:
+                continue
+
+            process = psutil.Process(self.kv_trans_processes[device_id].pid)
+            if not (process.is_running() and process.status() != psutil.STATUS_ZOMBIE):
+                self.kv_trans_process_alive[device_id] = False
+                logger.error(f"kv trans process for device: {device_id} dead!!!")
+            else:
+                at_least_one_alive = True
+
+        if not at_least_one_alive:
             if raise_exception:
-                raise Exception(f"trans process: {self.kv_trans_process.pid} is dead")
+                raise Exception("All trans process are dead!!!")
+
         return
 
     def get_next_device_index(self):
-        counts = [0 for _ in range(self.node_world_size)]
+        counts = [
+            0 if self.kv_trans_process_alive[device_id] else (1 << 20) for device_id in range(self.node_world_size)
+        ]
         for obj in self.node_id_to_trans_obj.values():
             counts[obj.device_index] += 1
         device_index = int(np.argmin(counts))