fix

Author: hiworldwzj

lightllm/common/mem_manager.py

@@ -2,6 +2,7 @@
 import os
 import torch
 import torch.distributed as dist
+import torch.multiprocessing as mp
 from typing import List, Union
 from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_for_dp
 from lightllm.server.pd_io_struct import KVMoveTask
@@ -435,6 +436,10 @@ def write_to_shm(self):
         """
         将 mem manager 写入到 shm中，方便pd分离等特性直接从中读取，不依赖进程间队列。
         """
+        from lightllm.server.router.model_infer.mode_backend.continues_batch.pd_mode.p2p_fix import reduce_tensor
+
+        mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__
+
         shm_name = f"{get_unique_server_name()}_mem_manager_{get_current_rank_in_node()}"
         obj_bytes = ForkingPickler.dumps(self).tobytes()
         shm = create_or_link_shm(name=shm_name, expected_size=len(obj_bytes) + 4, force_mode="create")

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

@@ -166,11 +166,6 @@ def exposed_fail_to_realese_forzen_tokens(self, group_req_ids: List[int]):
         release_acquired_lock()
         return
 
-    def exposed_put_mem_manager_to_shm(self):
-        self.backend.model.mem_manager.create_shm()
-        logger.info("put mem manager to shm ok")
-        return
-
     def exposed_unfrozen_time_out_reqs_tokens(self):
         acquire_lock_until_ready(self.backend.lock_nccl_group)
         if self.backend.dp_world_size == 1:

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

@@ -86,7 +86,6 @@ def __init__(self, args, info_queue: mp.Queue):
     # _put_kv_received_to_radix_cache
     # _fail_to_realese_forzen_tokens
     # _unfrozen_time_out_reqs_tokens
-    # _put_mem_manager_to_shm
     # 上述接口都是 kv move manager 与推理进程进行交互的接口，主要用于申请锁定kv资源或者释放
     # kv资源的接口
     # ==================================================================================
@@ -154,12 +153,6 @@ def _unfrozen_time_out_reqs_tokens(self) -> None:
             asyncio.run(self.wait_all_future_finish(futures))
         return
 
-    def _put_mem_manager_to_shm(self) -> None:
-        with self.infer_rpyc_lock:
-            for obj in self.infer_rpyc_objs:
-                obj.put_mem_manager_to_shm()
-        return
-
     # ==================================================================================
     # put_to_fail_release_task_queue 将因为一些原因失败，需要释放锁定的kv资源的请求放入到
     # 对应的处理队列中，handle_fail_release_task_loop 是一个循环的线程，专门处理这些失败的请求

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

@@ -281,11 +281,6 @@ def init_all(self, device_id: int, manager: "DecodeKVMoveManager"):
                 self.task_out_queue,
             )
             assert self.task_out_queue.get(timeout=30) == "proc_start"
-            # 确保在子进程读取共享内存之前，主进程已经将 mem_manager 写入共享内存
-            if self.device_id == 0:
-                manager._put_mem_manager_to_shm()
-            # 通知子进程可以从共享内存读取 mem_manager
-            self.task_in_queue.put("mem_managers_ready")
             assert self.task_out_queue.get(timeout=60) == "get_mem_managers_ok"
 
             return True

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

@@ -111,12 +111,11 @@ def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.
         graceful_registry(inspect.currentframe().f_code.co_name)
         task_out_queue.put("proc_start")
 
-        # 等待主进程将 mem_manager 写入共享内存后的信号
-        assert task_in_queue.get(timeout=60) == "mem_managers_ready"
-
         # 从共享内存读取所有rank的mem_manager
         node_world_size = args.tp // args.nnodes
-        mem_managers: List[MemoryManager] = [MemoryManager.from_shm(rank, device_id) for rank in range(node_world_size)]
+        mem_managers: List[MemoryManager] = [
+            MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size)
+        ]
 
         task_out_queue.put("get_mem_managers_ok")
         connect_id_to_comm: Dict[str, PyNcclCommunicator] = {}

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_infer_rpyc.py

@@ -45,8 +45,3 @@ def exposed_remove_req_refs_from_prompt_cache(self, group_req_ids: List[int]):
                     )
         release_acquired_lock()
         return
