fix

Author: hiworldwzj

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_trans_obj.py

@@ -51,7 +51,7 @@ def create(
         self.pd_prefill_nccl_port = pd_prefill_nccl_port
 
         self.manager = manager
-        self.timer_checker = TimeChecker(3)
+        self.timer_checker = TimeChecker(6)
 
         with self.kv_trans_process.device_lock:
             self.kv_trans_process.task_in_queue.put(

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

@@ -4,6 +4,7 @@
 import inspect
 import torch.multiprocessing as mp
 from torch.distributed import TCPStore
+from datetime import timedelta
 from typing import List, Dict, Union
 from lightllm.utils.log_utils import init_logger
 from lightllm.common.mem_manager import MemoryManager
@@ -52,12 +53,17 @@ def _handle_prefill_join(
     node_info: PDTransJoinInfo, task_out_queue: mp.Queue, connect_id_to_comm: Dict[str, PyNcclCommunicator]
 ):
     try:
+        logger.info(f"connect start {node_info}")
         store_client = TCPStore(
-            host_name=node_info.pd_prefill_nccl_ip, port=node_info.pd_prefill_nccl_port, is_master=False, use_libuv=True
-        )
-        group = StatelessP2PProcessGroup.create(
-            src_id=node_info.prefill_id, dest_id=node_info.decode_id, is_server=False, store=store_client
+            host_name=node_info.pd_prefill_nccl_ip, port=node_info.pd_prefill_nccl_port, is_master=False, use_libuv=True, timeout=timedelta(seconds=30)
         )
+        src_id = node_info.prefill_id
+        dest_id = node_info.connect_id
+        logger.info(f"connect src_id {src_id} dest_id {dest_id}")
+        group = StatelessP2PProcessGroup.create(src_id=src_id,
+                                                dest_id=dest_id, 
+                                                is_server=False,
+                                                store=store_client)
         comm = PyNcclCommunicator(group, node_info.decode_device_id)
         connect_id_to_comm[node_info.connect_id] = comm
         logger.info(f"{node_info} kv trans connected")
@@ -68,6 +74,13 @@ def _handle_prefill_join(
 
 
 def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, mem_queues: List[mp.Queue]):
+    import os
+
+    # os.environ["NCCL_DEBUG"] = "INFO"
+    os.environ["NCCL_MAX_NCHANNELS"] = "2"
+    os.environ["NCCL_NSOCKS_PER_CHANNEL"] = "1"
+    os.environ["NCCL_SOCKET_NTHREADS"] = "1"
+    torch.backends.cudnn.enabled = False
 
     dp_size_in_node = max(1, args.dp // args.nnodes)
 

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_trans_obj.py

@@ -51,7 +51,7 @@ def create(
         self.prefill_node_id = prefill_node_id
         self.device_index = device_index
         self.manager = manager
-        self.timer_checker = TimeChecker(3)
+        self.timer_checker = TimeChecker(6)
 
         con = rpyc.connect(
             host=decode_node_ip, port=decode_node_rpyc_port, config={"allow_pickle": True}, keepalive=True

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_trans_process.py

@@ -4,6 +4,7 @@
 import inspect
 import torch.multiprocessing as mp
 from torch.distributed import TCPStore
+from datetime import timedelta
 from typing import List, Dict, Union
 from lightllm.utils.log_utils import init_logger
 from lightllm.common.mem_manager import MemoryManager
@@ -56,7 +57,14 @@ def _handle_decode_join(
     store: TCPStore,
 ):
     try:
-        group = StatelessP2PProcessGroup.create(node_info.prefill_id, node_info.decode_id, True, store)
+        logger.info(f"connect start {node_info}")
+        src_id = node_info.prefill_id
+        dest_id = node_info.connect_id
+        logger.info(f"connect src_id {src_id} dest_id {dest_id}")
+        group = StatelessP2PProcessGroup.create(src_id=src_id,
+                                                dest_id=dest_id, 
+                                                is_server=True,
+                                                store=store)
         comm = PyNcclCommunicator(group, node_info.prefill_device_id)
         connect_id_to_comm[node_info.connect_id] = comm
         logger.info(f"{node_info} kv trans connected!")
@@ -75,10 +83,18 @@ def _init_env(
     task_out_queue: mp.Queue,
     mem_queues: List[mp.Queue],
 ):
+    import os
+
+    # os.environ["NCCL_DEBUG"] = "INFO"
+    os.environ["NCCL_MAX_NCHANNELS"] = "2"
+    os.environ["NCCL_NSOCKS_PER_CHANNEL"] = "1"
+    os.environ["NCCL_SOCKET_NTHREADS"] = "1"
+    torch.backends.cudnn.enabled = False
+
     try:
         torch.cuda.set_device(device_id)
         graceful_registry(inspect.currentframe().f_code.co_name)
-        master_store = TCPStore(host_name=store_ip, port=store_port, is_master=True, use_libuv=True)
+        master_store = TCPStore(host_name=store_ip, port=store_port, is_master=True, use_libuv=True, timeout=timedelta(seconds=30))
         dp_size_in_node = max(1, args.dp // args.nnodes)
         task_out_queue.put("proc_start")
         mem_managers: List[MemoryManager] = [mem_queue.get(timeout=60) for mem_queue in mem_queues]