Dp balancer (#991)

Author: shihaobai

lightllm/server/api_cli.py

@@ -144,6 +144,13 @@ def make_argument_parser() -> argparse.ArgumentParser:
                         using the deepseekv2 model, set dp to be equal to the tp parameter. In other cases, please
                         do not set it and keep the default value as 1.""",
     )
+    parser.add_argument(
+        "--dp_balancer",
+        type=str,
+        default="bs_balancer",
+        choices=["round_robin", "bs_balancer"],
+        help="the dp balancer type, default is bs_balancer",
+    )
     parser.add_argument(
         "--max_req_total_len", type=int, default=16384, help="the max value for req_input_len + req_output_len"
     )

lightllm/server/router/batch.py

@@ -40,6 +40,15 @@ def get_req_list_for_dp(self, dp_index: int):
                 req_list.append(req)
         return req_list
 
+    def get_all_dp_req_num(self) -> List[int]:
+        if self.dp_size_in_node == 1:
+            return [len(self.reqs)]
+
+        all_dp_req_num = [0 for _ in range(self.dp_size_in_node)]
+        for req in self.reqs:
+            all_dp_req_num[req.sample_params.suggested_dp_index] += 1
+        return all_dp_req_num
+
     def filter_out_finished_req(self, shm_req_manager: ShmReqManager):
         unfinished_req_ids = []
         for req in self.reqs:

lightllm/server/router/manager.py

@@ -197,10 +197,6 @@ async def wait_to_model_ready(self):
         return
 
     def _get_schedule_time_interval(self):
-        if self.running_batch is None:
-            # 没有运行中的 batch 时，每 10ms 触发一次请求调度
-            return 0.01
-
         # dp 模式，为了更好的配平，需要更长的调度间隔，以便于能收到更多的请求
         return self.schedule_time_interval
 

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

@@ -75,6 +75,7 @@ def init_model(self, kvargs):
         self.chunked_prefill_size = self.args.chunked_prefill_size
         self.return_all_prompt_logprobs = self.args.return_all_prompt_logprobs
         self.use_dynamic_prompt_cache = not self.args.disable_dynamic_prompt_cache
+        self.batch_max_tokens = self.args.batch_max_tokens
         self.eos_id: List[int] = kvargs.get("eos_id", [2])
         self.disable_cudagraph = self.args.disable_cudagraph
         self.is_multinode_tp = self.args.nnodes > 1 and self.args.dp == 1
@@ -395,6 +396,7 @@ def _get_classed_reqs(
         # 请求，其逻辑是不适合的。
         pause_max_req_num = 2
         wait_pause_count = 0
+        prefill_tokens = 0
 
         # 因为会使用到 radix cache 和 mem_manager 的计数信息
         # 所以需要加锁保护。
@@ -443,7 +445,10 @@ def _get_classed_reqs(
                         wait_pause_count += 1
             else:
                 token_num = req_obj.prefill_need_token_num(is_chuncked_prefill=not self.disable_chunked_prefill)
+                if prefill_tokens + token_num > self.batch_max_tokens:
+                    continue
                 if token_num <= can_alloc_token_num:
+                    prefill_tokens += token_num
                     prefill_reqs.append(req_obj)
                     can_alloc_token_num -= token_num
                 else:

lightllm/server/router/req_queue/base_queue.py

@@ -26,11 +26,6 @@ def __init__(self, args, router, dp_index, dp_size_in_node) -> None:
         self.router_token_ratio = args.router_token_ratio  # ratio to determine whether the router is busy
         self.router_max_new_token_len = args.router_max_new_token_len
 
-    def append(self, req: Req):
-        req.sample_params.suggested_dp_index = self.dp_index
-        self.waiting_req_list.append(req)
-        return
-
     def extend(self, req_group: List[Req]):
         for req in req_group:
             req.sample_params.suggested_dp_index = self.dp_index

lightllm/server/router/req_queue/dp_balancer/__init__.py

@@ -0,0 +1,13 @@
+from .roundrobin import RoundRobinDpBalancer
+from typing import List
+from lightllm.server.router.req_queue.base_queue import BaseQueue
+from .bs import DpBsBalancer
+
+
+def get_dp_balancer(args, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+    if args.dp_balancer == "round_robin":
+        return RoundRobinDpBalancer(dp_size_in_node, inner_queues)
+    elif args.dp_balancer == "bs_balancer":
+        return DpBsBalancer(dp_size_in_node, inner_queues)
+    else:
+        raise ValueError(f"Invalid dp balancer: {args.dp_balancer}")

lightllm/server/router/req_queue/dp_balancer/base.py

@@ -0,0 +1,23 @@
+import random
+from abc import ABC, abstractmethod
+from typing import List, Union
+from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.batch import Batch, Req
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class DpBalancer(ABC):
+    """
+    DP负载均衡器基类
+    定义了负载均衡策略的接口，子类可以实现不同的负载均衡算法
+    """
+
+    def __init__(self, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+        self.dp_size_in_node = dp_size_in_node
+        self.inner_queues = inner_queues
+
+    @abstractmethod
+    def assign_reqs_to_dp(self, current_batch: Batch, reqs_waiting_for_dp_index: List[List[Req]]) -> None:
+        pass

lightllm/server/router/req_queue/dp_balancer/bs.py

@@ -0,0 +1,45 @@
+import random
+from typing import List, Union
+from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.batch import Batch, Req
+from lightllm.utils.log_utils import init_logger
+from .base import DpBalancer
+
+logger = init_logger(__name__)
+
+
+class DpBsBalancer(DpBalancer):
+    """
+    This balancer is main to balance the batch size of each dp rank.
+    Because, for dp mode, if it exists a dp rank without any request, it will
+    padding a request and cause the waste of GPU compute resource.
+    """
+
+    def __init__(self, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+        super().__init__(dp_size_in_node, inner_queues)
+
+    def assign_reqs_to_dp(self, current_batch: Batch, reqs_waiting_for_dp_index: List[List[Req]]) -> None:
+        if len(reqs_waiting_for_dp_index) == 0:
+            return
+        # calculate the total load of each dp rank
+        all_dp_req_num = [0 for _ in range(self.dp_size_in_node)]
+        if current_batch is not None:
+            all_dp_req_num = current_batch.get_all_dp_req_num()
+        total_load_per_dp = [
+            all_dp_req_num[i] + len(self.inner_queues[i].waiting_req_list) for i in range(self.dp_size_in_node)
+        ]
+        for req_group in reqs_waiting_for_dp_index:
+            # find the dp rank with minimum load
+            min_load = min(total_load_per_dp)
+            select_dp_indexes = [i for i in range(self.dp_size_in_node) if total_load_per_dp[i] == min_load]
+            suggested_dp_index = random.choice(select_dp_indexes)
+
+            # assign the request to the dp rank and update the load count
+            for req in req_group:
+                req.sample_params.suggested_dp_index = suggested_dp_index
+            self.inner_queues[suggested_dp_index].extend(req_group)
+            # update the load count for this dp rank
+            total_load_per_dp[suggested_dp_index] += len(req_group)
+
+        reqs_waiting_for_dp_index.clear()
+        return

lightllm/server/router/req_queue/dp_balancer/roundrobin.py

@@ -0,0 +1,51 @@
+import random
+from typing import List, Union
+from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.batch import Batch, Req
+from lightllm.utils.log_utils import init_logger
+from .base import DpBalancer
+
+logger = init_logger(__name__)
+
+
+class RoundRobinDpBalancer(DpBalancer):
+    """
+    轮询负载均衡器
+    在队列长度最小的DP中进行轮询选择
+    """
+
+    def __init__(self, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+        super().__init__(dp_size_in_node, inner_queues)
+        self.pre_select_dp_index = self.dp_size_in_node - 1
+
+    def get_suggest_dp_index(
+        self,
+    ) -> int:
+        min_length = min(len(queue.waiting_req_list) for queue in self.inner_queues)
+        select_dp_indexes = [
+            i for i, queue in enumerate(self.inner_queues) if len(queue.waiting_req_list) == min_length
+        ]
+
+        # 如果没有可选择的索引，随机选择一个
+        if not select_dp_indexes:
+            self.pre_select_dp_index = random.randint(0, self.dp_size_in_node - 1)
+            return self.pre_select_dp_index
+
+        # 轮询选择
+        for i in range(self.dp_size_in_node):
+            next_dp_index = (self.pre_select_dp_index + i + 1) % self.dp_size_in_node
+            if next_dp_index in select_dp_indexes:
+                self.pre_select_dp_index = next_dp_index
+                return self.pre_select_dp_index
+
+        self.pre_select_dp_index = random.choice(select_dp_indexes)
+        return self.pre_select_dp_index
+
+    def assign_reqs_to_dp(self, current_batch: Batch, reqs_waiting_for_dp_index: List[List[Req]]) -> None:
+        for req_group in reqs_waiting_for_dp_index:
+            suggested_dp_index = self.get_suggest_dp_index()
+            for req in req_group:
+                req.sample_params.suggested_dp_index = suggested_dp_index
+            self.inner_queues[suggested_dp_index].extend(req_group)
+        reqs_waiting_for_dp_index.clear()
+        return

lightllm/server/router/req_queue/dp_base_queue.py

@@ -2,6 +2,7 @@
 from typing import List
 from ..batch import Batch, Req
 from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.req_queue.dp_balancer import get_dp_balancer
 from lightllm.common.basemodel.infer_lock import g_router_lock
 from lightllm.utils.log_utils import init_logger
 
@@ -12,14 +13,18 @@ class DpQueue:
     def __init__(self, args, router, base_queue_class, dp_size_in_node) -> None:
         self.dp_size_in_node = dp_size_in_node
         self.base_queue_class = base_queue_class
-        self.pre_select_dp_index = self.dp_size_in_node - 1
         from lightllm.server.router.manager import RouterManager
 
         self.router: RouterManager = router
         self.inner_queues: List[BaseQueue] = [
             base_queue_class(args, router, dp_index, dp_size_in_node) for dp_index in range(self.dp_size_in_node)
         ]
-
+        # 在调度这放松，在推理时约束。
+        # 避免prefill 模式下的情况下，推理完成了，调度没及时获取信息，导致调度bs 过小
+        for queue in self.inner_queues:
+            queue.batch_max_tokens = int(args.batch_max_tokens * 2)
+        self.dp_balancer = get_dp_balancer(args, dp_size_in_node, self.inner_queues)
+        self.reqs_waiting_for_dp_index: List[List[Req]] = []
         return
 
     def get_dp_queue(self, dp_index: int):
@@ -31,6 +36,7 @@ def get_wait_req_num(self):
 
     # @calculate_time(show=True, min_cost_ms=10)
     def generate_new_batch(self, current_batch: Batch):
+        self.dp_balancer.assign_reqs_to_dp(current_batch, self.reqs_waiting_for_dp_index)
         batches = [
             self.inner_queues[dp_index].generate_new_batch(current_batch) for dp_index in range(self.dp_size_in_node)
         ]
@@ -45,31 +51,13 @@ def _merge_batch(self, dp_batches: List[Batch]):
                 merged_batch = iter_batch
         return merged_batch
 
-    def append(self, req: Req):
-        suggested_dp_index = req.sample_params.suggested_dp_index
+    def extend(self, req_group: List[Req]):
+        suggested_dp_index = req_group[0].sample_params.suggested_dp_index
         if suggested_dp_index >= self.dp_size_in_node or suggested_dp_index < 0:
-            logger.warning(f"input req {req.request_id} dp index {suggested_dp_index} is invalid")
-            suggested_dp_index = self._get_suggest_dp_index()
-            self.pre_select_dp_index = suggested_dp_index
-            req.sample_params.suggested_dp_index = suggested_dp_index
-            self.inner_queues[suggested_dp_index].append(req)
+            # 同一个组的，要分配在同一个 dp 上
+            self.reqs_waiting_for_dp_index.append(req_group)
         else:
-            self.inner_queues[suggested_dp_index].append(req)
-        return
-
-    def extend(self, req_group: List[Req]):
-        # 同一个组的，要分配在同一个 dp 上，效率最高
-        index = self._get_suggest_dp_index()
-        for req in req_group:
-            suggested_dp_index = req.sample_params.suggested_dp_index
-            if suggested_dp_index >= self.dp_size_in_node or suggested_dp_index < 0:
-                logger.warning(f"input req {req.request_id} dp index {suggested_dp_index} is invalid")
-                self.pre_select_dp_index = index
-                req.sample_params.suggested_dp_index = index
-                self.inner_queues[index].append(req)
-            else:
-                self.inner_queues[suggested_dp_index].append(req)
-
+            self.inner_queues[suggested_dp_index].extend(req_group)
         return
 
     def is_busy(self):
@@ -87,21 +75,3 @@ def update_token_load(self, current_batch: Batch, force_update=False):
                     self.router.shared_token_load.set_estimated_peak_token_count(estimated_peak_token_count, dp_index)
                     self.router.shared_token_load.set_dynamic_max_load(dynamic_max_load, dp_index)
         return
-
-    def _get_suggest_dp_index(self):
-        min_length = min(len(queue.waiting_req_list) for queue in self.inner_queues)
-        select_dp_indexes = [
-            i for i, queue in enumerate(self.inner_queues) if len(queue.waiting_req_list) == min_length
-        ]
-
-        # multi thread safe keep
-        if not select_dp_indexes:
-            return random.randint(0, self.dp_size_in_node - 1)
-
-        # round_robin select.
-        for i in range(self.dp_size_in_node):
-            next_dp_index = (self.pre_select_dp_index + i + 1) % self.dp_size_in_node
-            if next_dp_index in select_dp_indexes:
-                return next_dp_index
-
-        return random.choice(select_dp_indexes)