修复新的rank管理机制下的pd 分离实现。

Author: hiworldwzj

lightllm/common/basemodel/infer_lock.py

@@ -17,28 +17,30 @@
 
 
 class InferStateLock:
-    def __init__(self, name):
+    def __init__(self, name, rank_in_dp:int, dp_rank_in_node:int, dp_world_size:int):
         self.infer_lock = threading.Lock()
+        self.dp_rank_in_node = dp_rank_in_node
+        # sync_world_size 应该是 min(dp_world_size, node_world_size)
+        self.dp_world_size = dp_world_size
+        self.rank_in_dp = rank_in_dp
         # 默认开 128 tp 的空间, 现在应该没什么卡能开这么大的tp 吧
-        self.lock_tp_infos = SharedArray(f"{name}_lock_tp_infos", shape=(129,), dtype=np.int64)
+        self.lock_tp_infos = SharedArray(f"{name}_dp_rank_{str(self.dp_rank_in_node)}_lock_tp_infos", shape=(self.dp_world_size + 1,), dtype=np.int64)
         self.lock_tp_infos.arr[:] = 0
-        self.rank_id = dist.get_rank()
-        self.world_size = dist.get_world_size()
 
     def add_cur_mark(self):
-        self.lock_tp_infos.arr[self.rank_id] += 1
+        self.lock_tp_infos.arr[self.rank_in_dp] += 1
 
     def get_cur_mark(self):
-        return self.lock_tp_infos.arr[self.rank_id]
+        return self.lock_tp_infos.arr[self.rank_in_dp]
 
     def get_max_mark_in_group(self):
-        return np.max(self.lock_tp_infos.arr[0 : self.world_size])
+        return np.max(self.lock_tp_infos.arr[0 : self.dp_world_size])
 
     def judge_cur_mark_equal_max_mark_in_group(self):
         return self.get_cur_mark() == self.get_max_mark_in_group()
 
     def judge_mark_in_group_all_same(self):
-        marks = self.lock_tp_infos.arr[0 : self.world_size]
+        marks = self.lock_tp_infos.arr[0 : self.dp_world_size]
         return bool(np.all(marks == marks[0]))
 
     def acquire_lock_and_update_cur_mark(self):
@@ -49,11 +51,11 @@ def release_lock(self):
         self.infer_lock.release()
 
     def set_group_wait_mark(self):
-        if self.rank_id == 0:
+        if self.rank_in_dp == 0:
             self.lock_tp_infos.arr[-1] = 1
 
     def unset_group_wait_mark(self):
-        if self.rank_id == 0:
+        if self.rank_in_dp == 0:
             self.lock_tp_infos.arr[-1] = 0
 
     def get_group_wait_mark(self):
@@ -63,7 +65,7 @@ def get_group_wait_mark(self):
 @dataclass
 class G_Infer_Lock:
     obj: InferStateLock = None
-    dp_size: int = None
+    dp_world_size: int = None
 
     def acquire(self):
         if self.obj is not None:
@@ -86,9 +88,8 @@ def release(self):
 
 # 下面两个函数需要配对使用
 def acquire_lock_until_ready(nccl_group):
-    # 在 deepseekv2 的tp dp 混合运行模式下, 不需要多个推理进程间做协调同步
-    # 所以直接加锁，解锁即可
-    if g_infer_state_lock.dp_size != 1:
+    # 单卡一tp不用过度加锁
+    if g_infer_state_lock.dp_world_size == 1:
         g_infer_state_lock.obj.infer_lock.acquire()
         return
 
@@ -118,7 +119,7 @@ def release_acquired_lock():
 @dataclass
 class G_Router_Lock:
     """
-    保护pd分离模式下, 一些数据的操作。
+    保护pd分离模式下, 一些调度相关信息数据的操作。
     """
 
     obj = None  # 进程锁对象

lightllm/common/deepseek2_mem_manager.py

@@ -40,9 +40,9 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         return
 
     def send_to_decode_node(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["Deepseek2MemoryManager"], dp_size: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["Deepseek2MemoryManager"], dp_size_in_node: int
     ):
-        assert dp_size == 1
+        assert dp_size_in_node == 1
 
         # 先将数据发送到指定的一张卡上的buffer，再发送。
         move_token_indexes = []
@@ -66,9 +66,9 @@ def _get_kv_move_data(self, token_indexes: List[int], layer_index: int):
         return move_buffer
 
     def receive_from_prefill_node(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
     ):
-        assert dp_size == 1
+        assert dp_size_in_node == 1
 
         # 先将数据接受到指定的一张卡上的buffer，再复制到其他的卡上。
         move_token_indexes = []
