完善rank信息的管理。

Author: hiworldwzj

lightllm/common/mem_manager.py

@@ -8,7 +8,7 @@
 from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
 from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory
 from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
-from lightllm.utils.dist_utils import get_global_rank
+from lightllm.utils.dist_utils import get_current_rank_in_node
 from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args
 
 
@@ -37,8 +37,8 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False
         # 用共享内存进行共享，router 模块读取进行精确的调度估计, nccl port 作为一个单机中单实列的标记。防止冲突。
         from lightllm.utils.envs_utils import get_unique_server_name
 
-        rank_id = get_global_rank()
-        self.shared_can_use_token_num = SharedInt(f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank_id}")
+        rank_in_node = get_current_rank_in_node()
+        self.shared_can_use_token_num = SharedInt(f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank_in_node}")
 
         self.shared_can_use_token_num.set_value(self.can_use_mem_size)
         self._init_buffers(
@@ -303,20 +303,16 @@ class ReadOnlyStaticsMemoryManager:
     def __init__(self) -> None:
         args = get_env_start_args()
         self.global_world_size = args.tp
-        node_world_size = args.tp // args.nnodes
-        rank_start = args.node_rank * node_world_size
-        rank_end = (args.node_rank + 1) * node_world_size
-        self.shared_tp_infos = {
-            rank: SharedInt(f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank}")
-            for rank in range(rank_start, rank_end)
-        }
-
-    def get_unrefed_token_num(self, dp_rank: int):
-        args = get_env_start_args()
-        if args.dp == 1 and args.nnodes > 1:
-            # 兼容多机 dp size=1 的情况
-            rank_id = args.tp // args.nnodes * args.node_rank
-            return self.shared_tp_infos[rank_id].get_value()
-        dp_size = args.dp
-        dp_world_size = self.global_world_size // dp_size
-        return self.shared_tp_infos[dp_rank * dp_world_size].get_value()
+        self.node_world_size = args.tp // args.nnodes
+        self.dp_world_size = self.global_world_size  // args.dp
+        # 兼容多机 dp size=1 纯 tp 模式的情况
+        self.is_multinode_tp = args.dp == 1 and args.nnodes > 1
+        self.shared_tp_infos = [
+            SharedInt(f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank_in_node}")
+            for rank_in_node in range(0, self.node_world_size, self.dp_world_size)
+        ]
+
+    def get_unrefed_token_num(self, dp_rank_in_node: int):
+        if self.is_multinode_tp:
+            return self.shared_tp_infos[0].get_value()
+        return self.shared_tp_infos[dp_rank_in_node].get_value()

lightllm/server/api_http.py

@@ -101,7 +101,8 @@ def set_args(self, args):
                 enable_multimodal=args.enable_multimodal,
                 metric_port=args.metric_port,
             )
