Vit sep

lightllm/models/vit/triton_kernel/flashattention_nopad.py

@@ -195,16 +195,16 @@ def flash_attention_v3_fwd(
             False,
             window_size[0],
             window_size[1],
-            0.0,
+            attention_chunk=0,
+            softcap=0.0,
             is_rotary_interleaved=False,
             scheduler_metadata=None,
             num_splits=1,
             pack_gqa=None,
             sm_margin=0,
             sinks=None,
         )
-
-        return
+        return o
 
 except ImportError:
     print("Failed to import _flash_attn_forward from hopper.flash_attn_interface.")

lightllm/server/api_cli.py

@@ -7,7 +7,16 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--run_mode",
         type=str,
-        choices=["normal", "prefill", "decode", "nixl_prefill", "nixl_decode", "pd_master", "config_server"],
+        choices=[
+            "normal",
+            "prefill",
+            "decode",
+            "nixl_prefill",
+            "nixl_decode",
+            "pd_master",
+            "config_server",
+            "visual_only",
+        ],
         default="normal",
         help="""set run mode, normal is started for a single server, prefill decode pd_master is for pd split run mode,
                 config_server is for pd split mode used to register pd_master node, and get pd_master node list,
@@ -61,6 +70,14 @@ def make_argument_parser() -> argparse.ArgumentParser:
         default=None,
         help="The port number for the config server in config_server mode.",
     )
+    parser.add_argument(
+        "--config_server_visual_redis_port",
+        type=int,
+        default=None,
+        help="""when run_mode is config_server, set this params will start a redis server,
+        when a llm infer node start to set this params, the visual infer module will start a
+        proxy module use config server to find  remote vit infer nodes to infer img""",
+    )
     parser.add_argument(
         "--nixl_pd_kv_page_num",
         type=int,
@@ -440,6 +457,15 @@ def make_argument_parser() -> argparse.ArgumentParser:
         default=None,
         help="List of NCCL ports to build a distributed environment for Vit, e.g., 29500 29501 29502",
     )
+    parser.add_argument(
+        "--visual_rpyc_port",
+        type=int,
+        default=None,
+        help="""
+            when run_mode is visual_only, set this port, make others to call local visual infer to
+            transfer image to embed.
+            """,
+    )
     parser.add_argument(
         "--enable_monitor_auth", action="store_true", help="Whether to open authentication for push_gateway"
     )
@@ -605,6 +631,12 @@ def make_argument_parser() -> argparse.ArgumentParser:
         default=0.03,
         help="""The interval of the schedule time, default is 30ms.""",
     )
+    parser.add_argument(
+        "--afs_image_embed_dir",
+        type=str,
+        default=None,
+        help="path for vit embed, when use vit remote infer mode",
+    )
     parser.add_argument(
         "--enable_cpu_cache",
         action="store_true",

lightllm/server/api_http.py

@@ -43,7 +43,7 @@
 from .multimodal_params import MultimodalParams
 from .httpserver.manager import HttpServerManager
 from .httpserver_for_pd_master.manager import HttpServerManagerForPDMaster
-from .api_lightllm import lightllm_get_score
+from .api_lightllm import lightllm_get_score, lightllm_get_image_embedding
 from lightllm.utils.envs_utils import get_env_start_args, get_lightllm_websocket_max_message_size
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.error_utils import ServerBusyError
@@ -92,6 +92,8 @@ def set_args(self, args: StartArgs):
             self.httpserver_manager = HttpServerManagerForPDMaster(
                 args=args,
             )
+        elif args.run_mode in ["visual", "visual_only"]:
+            self.metric_client = MetricClient(args.metric_port)
         else:
             init_tokenizer(args)  # for openai api
             SamplingParams.load_generation_cfg(args.model_dir)
@@ -136,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 == "pd_master":
+    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":
@@ -221,6 +223,18 @@ async def get_score(request: Request) -> Response:
         return create_error_response(HTTPStatus.EXPECTATION_FAILED, str(e))
 
 
+@app.post("/get_image_embedding")
+async def get_image_embed(request: Request) -> Response:
+    try:
+        return await lightllm_get_image_embedding(request, g_objs.httpserver_manager)
+    except ServerBusyError as e:
+        logger.error("%s", str(e), exc_info=True)
+        return create_error_response(HTTPStatus.SERVICE_UNAVAILABLE, str(e))
+    except Exception as e:
+        logger.error("An error occurred: %s", str(e), exc_info=True)
+        return create_error_response(HTTPStatus.EXPECTATION_FAILED, str(e))
+
+
 @app.post("/")
 async def compat_generate(request: Request) -> Response:
     if get_env_start_args().run_mode in ["prefill", "decode", "nixl_prefill", "nixl_decode"]:
@@ -359,6 +373,8 @@ async def startup_event():
     logger.info("server start up")
     loop = asyncio.get_event_loop()
     g_objs.set_args(get_env_start_args())
+    if g_objs.httpserver_manager is None:
+        return
     loop.create_task(g_objs.httpserver_manager.handle_loop())
     logger.info(f"server start up ok, loop use is {asyncio.get_event_loop()}")
     return

lightllm/server/api_lightllm.py

@@ -1,10 +1,11 @@
 import collections
 from typing import AsyncGenerator
 from fastapi import BackgroundTasks, Request
-from fastapi.responses import Response, StreamingResponse
+from fastapi.responses import Response, StreamingResponse, JSONResponse
 from lightllm.server.core.objs.sampling_params import SamplingParams
 from .multimodal_params import MultimodalParams
 from .httpserver.manager import HttpServerManager
+from lightllm.utils.envs_utils import get_env_start_args
 import ujson as json
 
 
@@ -150,3 +151,21 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
 
     background_tasks = BackgroundTasks()
     return StreamingResponse(stream_results(), media_type="text/event-stream", background=background_tasks)
+
+
+async def lightllm_get_image_embedding(request: Request, httpserver_manager: HttpServerManager) -> Response:
+    request_dict = await request.json()
+    args = get_env_start_args()
+    assert not args.disable_vision
+    assert args.enable_remote_vit
+    sample_params_dict = request_dict["parameters"]
+    sampling_params = SamplingParams()
+    sampling_params.init(tokenizer=None, **sample_params_dict)
+    sampling_params.verify()
+    multimodal_params_dict = request_dict.get("multimodal_params", {})
+    assert not multimodal_params_dict.get("audios")
+    multimodal_params = MultimodalParams(**multimodal_params_dict)
+
+    await httpserver_manager.get_image_embeding(sampling_params, multimodal_params, request=request)
+
+    return JSONResponse({"message": "OK"}, status_code=200)

lightllm/server/api_server.py

@@ -5,11 +5,13 @@
     torch.multiprocessing.set_start_method("spawn")  # this code will not be ok for settings to fork to subprocess
     parser = make_argument_parser()
     args = parser.parse_args()
-    from .api_start import pd_master_start, normal_or_p_d_start, config_server_start
+    from .api_start import pd_master_start, normal_or_p_d_start, visual_start, config_server_start
 
     if args.run_mode == "pd_master":
         pd_master_start(args)
     elif args.run_mode == "config_server":
         config_server_start(args)
+    elif args.run_mode in ["visual", "visual_only"]:
+        visual_start(args)
     else:
         normal_or_p_d_start(args)

lightllm/server/api_start.py

@@ -15,12 +15,37 @@
 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.redis_utils import start_redis_service
 from lightllm.utils.shm_size_check import check_recommended_shm_size
 from lightllm.utils.config_utils import has_audio_module, has_vision_module
 
 logger = init_logger(__name__)
 
 
+def _ensure_remote_vit_embed_dir(image_embed_dir: str) -> None:
+    if os.path.exists(image_embed_dir):
+        if not os.path.isdir(image_embed_dir):
+            raise ValueError(f"image_embed_dir is not a directory: {image_embed_dir}")
+        return
+
+    os.makedirs(image_embed_dir, mode=0o777, exist_ok=True)
+    os.chmod(image_embed_dir, 0o777)
+
+
+def _prepare_remote_vit_embed_dir(args):
+    remote_vit_mode = args.enable_remote_vit or args.run_mode in ["visual", "visual_only"]
+    if not remote_vit_mode:
+        return
+
+    if not args.image_embed_dir:
+        raise ValueError("remote vit mode requires --image_embed_dir to be set")
+
+    args.image_embed_dir = os.path.abspath(args.image_embed_dir)
+    _ensure_remote_vit_embed_dir(args.image_embed_dir)
+
+    logger.info(f"using image_embed_dir: {args.image_embed_dir}")
+
+
 def setup_signal_handlers(http_server_process, process_manager):
     def signal_handler(sig, frame):
         if sig == signal.SIGINT:
@@ -57,11 +82,12 @@ def signal_handler(sig, frame):
     signal.signal(signal.SIGINT, signal_handler)
 
     logger.info(f"start process pid {os.getpid()}")
-    logger.info(f"http server pid {http_server_process.pid}")
+    if http_server_process:
+        logger.info(f"http server pid {http_server_process.pid}")
     return
 
 
-def normal_or_p_d_start(args):
+def normal_or_p_d_start(args, only_prepare=False):
     from lightllm.server.core.objs.start_args_type import StartArgs
 
     args: StartArgs = args
@@ -73,7 +99,7 @@ def normal_or_p_d_start(args):
 
         enable_mps()
 
-    if args.run_mode not in ["normal", "prefill", "decode", "nixl_prefill", "nixl_decode"]:
+    if args.run_mode not in ["normal", "prefill", "decode", "nixl_prefill", "nixl_decode", "visual", "visual_only"]:
         return
 
     # 通过模型的参数判断是否是多模态模型，包含哪几种模态, 并设置是否启动相应得模块
@@ -168,6 +194,7 @@ def normal_or_p_d_start(args):
         assert args.mtp_draft_model_dir is None
         assert args.mtp_step == 0
 
+    _prepare_remote_vit_embed_dir(args)
     # 检查GPU数量是否足够
     if args.visual_gpu_ids is None:
         args.visual_gpu_ids = list(range(args.visual_dp * args.visual_tp))
@@ -229,20 +256,27 @@ def normal_or_p_d_start(args):
         args.data_type = get_dtype(args.model_dir)
         assert args.data_type in ["fp16", "float16", "bf16", "bfloat16", "fp32", "float32"]
 
+    if only_prepare:
+        return
+
     already_uesd_ports = [args.port]
     if args.nccl_port is not None:
         already_uesd_ports.append(args.nccl_port)
     if args.pd_decode_rpyc_port is not None:
         already_uesd_ports.append(args.pd_decode_rpyc_port)
+    if args.visual_nccl_ports is not None:
+        already_uesd_ports.extend(args.visual_nccl_ports[: args.visual_dp])
 
     # 提前锁定端口，防止在单个机器上启动多个实列的时候，要到模型启动的时候才能
     # 捕获到端口设置冲突的问题
     ports_locker = PortLocker(already_uesd_ports)
     ports_locker.lock_port()
 
     node_world_size = args.tp // args.nnodes
+    need_visual_nccl_ports = 0 if args.visual_nccl_ports is not None else args.visual_dp
     can_use_ports = alloc_can_use_network_port(
-        num=10 + node_world_size + args.visual_dp * (args.visual_tp + 1), used_ports=already_uesd_ports
+        num=10 + node_world_size + args.visual_dp * args.visual_tp + need_visual_nccl_ports,
+        used_ports=already_uesd_ports,
     )
     logger.info(f"alloced ports: {can_use_ports}")
     (
@@ -265,8 +299,12 @@ def normal_or_p_d_start(args):
         tp_ports_for_dp = can_use_ports[0 : args.visual_tp]
         visual_model_tp_ports.append(tp_ports_for_dp)
         can_use_ports = can_use_ports[args.visual_tp :]
-        visual_nccl_ports.append(can_use_ports[0])
-        can_use_ports = can_use_ports[1:]
+        if args.visual_nccl_ports is None:
+            visual_nccl_ports.append(can_use_ports[0])
+            can_use_ports = can_use_ports[1:]
+
+    if args.visual_nccl_ports is not None:
+        args.visual_nccl_ports = args.visual_nccl_ports[: args.visual_dp]
 
     # 将申请好的端口放入args参数中
     if args.nccl_port is None:
@@ -316,7 +354,7 @@ def normal_or_p_d_start(args):
             start_args=[(args,)],
         )
 
-    if not args.disable_vision:
+    if not args.disable_vision and not args.enable_remote_vit:
         from .visualserver.manager import start_visual_process
 
         process_manager.start_submodule_processes(
@@ -463,13 +501,89 @@ def pd_master_start(args):
     http_server_process.wait()
 
 
+def visual_start(args):
+    normal_or_p_d_start(args, only_prepare=True)
+
+    already_uesd_ports = [args.remote_vit_port]
+    if args.nccl_port is not None:
+        already_uesd_ports.append(args.nccl_port)
+    if args.visual_nccl_ports is not None:
+        already_uesd_ports.extend(args.visual_nccl_ports[: args.visual_dp])
+
+    need_visual_nccl_ports = 0 if args.visual_nccl_ports is not None else args.visual_dp
+    can_use_ports = alloc_can_use_network_port(
+        num=5 + args.visual_dp * args.visual_tp + need_visual_nccl_ports,
+        used_ports=already_uesd_ports,
+    )
+    logger.info(f"alloced ports: {can_use_ports}")
+    (
+        router_port,
+        visual_port,
+        audio_port,
+        cache_port,
+        metric_port,
+    ) = can_use_ports[0:5]
+    can_use_ports = can_use_ports[5:]
+
+    visual_nccl_ports = []
+    for _ in range(args.visual_dp):
+        if args.visual_nccl_ports is None:
+            visual_nccl_ports.append(can_use_ports[0])
+            can_use_ports = can_use_ports[1:]
+
+    if args.visual_nccl_ports is not None:
+        args.visual_nccl_ports = args.visual_nccl_ports[: args.visual_dp]
+    else:
+        args.visual_nccl_ports = visual_nccl_ports
+
+    args.router_port = router_port
+    args.visual_port = visual_port
+    args.audio_port = audio_port
+    args.cache_port = cache_port
+    args.metric_port = metric_port
+    args.visual_node_id = uuid.uuid4().int
+
+    logger.info(f"all start args:{args}")
+
+    set_env_start_args(args)
+
+    from .visualserver.manager import start_visual_process
+
+    process_manager.start_submodule_processes(
+        start_funcs=[
+            start_cache_manager,
+        ],
+        start_args=[(args,)],
+    )
+    process_manager.start_submodule_processes(
+        start_funcs=[
+            start_visual_process,
+        ],
+        start_args=[
+            (args,),
+        ],
+    )
+    setup_signal_handlers(None, process_manager)
+    try:
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        logger.info("Received keyboard interrupt, shutting down...")
+        process_manager.terminate_all_processes()
+        logger.info("All processes have been terminated gracefully.")
+        sys.exit(0)
+
+
 def config_server_start(args):
     set_unique_server_name(args)
     if args.run_mode != "config_server":
         return
 
     logger.info(f"all start args:{args}")
 
+    if args.start_redis:
+        start_redis_service(args)
+
     set_env_start_args(args)
 
     command = [