-
-    def exposed_put_mem_manager_to_shm(self):
-        self.backend.model.mem_manager.create_shm()
-        logger.info("put mem manager to shm ok")
-        return

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_kv_move_manager.py

@@ -142,8 +142,7 @@ def check_trans_process_loop(self):
             raise e
 
     # ==================================================================================
-    # 与推理进程交互接口,  _remove_req_refs_from_prompt_cache 和
-    # _put_mem_manager_to_shm 都是通过 rpyc 与推理进程进行交互的接口
+    # 与推理进程交互接口,  _remove_req_refs_from_prompt_cache
     # ==================================================================================
 
     def _remove_req_refs_from_prompt_cache(self, tasks: List[KVMoveTask]):
@@ -163,12 +162,6 @@ def _remove_req_refs_from_prompt_cache(self, tasks: List[KVMoveTask]):
             asyncio.run(self.wait_all_future_finish(futures))
         return
 
-    def _put_mem_manager_to_shm(self):
-        with self.infer_rpyc_lock:
-            for obj in self.infer_rpyc_objs:
-                obj.put_mem_manager_to_shm()
-        return
-
     async def wait_all_future_finish(self, futures: List[AsyncResult]):
         await asyncio.gather(*[asyncio.to_thread(future.wait) for future in futures])
         return

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_trans_obj.py

@@ -355,9 +355,6 @@ def init_all(self, device_id: int, manager: "PrefillKVMoveManager"):
                 self.task_out_queue,
             )
             assert self.task_out_queue.get(timeout=30) == "proc_start"
-            if self.device_id == 0:
-                manager._put_mem_manager_to_shm()
-            self.task_in_queue.put("mem_managers_ready")
             assert self.task_out_queue.get(timeout=60) == "get_mem_managers_ok"
 
             return True

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_trans_process.py

@@ -116,12 +116,11 @@ def _init_env(
         )
         task_out_queue.put("proc_start")
 
-        # 等待主进程将 mem_manager 写入共享内存后的信号
-        assert task_in_queue.get(timeout=60) == "mem_managers_ready"
-
         # 从共享内存读取所有rank的mem_manager
         node_world_size = args.tp // args.nnodes
-        mem_managers: List[MemoryManager] = [MemoryManager.from_shm(rank, device_id) for rank in range(node_world_size)]
+        mem_managers: List[MemoryManager] = [
+            MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size)
+        ]
         task_out_queue.put("get_mem_managers_ok")
         connect_id_to_comm: Dict[str, PyNcclCommunicator] = {}
 

lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py

@@ -27,7 +27,6 @@
 from lightllm.common.basemodel.triton_kernel.mtp_utils import mtp_scatter_next_token_ids
 from .control_state import DPControlState
 from lightllm.common.mem_manager import MemoryManager
-import torch.multiprocessing as mp
 
 min_trans_token_num = int(os.getenv("LIGHTLLM_MIN_TRANS_TOKEN_NUM", "512"))
 dp_kv_transfer_req_num = int(os.getenv("LIGHTLLM_DP_KV_TRANSFER_REQ_NUM", "16"))
@@ -74,15 +73,9 @@ def init_custom(self):
         if self.enable_dp_prompt_cache_fetch:
             torch.cuda.set_device(get_current_device_id())
 
-            from lightllm.server.router.model_infer.mode_backend.continues_batch.pd_mode.p2p_fix import reduce_tensor
             from lightllm.server.core.objs.shm_array import ShmArray
             from lightllm.utils.envs_utils import get_unique_server_name
 
-            mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__
-
-            # Create shared memory for mem_manager
-            self.model.mem_manager.create_shm(use_for_pd_trans=False)
-
             # Create shared ShmArray for kv_indexes transfer
             # Use a small buffer to save shared memory
             self.dp_kv_transfer_req_num = dp_kv_transfer_req_num
@@ -101,9 +94,7 @@ def init_custom(self):
             self.mem_managers = []
             for rank_idx in range(self.node_world_size):
                 if rank_idx != self.rank_in_node:
-                    self.mem_managers.append(
-                        MemoryManager.from_shm(rank_idx, self.rank_in_node, use_for_pd_trans=False)
-                    )
+                    self.mem_managers.append(MemoryManager.loads_from_shm(self.rank_in_node))
                 else:
                     self.mem_managers.append(self.model.mem_manager)