-            self.shared_token_load = TokenLoad(f"{get_unique_server_name()}_shared_token_load", args.dp)
+            dp_size_in_node = max(1, args.dp // args.nnodes) # 兼容多机存粹tp的运行模式，这时候 1 // 2 == 0, 需要兼容
+            self.shared_token_load = TokenLoad(f"{get_unique_server_name()}_shared_token_load", dp_size_in_node)
 
 
 g_objs = G_Objs()

lightllm/server/router/dynamic_prompt/radix_cache.py

@@ -96,7 +96,7 @@ class RadixCache:
     unique_name 主要用于解决单机，多实列部署时的shm冲突
     """
 
-    def __init__(self, unique_name, total_token_num, tp_id, mem_manager: MemoryManager = None):
+    def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager: MemoryManager = None):
         self.mem_manager = mem_manager
         self._key_dtype = torch.int64
         self._value_dtype = torch.int64
@@ -109,9 +109,9 @@ def __init__(self, unique_name, total_token_num, tp_id, mem_manager: MemoryManag
         self.evict_tree_set: Set[TreeNode] = SortedSet(key=lambda x: x.get_compare_key())  # 自定义比较器
         self.evict_tree_set.add(self.root_node)
 
-        self.refed_tokens_num = SharedArray(f"{unique_name}_refed_tokens_num_{tp_id}", (1,), dtype=np.int64)
+        self.refed_tokens_num = SharedArray(f"{unique_name}_refed_tokens_num_{rank_in_node}", (1,), dtype=np.int64)
         self.refed_tokens_num.arr[0] = 0
-        self.tree_total_tokens_num = SharedArray(f"{unique_name}_tree_total_tokens_num_{tp_id}", (1,), dtype=np.int64)
+        self.tree_total_tokens_num = SharedArray(f"{unique_name}_tree_total_tokens_num_{rank_in_node}", (1,), dtype=np.int64)
         self.tree_total_tokens_num.arr[0] = 0
 
     def insert(self, key, value=None):
@@ -345,9 +345,9 @@ class _RadixCacheReadOnlyClient:
     router 端只读用的客户端，用于从共享内存中读取树结构中的信息，用于进行prompt cache 的调度估计。
     """
 
-    def __init__(self, unique_name, total_token_num, tp_id):
-        self.refed_tokens_num = SharedArray(f"{unique_name}_refed_tokens_num_{tp_id}", (1,), dtype=np.int64)
-        self.tree_total_tokens_num = SharedArray(f"{unique_name}_tree_total_tokens_num_{tp_id}", (1,), dtype=np.int64)
+    def __init__(self, unique_name, total_token_num, rank_in_node):
+        self.refed_tokens_num = SharedArray(f"{unique_name}_refed_tokens_num_{rank_in_node}", (1,), dtype=np.int64)
+        self.tree_total_tokens_num = SharedArray(f"{unique_name}_tree_total_tokens_num_{rank_in_node}", (1,), dtype=np.int64)
 
     def get_refed_tokens_num(self):
         return self.refed_tokens_num.arr[0]
@@ -360,115 +360,16 @@ def get_unrefed_tokens_num(self):
 
 
 class RadixCacheReadOnlyClient:
-    def __init__(self, unique_name, total_token_num, tp_size):
-        self.tp_clients: List[_RadixCacheReadOnlyClient] = [
-            _RadixCacheReadOnlyClient(unique_name, total_token_num, tp_id) for tp_id in range(tp_size)
+    def __init__(self, unique_name, total_token_num, node_world_size, dp_world_size):
+        self.dp_rank_clients: List[_RadixCacheReadOnlyClient] = [
+            _RadixCacheReadOnlyClient(unique_name, total_token_num, rank_in_node) for rank_in_node in range(0, node_world_size, dp_world_size)
         ]
 
-    def get_refed_tokens_num(self, index):
-        return self.tp_clients[index].get_refed_tokens_num()
+    def get_refed_tokens_num(self, dp_rank_in_node):
+        return self.dp_rank_clients[dp_rank_in_node].get_refed_tokens_num()
 
-    def get_tree_total_tokens_num(self, index):
-        return self.tp_clients[index].get_tree_total_tokens_num()
+    def get_tree_total_tokens_num(self, dp_rank_in_node):
+        return self.dp_rank_clients[dp_rank_in_node].get_tree_total_tokens_num()
 
-    def get_unrefed_tokens_num(self, index):
-        return self.tp_clients[index].get_unrefed_tokens_num()
-
-
-# ///////////////////////////////////////////////////////////////////////////////
-
-if __name__ == "__main__":
-    # test 1
-    def test1():
-        tree = RadixCache("unique_name", 100, 0)
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        assert ans == 0
-        tree.print_self()
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        assert ans == 5
-        tree.print_self()
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        assert ans == 8
-        tree.print_self()
-
-        assert tree.get_refed_tokens_num() == 0
-        assert tree.get_tree_total_tokens_num() == 13
-
-        # print("evict")
-        tree.evict(9, lambda x: x)
-        tree.print_self()
-        assert tree.get_refed_tokens_num() == 0 and tree.get_tree_total_tokens_num() == 0
-
-    test1()
-
-    # test 2
-    def test2():
-        tree = RadixCache("unique_name", 100, 1)
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        tree.print_self()
-
-        tree_node, size, values = tree.match_prefix(
-            torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device="cpu"), update_refs=False
-        )
-        assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5
-        tree_node, size, values = tree.match_prefix(
-            torch.tensor([0, 1, 2, 3, 4, 9], dtype=torch.int64, device="cpu"), update_refs=False
-        )
-        assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5
-        tree_node, size, values = tree.match_prefix(
-            torch.tensor([0, 1, 2, 3, 4, 7, 8], dtype=torch.int64, device="cpu"), update_refs=False
-        )
-        assert tree_node.node_prefix_total_len == 7 and size == 7 and len(values) == 7
-        tree_node, size, values = tree.match_prefix(
-            torch.tensor([0, 1, 2, 3, 4, 7, 9], dtype=torch.int64, device="cpu"), update_refs=False
-        )
-        assert tree_node.node_prefix_total_len == 6 and size == 6 and len(values) == 6
-        print(ans)
-        return
-
-    # test2()
-
-    # test 3
-    def test3():
-        tree = RadixCache("unique_name", 100, 2)
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        tree.print_self()
-
-        tree_node, size, values = tree.match_prefix(
-            torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device="cpu"), update_refs=True
-        )
-        assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5
-        assert tree.get_refed_tokens_num() == 5 and tree.get_tree_total_tokens_num() == 13
-
-        tree_node, size, values = tree.match_prefix(
-            torch.tensor([0, 1, 2, 3, 4, 7, 9], dtype=torch.int64, device="cpu"), update_refs=True
-        )
-        assert tree_node.node_prefix_total_len == 6 and size == 6 and len(values) == 6
-        assert tree.get_refed_tokens_num() == 6 and tree.get_tree_total_tokens_num() == 13
-
-        tree.print_self()
-        tree.evict(2, lambda x: x)
-        assert tree.get_refed_tokens_num() == 6 and tree.get_tree_total_tokens_num() == 8
-        tree.print_self()
-
-        tree.dec_node_ref_counter(tree_node)
-        tree.print_self()
-        print(ans)
-        return
-
-    test3()
-
-    def test4():
-
-        tree = RadixCache("unique_name", 100, 2)
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        ans = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu"))
-        tree.print_self()
-
-        tree.clear_tree_nodes()
-        print(ans)
-        return
-
-    test4()
+    def get_unrefed_tokens_num(self, dp_rank_in_node):
+        return self.dp_rank_clients[dp_rank_in_node].get_unrefed_tokens_num()

