refine

Author: niushengxiao

lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py

@@ -31,7 +31,6 @@ class LlamaMultimodalPreLayerInfer(LlamaPreLayerInfer):
     def __init__(self, network_config):
         super().__init__(network_config)
         self.args = get_env_start_args()
-        self.cache_client = None
         if self.args.enable_remote_vit:
             self.cache_client = rpyc.connect("localhost", self.args.cache_port, config={"allow_pickle": True})
             self.cache_client._channel.stream.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
@@ -52,12 +51,14 @@ def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_wei
         img_token_lens = []
         img_start_locs_in_cache = []
         unique_uids = []
+        all_uids = []
         device = layer_weight.wte_weight_.weight.device
         dtype = layer_weight.wte_weight_.weight.dtype
         hidden_size = layer_weight.wte_weight_.weight.shape[1]
 
         for _, p in enumerate(infer_state.multimodal_params):
             for img in p["images"] + p["audios"]:
+                all_uids.append(img["uuid"])
                 # skip the same image
                 if img["token_id"] in img_start_token_ids:
                     continue
@@ -77,17 +78,12 @@ def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_wei
         )
 
         if self.args.enable_remote_vit:
-            release_ids = []
-            for _, p in enumerate(infer_state.multimodal_params):
-                for img in p["images"] + p["audios"]:
-                    release_ids.append(img["uuid"])
-
             for uid, start_index_in_embed_cache in zip(unique_uids, img_start_locs_in_cache):
                 embed_tensor = load_tensor_afs(get_shm_name_embed(uid), self.args.image_embed_dir)
                 self._copy_loaded_embed_to_cache(embed_tensor, cpu_embed_cache_tensor, start_index_in_embed_cache)
 
-            if release_ids:
-                self.cache_client.root.release(release_ids)
+            if all_uids:
+                self.cache_client.root.release(all_uids)
 
         assert cpu_embed_cache_tensor.shape[2] == hidden_size, (
             f"Dimension mismatch: text weight dimension is {hidden_size}, "

lightllm/server/api_http.py

@@ -92,7 +92,7 @@ def set_args(self, args: StartArgs):
             self.httpserver_manager = HttpServerManagerForPDMaster(
                 args=args,
             )
-        elif args.run_mode == "visual":
+        elif args.run_mode in ["visual", "visual_only"]:
             self.metric_client = MetricClient(args.metric_port)
         else:
             init_tokenizer(args)  # for openai api
@@ -138,7 +138,7 @@ def get_model_name():
 @app.get("/health", summary="Check server health")
 @app.head("/health", summary="Check server health")
 async def healthcheck(request: Request):
-    if g_objs.args.run_mode in ["pd_master", "visual"]:
+    if g_objs.args.run_mode in ["pd_master", "visual", "visual_only"]:
         return JSONResponse({"message": "Ok"}, status_code=200)
 
     if os.environ.get("DEBUG_HEALTHCHECK_RETURN_FAIL") == "true":

lightllm/server/api_server.py

@@ -11,7 +11,7 @@
         pd_master_start(args)
     elif args.run_mode == "config_server":
         config_server_start(args)
-    elif args.run_mode == "visual":
+    elif args.run_mode in ["visual", "visual_only"]:
         visual_start(args)
     else:
         normal_or_p_d_start(args)

lightllm/server/api_start.py

@@ -5,7 +5,7 @@
 import subprocess
 import signal
 from lightllm.utils.net_utils import alloc_can_use_network_port, PortLocker
-from lightllm.utils.start_utils import process_manager, kill_recursive, is_multimodal_mode
+from lightllm.utils.start_utils import process_manager, kill_recursive
 from .metrics.manager import start_metric_manager
 from .embed_cache.manager import start_cache_manager
 from lightllm.utils.log_utils import init_logger
@@ -194,7 +194,6 @@ def normal_or_p_d_start(args, only_prepare=False):
         assert args.mtp_draft_model_dir is None
         assert args.mtp_step == 0
 
-    args.enable_multimodal = is_multimodal_mode(args)
     _prepare_remote_vit_embed_dir(args)
     # 检查GPU数量是否足够
     if args.visual_gpu_ids is None:
@@ -355,27 +354,27 @@ def normal_or_p_d_start(args, only_prepare=False):
             start_args=[(args,)],
         )
 
