fix enable dp.

hiworldwzj · web-flow · commit 93d878c9285b · 2025-03-28T12:12:00.000+08:00
diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py
@@ -22,7 +22,6 @@ def __init__(self, layer_num, network_config, mode):
         self.tp_o_head_num_ = -1
         self.head_dim_ = -1
         self.embed_dim_ = -1
-        self.enable_dp = os.getenv("ENABLE_DP", "0").upper() in ["ON", "TRUE", "1"]
         return
 
     def _att_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
diff --git a/lightllm/models/deepseek2/infer_struct.py b/lightllm/models/deepseek2/infer_struct.py
@@ -9,7 +9,6 @@ class Deepseek2InferStateInfo(LlamaInferStateInfo):
     def __init__(self):
         super().__init__()
         self.kv_starts = None
-        self.enable_dp = os.getenv("ENABLE_DP", "0").upper() in ["ON", "TRUE", "1"]
 
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
         super().init_some_extra_state(model, input_ids)
diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py
@@ -66,7 +66,6 @@ def __init__(self, layer_num, network_config, mode=[]):
                 self.softmax_scale = self.softmax_scale * mscale * mscale
         self.enable_cc_method = not os.getenv("DISABLE_CC_METHOD", "False").upper() in ["ON", "TRUE", "1"]
         super().__init__(layer_num, network_config, mode)
-        self.enable_dp = os.getenv("ENABLE_DP", "0").upper() in ["ON", "TRUE", "1"]
         self.num_heads = network_config["num_attention_heads"]
         self.num_kv_heads = network_config["num_key_value_heads"]
         return
@@ -78,9 +77,8 @@ def _bind_func(self):
 
     def _bind_ffn(self):
         if self.is_moe:
-            if self.enable_dp:
-                moe_mode = os.environ.get("MOE_MODE", "TP")
-                if moe_mode == "EP":
+            moe_mode = os.environ.get("MOE_MODE", "TP")
+            if moe_mode == "EP":
                     self._ffn = partial(Deepseek2TransformerLayerInfer._moe_ffn_edp, self)
             else:
                 self._ffn = partial(Deepseek2TransformerLayerInfer._moe_ffn, self)
diff --git a/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py
@@ -19,7 +19,6 @@
 
 class Deepseek2TransformerLayerWeight(TransformerLayerWeight):
     def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
-        self.enable_dp = os.getenv("ENABLE_DP", "0").upper() in ["ON", "TRUE", "1"]
         self.enable_cc_method = not os.getenv("DISABLE_CC_METHOD", "False").upper() in ["ON", "TRUE", "1"]
         super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
         return
diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py
@@ -22,7 +22,7 @@
 class FlashInferStateExtraInfo:
     def __init__(self, model):
         num_heads = model.config["num_attention_heads"]
-        self.tp_q_head_num = num_heads if enable_env_vars("ENABLE_DP") else num_heads // get_dp_world_size()
+        self.tp_q_head_num = num_heads // get_dp_world_size()
         self.qk_nope_head_dim = model.qk_nope_head_dim
         self.qk_rope_head_dim = model.qk_rope_head_dim
         self.kv_lora_rank = model.kv_lora_rank
diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py
@@ -22,7 +22,6 @@ def __init__(self, network_config, mode):
         self.eps_ = network_config["rms_norm_eps"]
         self.vocab_size_ = network_config["vocab_size"]
         self.embed_dim_ = network_config["n_embed"]
-        self.enable_dp = os.getenv("ENABLE_DP", "0").upper() in ["ON", "TRUE", "1"]
         return
 
     def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor:
diff --git a/lightllm/models/llama/layer_infer/pre_layer_infer.py b/lightllm/models/llama/layer_infer/pre_layer_infer.py
@@ -16,12 +16,8 @@ class LlamaPreLayerInfer(PreLayerInferTpl):
 
     def __init__(self, network_config, mode):
         super().__init__(network_config, mode)
