Update vllm arguments

xusenlinzy · Aug 22, 2023 · abaf4f2 · abaf4f2
1 parent 9c6fc00
commit abaf4f2
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 11 deletions.
diff --git a/api/config.py b/api/config.py
@@ -1,5 +1,5 @@
 import os
-
+from loguru import logger
 import dotenv
 
 dotenv.load_dotenv()
@@ -11,29 +11,39 @@
     'MODEL_NAME': '',
     'MODEL_PATH': '',
     'ADAPTER_MODEL_PATH': '',
+
     'DEVICE': 'cuda',
     'DEVICE_MAP': "",
     'GPUS': '',
     'NUM_GPUs': 1,
-    'QUANTIZE': 16,
+
     'EMBEDDING_NAME': '',
-    'CONTEXT_LEN': '',
+    'EMBEDDING_SIZE': '',
+    'EMBEDDING_DEVICE': 'cuda',
+
+    'QUANTIZE': 16,
     'LOAD_IN_8BIT': 'False',
     'LOAD_IN_4BIT': 'False',
     'USING_PTUNING_V2': 'False',
+
+    'CONTEXT_LEN': '',
     'STREAM_INTERVERL': 2,
     'PROMPT_NAME': '',
+
     'PATCH_TYPE': '',
     'TRAINING_LENGTH': 4096,
     'WINDOW_SIZE': 512,
+
     'API_PREFIX': '/v1',
+
     'USE_VLLM': 'False',
     'TRUST_REMOTE_CODE': "False",
     'TOKENIZE_MODE': "auto",
     'TENSOR_PARALLEL_SIZE': 1,
     'DTYPE': "half",
-    'EMBEDDING_SIZE': '',
-    'EMBEDDING_DEVICE': 'cuda',
+    "GPU_MEMORY_UTILIZATION": 0.9,
+    "MAX_NUM_BATCHED_TOKENS": 5120,
+    "MAX_NUM_SEQS": 256,
 }
 
 
@@ -61,15 +71,19 @@ def __init__(self):
         self.GPUS = get_env('GPUS')
         self.NUM_GPUs = int(get_env('NUM_GPUs'))
 
-        self.QUANTIZE = int(get_env('QUANTIZE'))
         self.EMBEDDING_NAME = get_env('EMBEDDING_NAME') if get_env('EMBEDDING_NAME') else None
-        self.CONTEXT_LEN = int(get_env('CONTEXT_LEN')) if get_env('CONTEXT_LEN') else None
+        self.EMBEDDING_SIZE = int(get_env('EMBEDDING_SIZE')) if get_env('EMBEDDING_SIZE') else None
+        self.EMBEDDING_DEVICE = get_env('EMBEDDING_DEVICE')
+
+        self.QUANTIZE = int(get_env('QUANTIZE'))
         self.LOAD_IN_8BIT = get_bool_env('LOAD_IN_8BIT')
         self.LOAD_IN_4BIT = get_bool_env('LOAD_IN_4BIT')
         self.USING_PTUNING_V2 = get_bool_env('USING_PTUNING_V2')
 
+        self.CONTEXT_LEN = int(get_env('CONTEXT_LEN')) if get_env('CONTEXT_LEN') else None
         self.STREAM_INTERVERL = int(get_env('STREAM_INTERVERL'))
         self.PROMPT_NAME = get_env('PROMPT_NAME') if get_env('PROMPT_NAME') else None
+
         self.PATCH_TYPE = get_env('PATCH_TYPE') if get_env('PATCH_TYPE') else None
         self.TRAINING_LENGTH = int(get_env('TRAINING_LENGTH'))
         self.WINDOW_SIZE = int(get_env('WINDOW_SIZE'))
@@ -81,13 +95,13 @@ def __init__(self):
         self.TOKENIZE_MODE = get_env('TOKENIZE_MODE')
         self.TENSOR_PARALLEL_SIZE = int(get_env('TENSOR_PARALLEL_SIZE'))
         self.DTYPE = get_env('DTYPE')
-
-        self.EMBEDDING_SIZE = int(get_env('EMBEDDING_SIZE')) if get_env('EMBEDDING_SIZE') else None
-        self.EMBEDDING_DEVICE = get_env('EMBEDDING_DEVICE')
+        self.GPU_MEMORY_UTILIZATION = float(get_env('GPU_MEMORY_UTILIZATION'))
+        self.MAX_NUM_BATCHED_TOKENS = int(get_env('MAX_NUM_BATCHED_TOKENS'))
+        self.MAX_NUM_SEQS = int(get_env('MAX_NUM_SEQS'))
 
 
 config = Config()
-print(f"Config: {config.__dict__}")
+logger.debug(f"Config: {config.__dict__}")
 if config.GPUS:
     if len(config.GPUS.split(",")) < config.NUM_GPUs:
         raise ValueError(

diff --git a/api/models.py b/api/models.py
@@ -69,6 +69,9 @@ def get_vllm_engine():
         trust_remote_code=config.TRUST_REMOTE_CODE,
         dtype=config.DTYPE,
         tensor_parallel_size=config.TENSOR_PARALLEL_SIZE,
+        gpu_memory_utilization=config.GPU_MEMORY_UTILIZATION,
+        max_num_batched_tokens=config.MAX_NUM_BATCHED_TOKENS,
+        max_num_seqs=config.MAX_NUM_SEQS,
     )
     engine = AsyncLLMEngine.from_engine_args(engine_args)
 
@@ -91,6 +94,7 @@ def get_vllm_engine():
     )
 
     engine_model_config = asyncio.run(engine.get_model_config())
+    engine.engine.scheduler_config.max_model_len = get_context_len(engine_model_config)
     engine.max_model_len = get_context_len(engine_model_config)
 
     return engine

diff --git a/docs/VLLM_SCRIPT.md b/docs/VLLM_SCRIPT.md
@@ -23,6 +23,7 @@ pip uninstall transformer-engine -y
 
 ### 环境变量含义
 
+
 + `MODEL_NAME`: 模型名称，如 `qwen`、`baichuan-13b-chat` 等
 
 
@@ -44,6 +45,15 @@ pip uninstall transformer-engine -y
 + `EMBEDDING_NAME`（可选项）: 嵌入模型的文件所在路径，推荐使用 `moka-ai/m3e-base` 或者 `BAAI/bge-large-zh`
 
 
++ `GPU_MEMORY_UTILIZATION`（可选项）: `GPU` 占用率
+
+
++ `MAX_NUM_BATCHED_TOKENS`（可选项）: 每个批处理的最大 `token` 数量
+
+
++ `MAX_NUM_SEQS`（可选项）: 批量大小
+
+
 ### 启动方式
 
 选择下面两种方式之一启动模型接口服务