-    if not args.disable_audio:
-        from .audioserver.manager import start_audio_process
+    if not args.disable_vision and not args.enable_remote_vit:
+        from .visualserver.manager import start_visual_process
 
         process_manager.start_submodule_processes(
             start_funcs=[
-                start_audio_process,
+                start_visual_process,
             ],
             start_args=[
-                (args,),
+                (args, visual_model_tp_ports),
             ],
         )
 
-    if not args.disable_vision and not args.enable_remote_vit:
-        from .visualserver.manager import start_visual_process
+    if not args.disable_audio:
+        from .audioserver.manager import start_audio_process
 
         process_manager.start_submodule_processes(
             start_funcs=[
-                start_visual_process,
+                start_audio_process,
             ],
             start_args=[
-                (args, visual_model_tp_ports),
+                (args,),
             ],
         )
 

lightllm/server/embed_cache/manager.py

@@ -54,7 +54,7 @@ def exposed_get_items_embed(self, ids: list[int], embeding_only: bool = False) -
 
 
 def get_cache_manager(args):
-    if args.enable_remote_vit or args.run_mode == "visual":
+    if args.enable_remote_vit or args.run_mode in ["visual", "visual_only"]:
         return MemoryCacheWithRedis(args)
     else:
         return InMemoryCache(args)

lightllm/server/httpserver/manager.py

