fix

Author: flyinglandlord

lightllm/server/api_start.py

@@ -16,7 +16,7 @@
 from .router.manager import start_router_process
 from lightllm.utils.process_check import is_process_active
 from lightllm.utils.multinode_utils import send_and_receive_node_ip
-from lightllm.utils.shm_size_check import check_shm_size, start_shm_size_warning_thread
+from lightllm.utils.shm_size_check import check_recommended_shm_size
 
 logger = init_logger(__name__)
 
@@ -64,15 +64,7 @@ def signal_handler(sig, frame):
 def normal_or_p_d_start(args):
     set_unique_server_name(args)
 
-    shm_size, require_shm_size, is_shm_sufficient = check_shm_size(args)
-    if not args.disable_shm_warning and not is_shm_sufficient:
-        start_shm_size_warning_thread(shm_size, require_shm_size)
-    else:
-        logger.info(
-            f"SHM check: Available={shm_size:.2f} GB,",
-            f"Required={require_shm_size:.2f} GB.",
-            f"Sufficient: {is_shm_sufficient}",
-        )
+    check_recommended_shm_size(args)
 
     if args.enable_mps:
         from lightllm.utils.device_utils import enable_mps

lightllm/utils/config_utils.py

@@ -13,6 +13,23 @@ def get_config_json(model_path: str):
     return json_obj
 
 
+def get_hidden_size(model_path: str):
+    # try to get hidden_size in config.json
+    config_json = get_config_json(model_path)
+    try:
+        hidden_size = config_json["hidden_size"]
+    except:
+        # for some multimodal model
+        try:
+            hidden_size = config_json.get("llm_config", {}).get("hidden_size")
+        except:
+            hidden_size = config_json.get("text_config", {}).get("hidden_size")
+    if not isinstance(hidden_size, int):
+        logger.error("cannot get hidden size from config.json, return None instead")
+        return None
+    return hidden_size
+
+
 def get_eos_token_ids(model_path: str):
     config_json = get_config_json(model_path)
     try:

lightllm/utils/shm_size_check.py

@@ -7,17 +7,54 @@
 import signal
 from lightllm.server.core.objs.out_token_circlequeue import LIGHTLLM_OUT_TOKEN_QUEUE_SIZE, LIGHTLLM_TOKEN_MAX_BYTES
 from lightllm.server.core.objs.req import ChunkedPrefillReq, TokenHealingReq
+from lightllm.server.multimodal_params import ImageItem
 from lightllm.server.tokenizer import get_tokenizer
-from lightllm.utils.config_utils import get_config_json
+from lightllm.utils.config_utils import get_hidden_size
 from lightllm.utils.log_utils import init_logger
 from transformers import AutoTokenizer
 
 logger = init_logger(__name__)
 
 
-def get_shm_size_gb():
+def check_recommended_shm_size(args):
+    shm_size, recommended_shm_size, is_shm_sufficient = _check_shm_size(args)
+    if not args.disable_shm_warning and not is_shm_sufficient:
+        _start_shm_size_warning_thread(shm_size, recommended_shm_size)
+    else:
+        logger.info(
+            f"SHM check: Available={shm_size:.2f} GB,"
+            f"Recommended={recommended_shm_size:.2f} GB."
+            f"Sufficient: {is_shm_sufficient}",
+        )
+
+
+def _check_shm_size(args):
+    RED = "\033[91m"
+    ENDC = "\033[0m"
+    shm_size = _get_system_shm_size_gb()
+    required_size = _get_recommended_shm_size_gb(args)  # 128G
+    if shm_size < required_size:
+        logger.warning(f"{RED}Available shm size is less than 128G: {shm_size:.2f}G{ENDC}")
+        return shm_size, required_size, False
+    else:  # shm_size >= required_size
+        return shm_size, required_size, True
+
+
+def _start_shm_size_warning_thread(shm_size, required_shm_size):
+    shm_warning_thread = threading.Thread(
+        target=_periodic_shm_warning,
+        args=(
+            shm_size,
+            required_shm_size,
+        ),
+        daemon=True,
+    )
+    shm_warning_thread.start()
+
+
+def _get_system_shm_size_gb():
     """
-    获取 /dev/shm 的总大小（以GB为单位）。
+    获取 /dev/shm 的总大小(以GB为单位)。
     """
     try:
         shm_path = "/dev/shm"
@@ -34,87 +71,70 @@ def get_shm_size_gb():
         return 0
 
 
-def get_required_shm_size_gb(args, max_image_resolution=(3940, 2160), dtype_size=2):
+def _get_recommended_shm_size_gb(args, max_image_resolution=(3940, 2160), dtype_size=2):
     """
-    获取所需的 /dev/shm 大小（以GB为单位）。
+    获取所需的 /dev/shm 大小(以GB为单位)。
     """
-    model_config = get_config_json(args.model_dir)
     tokenizer = get_tokenizer(args.model_dir, trust_remote_code=True)
 
