fix

Author: shihaobai

lightllm/common/image_cache_manager.py

@@ -58,7 +58,7 @@ def filter(self, uuid_list):
         """
         for uuid in uuid_list:
             if uuid in self._gpu_cache:
-                embed_cpu = self._gpu_cache[uuid].cpu(non_blocking=True)
+                embed_cpu = self._gpu_cache[uuid].cpu()
                 # Move to CPU cache and remove from GPU cache
                 self._gpu_cache.pop(uuid)
                 if uuid in self._cpu_cache:
@@ -68,6 +68,8 @@ def filter(self, uuid_list):
                     self._cpu_cache.popitem(last=False)
             elif uuid in self._cpu_cache:
                 self._cpu_cache.move_to_end(uuid)
+        print(self._gpu_cache.keys())
+        print(self._cpu_cache.keys())
         return
 
 

lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py

@@ -68,7 +68,6 @@ def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_wei
                 if self.disable_extra_process_for_multimodal:
                     img_embed = image_cache_manager.get_embed(img["uuid"])
                     img_weight.append(img_embed.reshape(img["token_num"], -1))
-                    print(img_weight[-1].shape)
                 else:
                     data = read_shm(get_shm_name_embed(img["uuid"]))
                     img_weight.append(bytes2tensor(data).cuda().reshape(img["token_num"], -1))

lightllm/models/vit/layer_infer/pre_layer_infer.py

@@ -16,6 +16,7 @@ def __init__(self, network_config, mode):
         self.tp_world_size_ = get_dp_world_size()
         self.network_config_ = network_config
         self.mode = mode
+        print(f"tp_rank_: {self.tp_rank_}, tp_world_size_: {self.tp_world_size_}")
         return
 
     def forward(self, pixel_values, layer_weight: ViTPreAndPostLayerWeight):

lightllm/models/vit/layer_infer/transformer_layer_infer.py

@@ -21,6 +21,7 @@ class ViTTransformerLayerInfer:
     def __init__(self, layer_num, network_config, mode=[]):
         self.tp_rank_ = get_current_rank_in_dp()
         self.tp_world_size_ = get_dp_world_size()
+        print(f"tp_rank_: {self.tp_rank_}, tp_world_size_: {self.tp_world_size_}")
         self.eps_ = network_config["layer_norm_eps"]
         self.head_num = network_config["num_attention_heads"]
         self.tp_padding_head_num = network_config["padding_head_num"] // self.tp_world_size_

lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py

@@ -5,10 +5,9 @@
 from lightllm.common.basemodel import PreAndPostLayerWeight
 from lightllm.utils.dist_utils import (
     get_current_device_id,
-    get_global_rank,
-    get_global_world_size,
+    get_current_rank_in_dp,
+    get_dp_world_size,
 )
-from lightllm.utils.envs_utils import get_env_start_args
 
 
 class ViTPreAndPostLayerWeight(PreAndPostLayerWeight):
@@ -18,10 +17,8 @@ def __init__(self, data_type, network_config, mode):
         self.image_size = self.network_config_["image_size"]
         self.patch_size = self.network_config_["patch_size"]
         self.llm_hidden_size = self.network_config_["llm_hidden_size"]
-        if get_env_start_args().disable_extra_process_for_multimodal:
-            self.tp_world_size_ = get_global_world_size()
-            self.tp_rank_ = get_global_rank()
-
+        self.tp_rank_ = get_current_rank_in_dp()
+        self.tp_world_size_ = get_dp_world_size()
         return
 
     def _cuda(self, cpu_tensor):

lightllm/models/vit/layer_weights/transformer_layer_weight.py

@@ -13,18 +13,12 @@
 )
 from lightllm.utils.dist_utils import (
     get_current_device_id,
-    get_global_rank,
-    get_global_world_size,
 )
-from lightllm.utils.envs_utils import get_env_start_args
 
 
 class ViTTransformerLayerWeight(TransformerLayerWeight):
     def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
         super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
-        if get_env_start_args().disable_extra_process_for_multimodal:
-            self.tp_world_size_ = get_global_world_size()
-            self.tp_rank_ = get_global_rank()
         return
 
     def _cuda(self, cpu_tensor):

lightllm/models/vit/model.py

@@ -38,11 +38,7 @@ class VisionTransformer:
     post_layer_infer_class = ViTPostLayerInfer
 
     def __init__(self, kvargs):
-        if get_env_start_args().disable_extra_process_for_multimodal:
-            # if we don't assign an extra process for visual model, the visual model uses tensor parallel by default.
-            self.tp_world_size_ = get_global_world_size()
-        else:
-            self.tp_world_size_ = get_dp_world_size()
+        self.tp_world_size_ = get_dp_world_size()
         self.weight_dir_ = kvargs["weight_dir"]
         self.load_way = kvargs.get("load_way", "HF")
         self.mode = [m.replace("int4weight", "w4a16").replace("int8weight", "w8a16") for m in kvargs.get("mode", [])]

lightllm/server/router/model_infer/infer_batch.py

@@ -9,6 +9,7 @@
 from dataclasses import dataclass, field
 from typing import List, Dict, Tuple, Optional, Union, Any
 from lightllm.common.req_manager import ReqManager
+from lightllm.common.image_cache_manager import image_cache_manager
 from lightllm.utils.infer_utils import mark_start, mark_end
 from lightllm.server.core.objs import Req, SamplingParams, FinishStatus, ShmReqManager
 from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode
@@ -131,6 +132,7 @@ def filter(self, finished_request_ids: List[int]):
 
         free_req_index = []
         free_token_index = []
+        image_uuid_list = []
         for request_id in finished_request_ids:
             req: InferReq = self.requests_mapping.pop(request_id)
             group_req_id = convert_sub_id_to_group_id(req.shm_req.request_id)
@@ -145,6 +147,10 @@ def filter(self, finished_request_ids: List[int]):
             # logger.info(f"infer release req id {req.shm_req.request_id}")
             req.shm_req.shm_infer_released = True
             self.shm_req_manager.put_back_req_obj(req.shm_req)
+            if req.multimodal_params is not None and get_env_start_args().disable_extra_process_for_multimodal:
+                for img in req.multimodal_params["images"]:
+                    image_uuid_list.append(img["uuid"])
+                image_cache_manager.filter(image_uuid_list)
 
         free_token_index = custom_cat(free_token_index)
         self.req_manager.free(free_req_index, free_token_index)