-        self.enable_dp = os.getenv("ENABLE_DP", "0").upper() in ["ON", "TRUE", "1"]
-        if not self.enable_dp:
-            tp_vob_ids = np.linspace(0, network_config["vocab_size"], self.tp_world_size_ + 1, dtype=np.int64)
-            self.vob_start_id_, self.vob_end_id_ = int(tp_vob_ids[self.tp_rank_]), int(tp_vob_ids[self.tp_rank_ + 1])
-        else:
-            self.vob_start_id_, self.vob_end_id_ = 0, network_config["vocab_size"]
+        tp_vob_ids = np.linspace(0, network_config["vocab_size"], self.tp_world_size_ + 1, dtype=np.int64)
+        self.vob_start_id_, self.vob_end_id_ = int(tp_vob_ids[self.tp_rank_]), int(tp_vob_ids[self.tp_rank_ + 1])
 
         return
 
diff --git a/lightllm/models/llama/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/llama/layer_weights/pre_and_post_layer_weight.py
@@ -7,7 +7,6 @@
 class LlamaPreAndPostLayerWeight(PreAndPostLayerWeight):
     def __init__(self, data_type, network_config, mode):
         super().__init__(data_type, network_config, mode)
-        self.enable_dp = os.getenv("ENABLE_DP", "0").upper() in ["ON", "TRUE", "1"]
         return
 
     def load_hf_weights(self, weights):
@@ -16,18 +15,12 @@ def load_hf_weights(self, weights):
         split_start = split_indexes[self.tp_rank_]
         split_end = split_indexes[self.tp_rank_ + 1]
         if "model.embed_tokens.weight" in weights:
-            if self.enable_dp:
-                self.wte_weight_ = self._cuda(weights["model.embed_tokens.weight"])
-            else:
-                self.wte_weight_ = self._cuda(weights["model.embed_tokens.weight"][split_start:split_end, :])
+            self.wte_weight_ = self._cuda(weights["model.embed_tokens.weight"][split_start:split_end, :])
             tie_word_embeddings = self.network_config_.get("tie_word_embeddings", False)
             if tie_word_embeddings:
                 self.lm_head_weight_ = self.wte_weight_
         if "lm_head.weight" in weights:
-            if self.enable_dp:
-                self.lm_head_weight_ = self._cuda(weights["lm_head.weight"])
-            else:
-                self.lm_head_weight_ = self._cuda(weights["lm_head.weight"][split_start:split_end, :])
+            self.lm_head_weight_ = self._cuda(weights["lm_head.weight"][split_start:split_end, :])
         if "model.norm.weight" in weights:
             self.final_norm_weight_ = self._cuda(weights["model.norm.weight"])
 
diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py
@@ -87,6 +87,7 @@ def normal_or_p_d_start(args):
 
     # 这些模式不能同时设置。
     assert [
+        args.disable_chunked_prefill,
         args.diverse_mode,
         args.token_healing_mode,
         args.use_reward_model,
diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py
@@ -523,7 +523,7 @@ async def abort(self, group_req_id: int):
         return
 
     async def recycle_resource_loop(self):
-        # pre_time_mark = time.time()
+        pre_time_mark = time.time()
 
         while True:
 
@@ -547,14 +547,14 @@ async def recycle_resource_loop(self):
                 await self._release_multimodal_resources(req_status.group_req_objs.multimodal_params)
 
             # 先保留这个关键得日志，用于方便定位重构中的问题。
-            # if time.time() - pre_time_mark > 20:
-            #     pre_time_mark = time.time()
-            #     for req_status in self.req_id_to_out_inf.values():
-            #         logger.info(
-            #             f"left req id {req_status.group_req_objs.group_req_id}"
-            #             f"can release {req_status.group_req_objs.shm_req_objs[0].can_released_mark} "
-            #             f"refcount {req_status.group_req_objs.shm_req_objs[0].ref_count}"
-            #         )
+            if time.time() - pre_time_mark > 120:
+                pre_time_mark = time.time()
+                for req_status in self.req_id_to_out_inf.values():
+                    logger.info(
+                        f"left req id {req_status.group_req_objs.group_req_id}"
+                        f"can release {req_status.group_req_objs.shm_req_objs[0].can_released_mark} "
+                        f"refcount {req_status.group_req_objs.shm_req_objs[0].ref_count}"
+                    )
         return
 
     async def handle_loop(self):
diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py
@@ -85,10 +85,6 @@ def init_model(self, kvargs):
         self.pd_rpyc_ports = kvargs.get("pd_rpyc_ports", None)
         max_total_token_num = kvargs["max_total_token_num"]
 
-        if self.dp_size > 1:
-            assert self.dp_size == self.world_size, "Currently only self-sustaining dp_size == tp_size"
-            os.environ["ENABLE_DP"] = "1"
-
         init_distributed_env(kvargs)
         self.init_rank_infos()
         group_size = 2 if self.args.enable_decode_microbatch_overlap else 1
diff --git a/lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_impl.py b/lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_impl.py
@@ -11,6 +11,7 @@
 from lightllm.server.core.objs import FinishStatus
 from lightllm.server.pd_io_struct import UpKVStatus
 from lightllm.utils.log_utils import init_logger
+from lightllm.server.router.model_infer.mode_backend.generic_pre_process import prepare_decode_inputs
 from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
 from .up_status import UpStatusManager
 from rpyc.utils.server import ThreadedServer
@@ -69,8 +70,7 @@ def decode(self):
         self._filter_reqs(aborted_reqs)
 
         if decode_reqs:
-            from lightllm.server.router.model_infer.mode_backend.generic_pre_process import prepare_decode_inputs
-
+            
             kwargs, run_reqs = prepare_decode_inputs(decode_reqs)
             logits = self.model.forward(**kwargs)
 
diff --git a/lightllm/server/router/model_infer/model_rpc.py b/lightllm/server/router/model_infer/model_rpc.py
@@ -138,8 +138,6 @@ def init_model(self, kvargs):
                 self.backend = ContinuesBatchBackendForDecodeNode(self.info_queue, self.mem_queue)
         elif kvargs.get("dp_size", 1) > 1:
             self.backend = DPChunkedPrefillBackend()
-        elif disable_chunked_prefill:
-            self.backend = ContinuesBatchBackend()
         elif use_reward_model:
             self.backend = RewardModelBackend()
         elif return_all_prompt_logprobs:
@@ -154,6 +152,8 @@ def init_model(self, kvargs):
             self.backend = XgrammarBackend()
         elif is_first_token_constraint_mode:
             self.backend = FirstTokenConstraintBackend()
+        elif disable_chunked_prefill:
+            self.backend = ContinuesBatchBackend()
         else:
             self.backend = ChunkedPrefillBackend()
 
diff --git a/lightllm/server/router/req_queue/dp_base_queue.py b/lightllm/server/router/req_queue/dp_base_queue.py
@@ -57,12 +57,13 @@ 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
         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")
-                index = self.round_robin_dp_id
-                self.round_robin_dp_id = (self.round_robin_dp_id + 1) % self.dp_size_in_node
+
                 req.sample_params.suggested_dp_index = index
                 self.inner_queues[index].append(req)
             else:
diff --git a/lightllm/utils/net_utils.py b/lightllm/utils/net_utils.py
@@ -46,6 +46,7 @@ def find_available_port(start_port, end_port):
 def get_hostname_ip():
     try:
         result = subprocess.run(["hostname", "-i"], capture_output=True, text=True, check=True)
+        # 兼容 hostname -i 命令输出多个 ip 的情况
         result = result.stdout.strip().split(" ")[0]
         logger.info(f"get hostname ip {result}")
         return result
diff --git a/test/model/model_infer.py b/test/model/model_infer.py
@@ -14,8 +14,6 @@ def test_model_inference(args, model_class):
     ans_queue = Queue()
     workers = []
     dp_size = args.get("dp", 1)
-    if dp_size > 1:
-        os.environ["ENABLE_DP"] = "1"
 
     for rank_id in range(args.tp):
         model_kvargs = {
