切换vllm 后端为 V0,加快启动和推理性能

Author: shell-nlp

gpt_server/model_backend/vllm_backend.py

@@ -19,9 +19,7 @@
 # 解决vllm中 ray集群在 TP>1时死的Bug
 import ray
 
-ray.init(ignore_reinit_error=True, num_cpus=4)
-
-os.environ["VLLM_USE_V1"] = "1"
+ray.init(ignore_reinit_error=True, num_cpus=8)
 
 
 class VllmBackend(ModelBackend):

gpt_server/model_worker/base/model_worker_base.py

@@ -259,6 +259,7 @@ def run(cls):
         logger.remove(0)
         log_level = os.getenv("log_level", "WARNING")
         logger.add(sys.stderr, level=log_level)
+        os.environ["VLLM_USE_V1"] = "0"
 
         host = args.host
         controller_address = args.controller_address

gpt_server/model_worker/spark_tts.py

@@ -10,8 +10,6 @@
 
 root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 
-
-os.environ["VLLM_USE_V1"] = "1"
 import httpx
 from fastapi import HTTPException
 import base64
@@ -69,6 +67,7 @@ def __init__(
             model_type="tts",
         )
         backend = os.environ["backend"]
+        gpu_memory_utilization = float(os.getenv("gpu_memory_utilization", 0.6))
         self.engine = AutoEngine(
             model_path=model_path,
             max_length=32768,
@@ -79,6 +78,7 @@ def __init__(
             wav2vec_attn_implementation="sdpa",  # 使用flash attn加速wav2vec
             llm_gpu_memory_utilization=0.6,
             seed=0,
+            llm_gpu_memory_utilization=gpu_memory_utilization,
         )
         loop = asyncio.get_running_loop()
         # ------------- 添加声音 -------------