@@ -97,11 +97,11 @@ def _write_kv_move_data(self, token_indexes: torch.Tensor, buffer_tensor: torch.
         self.kv_buffer[layer_index : layer_index + 1, token_indexes, :, :] = buffer_tensor
         return
 
-    def send_to_decode_node_p2p(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int):
+    def send_to_decode_node_p2p(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int):
         """
         使用 p2p triton kernel 进行数据复制和传输的实现方式。
         """
-        assert dp_size == 1
+        assert dp_size_in_node == 1
 
         move_token_indexes = []
         for task in move_tasks:
@@ -124,9 +124,9 @@ def _get_kv_move_data_p2p(self, token_indexes: torch.Tensor, layer_index: int, k
         return move_buffer
 
     def receive_from_prefill_node_p2p(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
     ):
-        assert dp_size == 1
+        assert dp_size_in_node == 1
 
         move_token_indexes = []
         for task in move_tasks:

lightllm/common/mem_manager.py

@@ -83,13 +83,8 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         self.kv_move_buf_indexes = torch.arange(0, max_req_total_len + 8, dtype=torch.int64, device="cuda")
         return
 
-    def send_to_decode_node(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int):
-        """
-        dp_size 是为 deepseekv2 类型，可以 dp 和 tp 混合模式运行的模型定制的参数，
-        普通tp模式下, dp_size 一定等于 1, dp_index 一定等于 0, 同时普通模式下, 这两个参数并不会
-        被真正使用
-        """
-        assert dp_size == 1
+    def send_to_decode_node(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int):
+        assert dp_size_in_node == 1
 
         # 先将数据发送到指定的一张卡上的buffer，再发送。
 
@@ -123,14 +118,9 @@ def _get_kv_move_data(self, token_indexes: List[int], layer_index: int):
         return move_buffer
 
     def receive_from_prefill_node(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
     ):
-        """
-        dp_size 是为 deepseekv2 类型，可以 dp 和 tp 混合模式运行的模型定制的参数，
-        普通tp模式下, dp_size 一定等于 1, 同时普通模式下, 这两个参数并不会
-        被真正使用
-        """
-        assert dp_size == 1
+        assert dp_size_in_node == 1
 
         # 先将数据接受到指定的一张卡上的buffer，再复制到其他的卡上。
 
@@ -160,11 +150,11 @@ def _write_kv_move_data(self, token_indexes: torch.Tensor, buffer_tensor: torch.
         self.kv_buffer[layer_index : layer_index + 1, token_indexes, :, :] = buffer_tensor
         return
 
-    def send_to_decode_node_p2p(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int):
+    def send_to_decode_node_p2p(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int):
         """
         使用 p2p triton kernel 进行数据复制和传输的实现方式。
         """
-        assert dp_size == 1
+        assert dp_size_in_node == 1
 
         # 先将数据发送到指定的一张卡上的buffer，再发送。
 
@@ -190,9 +180,9 @@ def _get_kv_move_data_p2p(self, token_indexes: torch.Tensor, layer_index: int, k
         return move_buffer
 
     def receive_from_prefill_node_p2p(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
     ):
-        assert dp_size == 1
+        assert dp_size_in_node == 1
 
         # 先将数据接受到指定的一张卡上的buffer，再复制到其他的卡上。
 

lightllm/server/api_start.py

@@ -164,14 +164,14 @@ def normal_or_p_d_start(args):
     # 捕获到端口设置冲突的问题
     ports_locker = PortLocker(already_uesd_ports)
     ports_locker.lock_port()
-
+    
+    node_world_size = args.tp // args.nnodes
     can_use_ports = alloc_can_use_network_port(
-        num=6 + args.tp + args.tp + args.visual_dp * args.visual_tp, used_nccl_ports=already_uesd_ports
+        num=6 + node_world_size + args.visual_dp * args.visual_tp, used_nccl_ports=already_uesd_ports
     )
     logger.info(f"alloced ports: {can_use_ports}")
     router_port, detokenization_port, detokenization_pub_port, visual_port, cache_port, metric_port = can_use_ports[0:6]
-    model_rpc_ports = can_use_ports[6 : 6 + args.tp]
-    can_use_ports = can_use_ports[6 + args.tp :]
+    can_use_ports = can_use_ports[6:]
 
     visual_model_tp_ports = []
     for _ in range(args.visual_dp):
@@ -188,7 +188,7 @@ def normal_or_p_d_start(args):
     args.metric_port = metric_port
 
     # 申请在 p d 分离模式下，会用的端口
-    args.pd_tp_infer_rpyc_ports = can_use_ports[0 : args.tp]
+    args.pd_node_infer_rpyc_ports = can_use_ports[0 : node_world_size]
     # p d 分离模式下用于标识节点的id
     args.pd_node_id = uuid.uuid4().int
     # p 节点用来建立torch kv 传输分布组的可用端口范围
@@ -231,7 +231,7 @@ def normal_or_p_d_start(args):
     process_manager.start_submodule_processes(
         start_funcs=[start_router_process, start_detokenization_process],
         start_args=[
-            (args, router_port, detokenization_port, model_rpc_ports, metric_port),
+            (args, router_port, detokenization_port, metric_port),
             (args, detokenization_port, detokenization_pub_port),
         ],
     )

lightllm/server/router/manager.py

@@ -39,7 +39,7 @@
 
 
 class RouterManager:
-    def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metric_port):
+    def __init__(self, args, router_port, detokenization_port, metric_port):
         self.args = args
         self.model_weightdir = args.model_dir
         self.world_size = args.tp
@@ -81,8 +81,7 @@ def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metr
 
         self.send_to_detokenization = context.socket(zmq.PUSH)
         self.send_to_detokenization.connect(f"{args.zmq_mode}127.0.0.1:{detokenization_port}")
-        self.model_rpc_ports = model_rpc_ports
-
+        
         if self.is_multinode_tp:
             self.mulitnode_group = dist.init_process_group(
                 backend="gloo",
@@ -173,7 +172,7 @@ async def wait_to_model_ready(self):
             "batch_max_tokens": self.args.batch_max_tokens,
             "quant_type": self.args.quant_type,
             "quant_cfg": self.args.quant_cfg,
-            "pd_rpyc_ports": self.args.pd_tp_infer_rpyc_ports,  # 非 pd 模式可以不设置
+            "pd_rpyc_ports": self.args.pd_node_infer_rpyc_ports,  # 非 pd 模式可以不设置
         }
 
         await self.model_rpc_client.init_model(kvargs=kvargs)
@@ -416,7 +415,7 @@ def clean_up(self):
         return
 
 
-def start_router_process(args, router_port, detokenization_port, model_rpc_ports, metric_port, pipe_writer):
+def start_router_process(args, router_port, detokenization_port, metric_port, pipe_writer):
     # 注册 graceful 退出的处理
     graceful_registry(inspect.currentframe().f_code.co_name)
     start_parent_check_thread()
@@ -426,7 +425,6 @@ def start_router_process(args, router_port, detokenization_port, model_rpc_ports
             args,
             router_port=router_port,
             detokenization_port=detokenization_port,
-            model_rpc_ports=model_rpc_ports,
             metric_port=metric_port,
         )
 

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -100,8 +100,8 @@ def init_model(self, kvargs):
 
         # 为 p d 分离模式添加的全局锁管理，用于做一些同步操作。 一定需要在
         # init_process_group 之后调用
-        g_infer_state_lock.obj = InferStateLock(name=get_unique_server_name())
-        g_infer_state_lock.dp_size = self.dp_size
+        g_infer_state_lock.obj = InferStateLock(name=get_unique_server_name(), rank_in_dp=self.rank_in_dp, dp_rank_in_node=self.dp_rank_in_node, dp_world_size=self.dp_world_size)
+        g_infer_state_lock.dp_world_size = self.dp_world_size
         self.infer_state_lock = g_infer_state_lock
         # 防止InferStateLock 中的全局共享信息被重复异常初始化,导致同步异常的问题。
         # 所以做一次barrier等待

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

@@ -18,6 +18,7 @@
 from lightllm.common.basemodel.infer_lock import g_infer_state_lock, g_router_lock
 from .decode_task_cache import g_success_kv_move_task_cache, KVMoveTask
 from lightllm.utils.device_utils import kv_trans_use_p2p
+from lightllm.utils.envs_utils import get_unique_server_name
 
 logger = init_logger(__name__)
 
@@ -29,10 +30,16 @@ def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
         self.mem_queue: mp.Queue = mem_queue
 
     def init_custom(self):
-        self.lock_nccl_group = dist.new_group(backend="gloo")
+        ranks = []
+        for i in range(self.dp_world_size):
+            ranks.append(i + self.global_dp_rank * self.dp_world_size)
+        
+        self.lock_nccl_group = dist.new_group(ranks=ranks, backend="gloo")
+        logger.info(f"lock_nccl_group ranks {self.lock_nccl_group.get_rank()}")
+        
         from .decode_infer_rpyc import PDDecodeInferRpcServer
 
-        socket_path = f"/tmp/decode_node_infer_rpyc_{self.pd_rpyc_ports[self.tp_rank]}"
+        socket_path = f"/tmp/{get_unique_server_name()}_decode_node_infer_rpyc_{self.pd_rpyc_ports[self.rank_in_node]}"
         if os.path.exists(socket_path):
             os.remove(socket_path)
 
@@ -141,8 +148,8 @@ def post_init(self, uninit_reqs: List[InferReq]):
 
         if self.is_master_in_dp:
             with g_router_lock.obj:
-                self.shared_token_load.add_frozened_token_count(-remove_count, self.tp_rank)
-                self.shared_token_load.add_estimated_peak_token_count(estimated_peak_token_count, self.tp_rank)
+                self.shared_token_load.add_frozened_token_count(-remove_count, self.dp_rank_in_node)
+                self.shared_token_load.add_estimated_peak_token_count(estimated_peak_token_count, self.dp_rank_in_node)
         return
 
     def filter_finished_reqs(self, finished_reqs: List[InferReq]):