+    # 估算input_token和logprob占用shm大小，由于是double和int64，所以固定占用8个字节
+    input_token_logprob_size_bytes = args.running_max_req_size * 8 * 2 * args.max_req_total_len
+
+    # 估算Req所需的shm大小
+    if args.token_healing_mode:
+        req_class_size = ctypes.sizeof(TokenHealingReq)
+    else:
+        req_class_size = ctypes.sizeof(ChunkedPrefillReq)
+    req_shm_size_bytes = req_class_size * args.running_max_req_size
+
     if not args.enable_multimodal:
-        # by default, 非多模态 24 GB
-        total_required_size_gb = 24
+        total_recommended_shm_size_gb = (req_shm_size_bytes + input_token_logprob_size_bytes) / (1024 ** 3) + 2
     else:
+        # 存储图片数据所需的shm大小
         num_channels = 3
         image_width, image_height = max_image_resolution
         image_size_bytes = image_width * image_height * num_channels
 
         # 假设加载最大分辨率图片时，通过 tokenizer 得到最多的 image_tokens
         if not hasattr(tokenizer, "get_image_token_length"):
             raise AttributeError("Tokenizer must have a 'get_image_token_length' method for multimodal models.")
-        max_image_tokens = tokenizer.get_image_token_length(None)
+        fake_image_item = ImageItem(
+            type="image_size",
+            data=max_image_resolution,
+        )
+        fake_image_item.image_w = fake_image_item._data[0]
+        fake_image_item.image_h = fake_image_item._data[1]
+        max_image_tokens = tokenizer.get_image_token_length(fake_image_item)
 
         # 估算图片 token 所需的资源
-        hidden_size = model_config.get("hidden_size")
+        hidden_size = get_hidden_size(args.model_dir)
         if hidden_size is None:
-            logger.warning("Model config not contain 'hidden_size', using 4096 by default.")
+            logger.warning(
+                "Model config not contain 'hidden_size', " "using 4096 by default to calculate recommended shm size."
+            )
             image_token_size_bytes = max_image_tokens * 4096 * dtype_size
         else:
             image_token_size_bytes = max_image_tokens * hidden_size * dtype_size
 
-        # 估算Req所需的shm大小
-        if args.token_healing_mode:
-            req_class_size = ctypes.sizeof(TokenHealingReq)
-        else:
-            req_class_size = ctypes.sizeof(ChunkedPrefillReq)
-        req_shm_size_bytes = req_class_size * args.running_max_req_size
-
-        # 估算OutTokenQueue所需shm大小
-        out_token_queue_size_bytes = LIGHTLLM_TOKEN_MAX_BYTES * LIGHTLLM_OUT_TOKEN_QUEUE_SIZE
-
-        total_required_size = (
+        total_recommended_shm_size_gb = (
             args.cache_capacity * (image_size_bytes + image_token_size_bytes)
             + req_shm_size_bytes
-            + out_token_queue_size_bytes
+            + input_token_logprob_size_bytes
         )
 
-        total_required_size_gb = total_required_size / (1024 ** 3) + 2
-
-    return total_required_size_gb
+        total_recommended_shm_size_gb = total_recommended_shm_size_gb / (1024 ** 3) + 2
 
-
-def check_shm_size(args):
-    RED = "\033[91m"
-    ENDC = "\033[0m"
-    shm_size = get_shm_size_gb()
-    required_size = get_required_shm_size_gb(args)  # 128G
-    if shm_size < required_size:
-        logger.warning(f"{RED}Available shm size is less than 128G: {shm_size:.2f}G{ENDC}")
-        return shm_size, required_size, False
-    else:  # shm_size >= required_size
-        return shm_size, required_size, True
+    return total_recommended_shm_size_gb
 
 
-def periodic_shm_warning(shm_size, required_shm_size):
+def _periodic_shm_warning(shm_size, required_shm_size):
     RED = "\033[91m"
     ENDC = "\033[0m"
     while True:
         logger.warning(
-            f"{RED}Insufficient shared memory (SHM) available.",
-            f"Required: {required_shm_size:.2f}G, Available: {shm_size:.2f}G.\n",
-            "If running in Docker, you can increase SHM size with the `--shm-size` flag, ",
+            f"{RED}Insufficient shared memory (SHM) available."
+            f"Required: {required_shm_size:.2f}G, Available: {shm_size:.2f}G.\n"
+            "If running in Docker, you can increase SHM size with the `--shm-size` flag, "
             f"like so: `docker run --shm-size=30g [your_image]`{ENDC}",
         )
         time.sleep(120)  # 每 120 秒打印一次警告日志
-
-
-def start_shm_size_warning_thread(shm_size, required_shm_size):
-    shm_warning_thread = threading.Thread(
-        target=periodic_shm_warning,
-        args=(
-            shm_size,
-            required_shm_size,
-        ),
-        daemon=True,
-    )
-    shm_warning_thread.start()