@@ -82,14 +82,15 @@ def __init__(
         if self.enable_multimodal:
             self.cache_client = rpyc.connect("localhost", args.cache_port, config={"allow_pickle": True})
             self.cache_client._channel.stream.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-            if not self.args.disable_vision:
-                from lightllm.server.visualserver.vit_connect import VITConnectionManager
 
-                self.vit_manager = VITConnectionManager(args, context, args.visual_port, self.cache_client)
+        if not self.args.disable_vision:
+            from lightllm.server.visualserver.vit_connect import VITConnectionManager
 
-            if not self.args.disable_audio:
-                self.send_to_audio = context.socket(zmq.PUSH)
-                self.send_to_audio.connect(f"{args.zmq_mode}127.0.0.1:{args.audio_port}")
+            self.vit_manager = VITConnectionManager(args, context, args.visual_port, self.cache_client)
+
+        if not self.args.disable_audio:
+            self.send_to_audio = context.socket(zmq.PUSH)
+            self.send_to_audio.connect(f"{args.zmq_mode}127.0.0.1:{args.audio_port}")
 
         if args.enable_cpu_cache and not self.args.enable_multimodal:
             self.send_to_multi_level_kv_cache = context.socket(zmq.PUSH)
@@ -151,7 +152,6 @@ async def _alloc_resource(self, items, uuids, token_nums, datas):
             if self.args.enable_remote_vit:
                 # 避免远端lru被逐出
                 self.cache_client.root.get_items_embed(uid_list, False)
-                return
 
             ready_flags = obtain(self.cache_client.root.get_items_data(uid_list))
             update_data_ids = []
@@ -592,25 +592,13 @@ async def transfer_to_next_module(
 
         if self.pd_mode.is_P_or_NORMAL():
             group_req_index = group_req_objs.to_group_req_index()
-            has_images = len(group_req_index.multimodal_params.images) > 0
-            has_audios = len(group_req_index.multimodal_params.audios) > 0
-
-            if has_images and not self.args.disable_vision:
-                free_mode = "all"
-                if self.args.enable_remote_vit and has_audios and not self.args.disable_audio:
-                    free_mode = "images"
-
-                await self.vit_manager.send_to_vit(
-                    group_req_index, protocol=pickle.HIGHEST_PROTOCOL, free_mode=free_mode
-                )
-
+            if not self.args.disable_vision:
+                await self.vit_manager.send_to_vit(group_req_index, protocol=pickle.HIGHEST_PROTOCOL)
                 if not self.args.enable_remote_vit:
                     return
 
-            if has_audios and not self.args.disable_audio:
+            if not self.args.disable_audio:
                 self.send_to_audio.send_pyobj(group_req_index, protocol=pickle.HIGHEST_PROTOCOL)
-                if self.args.enable_remote_vit:
-                    group_req_index.multimodal_params.free()
                 return
 
             if self.args.enable_cpu_cache:

lightllm/server/multimodal_params.py

@@ -26,7 +26,6 @@ def __init__(self, **kwargs):
         self.token_num = None
         # the audio length
         self.audio_length = None
-        self.afs_embed = False
 
         self._preload_data = None
         self.extra_params = {}
@@ -55,11 +54,10 @@ async def preload(self, request: Request):
 
     def read(self):
         assert self._preload_data is not None
-        return self._preload_data
-
-    def free(self):
+        ans = self._preload_data
         self._preload_data = None
         self._data = None
+        return ans
 
     def to_dict(self):
         ret = {}
@@ -167,23 +165,10 @@ def __init__(
         self.audios = [AudioItem(**a) for a in audios]
         return
 
-    def free(self):
-        for image in self.images:
-            image.free()
-        for audio in self.audios:
-            audio.free()
-
     def free_images(self):
         for image in self.images:
             image.free()
 
-    def free_audios(self):
-        for audio in self.audios:
-            audio.free()
-
-    def get_all_uuids(self):
-        return [image.uuid for image in self.images] + [audio.uuid for audio in self.audios]
-
     async def verify_and_preload(self, request: Request):
         for image in self.images:
             await image.preload(request)

lightllm/server/visualserver/manager.py

@@ -239,7 +239,7 @@ async def loop_for_netio_req(self):
                         self.waiting_reqs.append(recv_req)
                     else:
                         assert False, f"Error Req Inf {recv_req}"
-                self.visual_recv_max_count = min(int(self.visual_recv_max_count * 1.3), 256)
+                self.visual_recv_max_count = int(min(self.visual_recv_max_count * 1.3, 256))
             except zmq.ZMQError:
                 # 当队列已经开始清空的时候，将一次接受数量下调
                 self.visual_recv_max_count = 64

lightllm/server/visualserver/vit_connect.py

@@ -159,7 +159,7 @@ def _get_vit_instance(self):
         self.current_vit_index = index
         return list(self.remote_vit_instances.values())[index]
 
-    async def send_to_vit(self, req: GroupReqIndexes, protocol=pickle.HIGHEST_PROTOCOL, free_mode: str = "all"):
+    async def send_to_vit(self, req: GroupReqIndexes, protocol=pickle.HIGHEST_PROTOCOL):
         """
         发送数据到VIT实例，支持本地和远程模式
         """
@@ -174,10 +174,7 @@ async def send_to_vit(self, req: GroupReqIndexes, protocol=pickle.HIGHEST_PROTOC
         if self.remote_vit:
             await self._wait_visual_embed_ready(req)
 
-        if free_mode == "all":
-            req.multimodal_params.free()
-        elif free_mode == "images":
-            req.multimodal_params.free_images()
+        req.multimodal_params.free_images()
 
     async def vit_handle_loop(self):
         """
@@ -223,7 +220,7 @@ async def _wait_visual_embed_ready(
         # 本地模式不需要等待
         if not self.remote_vit:
             return
-        uuids = req.multimodal_params.get_all_uuids()
+        uuids = [image.uuid for image in req.multimodal_params.images]
 
         async def wait_for_embeds():
             while not all(self.cache_client.root.get_items_embed(uuids, True)):

lightllm/utils/start_utils.py

@@ -111,12 +111,4 @@ def kill_recursive(proc):
         logger.warning(f"Process {proc.pid} does not exist.")
 
 
-def is_multimodal_mode(args):
-    from transformers import PretrainedConfig
-
-    model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir)
-    is_multimodal = "visual" in model_cfg or "vision_config" in model_cfg
-    return is_multimodal
-
-
 process_manager = SubmoduleManager()