use rpyc in model_rpc.py (#1221)

Author: hiworldwzj

lightllm/models/bloom/layer_weights/transformer_layer_weight.py

@@ -56,6 +56,10 @@ def _parse_config(self):
         self.n_embed = self.network_config_["n_embed"]
         self.n_head = self.network_config_["num_attention_heads"]
         self.n_inter = self.network_config_["n_embed"] * 4
+        self.q_head_num_ = self.n_head
+        self.k_head_num_ = self.n_head
+        self.v_head_num_ = self.n_head
+        self.o_head_num_ = self.n_head
         self.n_kv_head = self.network_config_["num_attention_heads"]
         self.head_dim = self.network_config_.get("head_dim", self.n_embed // self.n_head)
         # 计算生成alibi

lightllm/server/api_cli.py

@@ -259,7 +259,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
     )
     parser.add_argument("--disable_dynamic_prompt_cache", action="store_true", help="disable dynamic prompt cache")
 
-    parser.add_argument("--chunked_prefill_size", type=int, default=8192, help="chunked prefill size")
+    parser.add_argument("--chunked_prefill_size", type=int, default=None, help="chunked prefill size")
     parser.add_argument("--disable_chunked_prefill", action="store_true", help="whether to disable chunked prefill")
     parser.add_argument("--diverse_mode", action="store_true", help="diversity generation mode")
     parser.add_argument("--token_healing_mode", action="store_true", help="code model infer mode")

lightllm/server/core/objs/__init__.py

@@ -1,6 +1,5 @@
 from .sampling_params import SamplingParams
 from .req import Req, FinishStatus
 from .shm_req_manager import ShmReqManager
-from .rpc_shm import RpcShmParams, RpcShmResults, ShmSyncStatusArray
 from .start_args_type import StartArgs
 from .atomic_lock import AtomicShmLock

lightllm/server/core/objs/rpc_shm.py

@@ -115,37 +115,28 @@ async def wait_to_model_ready(self):
         self.model_rpc_servers = []
         # 用于 kv move 管理进程 和 推理进程进行task信息的交互。
         self.info_queue: mp.Queue = mp.Queue()
-        self.rpc_event = multiprocessing.Event()
-        self.rpc_finished_event = multiprocessing.Event()
 
         assert (self.world_size % self.nnodes) == 0
         node_world_size = self.world_size // self.nnodes
 
         # Create tasks for parallel startup
         tasks = []
         for rank_id in range(self.node_rank * node_world_size, (self.node_rank + 1) * node_world_size):
+            rank_in_node = rank_id % node_world_size
             task = asyncio.create_task(
                 start_model_process(
                     args=self.args,
                     rank=rank_id,
-                    rank_in_node=rank_id % node_world_size,
+                    rank_in_node=rank_in_node,
                     node_world_size=node_world_size,
-                    rpc_event=self.rpc_event,
-                    rpc_finished_event=self.rpc_finished_event,
                     info_queue=self.info_queue,
                     router_lock=self.router_lock,
                 )
             )
             tasks.append(task)
 
         # Wait for all tasks to complete in parallel
-        self.model_rpc_servers = await asyncio.gather(*tasks)
-
-        self.model_rpc_client = ModelRpcClient(
-            rpc_event=self.rpc_event,
-            rpc_finished_event=self.rpc_finished_event,
-        )
-
+        self.model_rpc_clients = await asyncio.gather(*tasks)
         kvargs = {
             "args": self.args,
             "rank_id": None,  # 由后续处理填充真实数据
@@ -178,10 +169,19 @@ async def wait_to_model_ready(self):
             "pd_rpyc_ports": self.args.pd_node_infer_rpyc_ports,  # 非 pd 模式可以不设置
         }
 
-        await self.model_rpc_client.init_model(kvargs=kvargs)
+        # Call init_model on all model processes
+        init_tasks = []
+        for model_rpc_client in self.model_rpc_clients:
+            init_tasks.append(model_rpc_client.init_model(kvargs=kvargs))
+        await asyncio.gather(*init_tasks)
 
         if self.max_total_token_num is None:
-            self.max_total_token_num = await self.model_rpc_client.get_max_total_token_num()
+            _tasks = []
+            for model_rpc_client in self.model_rpc_clients:
+                _tasks.append(model_rpc_client.get_max_total_token_num())
+            _nums = await asyncio.gather(*_tasks)
+            assert max(_nums) == min(_nums), "all rank must have same token num"
+            self.max_total_token_num = _nums[0]
             self.args.max_total_token_num = self.max_total_token_num
         if not self.args.disable_dynamic_prompt_cache:
             self.radix_cache_client = RadixCacheReadOnlyClient(