lightllm/server/router/manager.py

@@ -43,9 +43,15 @@ def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metr
         self.args = args
         self.model_weightdir = args.model_dir
         self.world_size = args.tp
+        self.node_world_size = self.world_size // args.nnodes
         self.nnodes = args.nnodes
         self.node_rank = args.node_rank
         self.dp_size = args.dp
+        # 兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容
+        self.dp_size_in_node = max(1, args.dp // self.nnodes)
+        self.is_multinode_tp = args.nnodes > 1 and args.dp == 1
+        # 判断是否是保守调度，保守调度不会发生暂停 req 的情况，但是有些场景可能影响吞吐
+        self.is_safe_schedule = args.router_token_ratio == 0.0
         self.load_way = args.load_way
         self.mode = args.mode
         self.max_total_token_num = args.max_total_token_num
@@ -56,8 +62,8 @@ def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metr
         self.radix_cache_client = None
 
         # 共享变量，用于存储router端调度分析得到的机器负载信息
-        self.shared_token_load = TokenLoad(f"{get_unique_server_name()}_shared_token_load", self.dp_size)
-        for dp_index in range(self.dp_size):
+        self.shared_token_load = TokenLoad(f"{get_unique_server_name()}_shared_token_load", self.dp_size_in_node)
+        for dp_index in range(self.dp_size_in_node):
             self.shared_token_load.set_estimated_peak_token_count(0, dp_index)
             self.shared_token_load.set_frozened_token_count(0, dp_index)
             self.shared_token_load.set_current_load(0.0, dp_index)
@@ -77,7 +83,7 @@ def __init__(self, args, router_port, detokenization_port, model_rpc_ports, metr
         self.send_to_detokenization.connect(f"{args.zmq_mode}127.0.0.1:{detokenization_port}")
         self.model_rpc_ports = model_rpc_ports
 
-        if args.nnodes > 1 and args.dp == 1:
+        if self.is_multinode_tp:
             self.mulitnode_group = dist.init_process_group(
                 backend="gloo",
                 init_method=f"tcp://{args.nccl_host}:{args.multinode_router_gloo_port}",
@@ -177,9 +183,9 @@ async def wait_to_model_ready(self):
             self.args.max_total_token_num = self.max_total_token_num
         if self.args.use_dynamic_prompt_cache:
             self.radix_cache_client = RadixCacheReadOnlyClient(
-                get_unique_server_name(), self.max_total_token_num, tp_size=self.world_size
+                get_unique_server_name(), self.max_total_token_num, node_world_size=self.node_world_size, dp_world_size=self.world_size // self.dp_size
             )
-        self.req_queue = build_req_queue(self.args, self, self.dp_size)
+        self.req_queue = build_req_queue(self.args, self, self.dp_size_in_node)
         logger.info(f"use req queue {self.req_queue.__class__.__name__}")
 
         if self.args.run_mode == "prefill":
@@ -223,7 +229,7 @@ async def loop_for_fwd(
             counter_count += 1
             if self.running_batch is not None:
                 if counter_count % 50 == 0:
-                    for dp_index in range(self.dp_size):
+                    for dp_index in range(self.dp_size_in_node):
                         token_ratio1 = self.get_used_tokens(dp_index) / self.max_total_token_num
                         token_ratio2 = (
                             self.max_total_token_num
@@ -244,7 +250,7 @@ async def loop_for_fwd(
                 self.metric_client.gauge_set(
                     "lightllm_batch_current_max_tokens",
                     int(
-                        sum(self.shared_token_load.get_dynamic_max_load(d_i) for d_i in range(self.dp_size))
+                        sum(self.shared_token_load.get_dynamic_max_load(d_i) for d_i in range(self.dp_size_in_node))
                         * self.max_total_token_num
                     ),
                 )
@@ -264,7 +270,7 @@ async def get_schedule_result(self, running_batch: Batch):
 
             def get_new_batch():
                 limit_router_queue_length = None
-                if self.nnodes > 1 and self.args.dp == 1:
+                if self.is_multinode_tp:
                     # 使用 all_reduce 获取最小值
                     limit_router_queue_length = len(self.req_queue.waiting_req_list)
                     limit_router_queue_length_tensor = torch.tensor(
@@ -381,7 +387,7 @@ def _can_decode(self, batch: Batch):
         # p d 分离模式下，目前只能使用保守调度，保证请求放入进行decode的时候
         # 显存token肯定是够用的。
         # deepseekv2 dp 模式下,采用保守调度，也肯定够用
-        if self.is_pd_run_mode or self.dp_size > 1:
+        if self.is_pd_run_mode or self.dp_size > 1 or self.is_safe_schedule:
             return True
 
         # 下面的判定条件，只在 dp 为 1 的情况下启用

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

@@ -38,13 +38,14 @@
 from lightllm.server.router.model_infer.infer_batch import InferReq, InferSamplingParams
 from lightllm.server.router.token_load import TokenLoad
 from lightllm.common.basemodel.infer_lock import g_infer_state_lock, InferStateLock
-from lightllm.utils.dist_utils import _init_distributed_env
+from lightllm.utils.dist_utils import init_distributed_env
 from lightllm.utils.envs_utils import get_unique_server_name
 from lightllm.server.core.objs import ShmReqManager
 from lightllm.server.router.model_infer.infer_batch import g_infer_context
 from lightllm.utils.dist_utils import get_global_rank, get_global_world_size, get_dp_size
-from lightllm.utils.dist_utils import get_dp_world_size, get_current_dp_rank, get_current_rank_in_dp
+from lightllm.utils.dist_utils import get_dp_world_size, get_global_dp_rank, get_current_rank_in_dp
 from lightllm.utils.dist_utils import get_current_device_id, get_current_rank_in_node, get_node_world_size
+from lightllm.utils.dist_utils import get_dp_rank_in_node
 import torch.distributed as dist
 
 
@@ -64,6 +65,8 @@ def init_model(self, kvargs):
         self.tp_rank = kvargs["rank_id"]
         self.world_size = kvargs["world_size"]
         self.dp_size = kvargs.get("dp_size", 1)
+        # dp_size_in_node 计算兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容
+        self.dp_size_in_node = max(1, self.dp_size // self.nnodes)
         self.load_way = kvargs["load_way"]
         self.mode = kvargs["mode"]
         self.enable_chunked_prefill = kvargs.get("enable_chunked_prefill", False)
@@ -85,10 +88,10 @@ def init_model(self, kvargs):
             assert self.dp_size == self.world_size, "Currently only self-sustaining dp_size == tp_size"
             os.environ["ENABLE_DP"] = "1"
 
-        _init_distributed_env(kvargs)
+        init_distributed_env(kvargs)
         self.init_rank_infos()
 
-        self.shared_token_load = TokenLoad(f"{get_unique_server_name()}_shared_token_load", self.dp_size)
+        self.shared_token_load = TokenLoad(f"{get_unique_server_name()}_shared_token_load", self.dp_size_in_node)
 
         from lightllm.distributed import custom_comm_ops
 
@@ -239,17 +242,17 @@ def decode(self):
         raise NotImplementedError()
 
     def pause_reqs(self, req_ids):
-        if self.dp_size != 1:
+        if self.dp_size_in_node != 1:
             req_ids = [req_id for req_id in req_ids if req_id in g_infer_context.requests_mapping]
 
         g_infer_context.pause_reqs(req_ids)
         return
 
     # 一些可以复用的单元功能函数
     def _init_reqs(self, reqs: List[Tuple], init_req_obj=True):
-        if self.dp_size != 1:
-            cur_dp_index = self.tp_rank
-            reqs = [req for req in reqs if req[3] == cur_dp_index]
+        if self.dp_size_in_node != 1:
+            dp_rank_in_node = self.dp_rank_in_node
+            reqs = [req for req in reqs if req[3] == dp_rank_in_node]
 
         g_infer_state_lock.acquire()
         g_infer_context.add_reqs(reqs, init_req_obj=init_req_obj)
@@ -280,7 +283,8 @@ def init_rank_infos(self):
         self.rank_in_node = get_current_rank_in_node()
         self.current_device_id = get_current_device_id()
         self.rank_in_dp = get_current_rank_in_dp()
-        self.dp_rank = get_current_dp_rank()
+        self.global_dp_rank = get_global_dp_rank()
+        self.dp_rank_in_node = get_dp_rank_in_node()
         self.dp_world_size = get_dp_world_size()
         self.global_rank = get_global_rank()
         self.global_world_size = get_global_world_size()

lightllm/server/router/req_queue/__init__.py

@@ -5,7 +5,7 @@
 from .dp_base_queue import DpQueue
 
 
-def build_req_queue(args, router, dp_size: int):
+def build_req_queue(args, router, dp_size_in_node: int):
     queue_class = None
     if args.run_mode == "decode":
         queue_class = ContinuesBatchQueueForPDDecode
@@ -22,7 +22,7 @@ def build_req_queue(args, router, dp_size: int):
     if queue_class is None:
         queue_class = ContinuesBatchQueue
 
-    if dp_size == 1:
-        return queue_class(args, router, 0, dp_size)
+    if dp_size_in_node == 1:
+        return queue_class(args, router, 0, dp_size_in_node)
     else:
-        return DpQueue(args, router, queue_class, dp_size)
+        return DpQueue(args, router, queue_class, dp_size_in_node)