fix tp moe and improve dp router. (#897)

Author: hiworldwzj

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe_weight_tp.py

@@ -3,7 +3,7 @@
 import threading
 from typing import Optional, Tuple, List, Dict, Any
 from .base_weight import BaseWeight
-from lightllm.utils.dist_utils import get_global_rank, get_current_device_id
+from lightllm.utils.dist_utils import get_current_rank_in_dp, get_current_device_id
 from lightllm.common.quantization import Quantcfg
 
 
@@ -37,7 +37,7 @@ def __init__(
         self.n_routed_experts = n_routed_experts
         self.split_inter_size = split_inter_size
         self.data_type_ = data_type
-        self.tp_rank_ = get_global_rank()
+        self.tp_rank_ = get_current_rank_in_dp()
         self.experts_up_projs = [None] * self.n_routed_experts
         self.experts_gate_projs = [None] * self.n_routed_experts
         self.experts_up_proj_scales = [None] * self.n_routed_experts

lightllm/server/router/req_queue/dp_base_queue.py

@@ -12,7 +12,7 @@ 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.round_robin_dp_id = 0
+        self.pre_select_dp_index = self.dp_size_in_node - 1
         from lightllm.server.router.manager import RouterManager
 
         self.router: RouterManager = router
@@ -52,8 +52,8 @@ def append(self, req: Req):
         suggested_dp_index = req.sample_params.suggested_dp_index
         if suggested_dp_index >= self.dp_size_in_node or suggested_dp_index < 0:
             logger.error(f"input req {req.request_id} dp index {suggested_dp_index} has error")
-            suggested_dp_index = self.round_robin_dp_id
-            self.round_robin_dp_id = (self.round_robin_dp_id + 1) % self.dp_size_in_node
+            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)
         else:
@@ -62,13 +62,12 @@ def append(self, req: Req):
 
     def extend(self, req_group: List[Req]):
         # 同一个组的，要分配在同一个 dp 上，效率最高
-        index = self.round_robin_dp_id
-        self.round_robin_dp_id = (self.round_robin_dp_id + 1) % self.dp_size_in_node
+        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.error(f"input req {req.request_id} dp index {suggested_dp_index} has error")
-
+                self.pre_select_dp_index = index
                 req.sample_params.suggested_dp_index = index
                 self.inner_queues[index].append(req)
             else:
@@ -94,3 +93,21 @@ 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)