新增vllm_args配置，更新测试模型参数配置，允许自定义测试数据路径。

settings.template.jsonc

@@ -78,5 +78,11 @@
         "temperature": 0.5,
         "max_length": 50,
         "top_p": 0.65
+    },
+    "vllm_args": {
+        "gpu_memory_utilization": 0.9
+    },
+    "test_model_args": {
+        "test_data_path": "dataset/test_data.json"
     }
 }

tests/full_pipe.jsonc

@@ -69,5 +69,8 @@
         "temperature": 0.5,
         "max_length": 50,
         "top_p": 0.65
+    },
+    "vllm_args": {
+        "gpu_memory_utilization": 0.90
     }
 }

weclone/core/inference/offline_infer.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Union, cast
 
 from llamafactory.data import get_template_and_fix_tokenizer
 from llamafactory.extras.misc import get_device_count
@@ -9,6 +9,9 @@ from vllm import LLM, SamplingParams
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import GuidedDecodingParams
 
+from weclone.utils.config import load_config
+from weclone.utils.config_models import VllmArgs
+
 # 这里不需要写太好，transforms库后续更新自带vllm
 
 
@@ -41,6 +44,7 @@ def vllm_infer(
     if pipeline_parallel_size > get_device_count():
         raise ValueError("Pipeline parallel size should be smaller than the number of gpus.")
 
+    wc_vllm_args = cast(VllmArgs, load_config("vllm"))
     model_args, data_args, _, generating_args = get_infer_args(
         {
             "model_name_or_path": model_name_or_path,
@@ -99,7 +103,7 @@ def vllm_infer(
         "disable_log_stats": True,
         "enable_lora": model_args.adapter_name_or_path is not None,
         "enable_prefix_caching": True,  # 是否启用前缀缓存
-        "gpu_memory_utilization": 0.95,
+        "gpu_memory_utilization": wc_vllm_args.gpu_memory_utilization,
         # "quantization": "bitsandbytes", # 是否启用vllm的 bitsandbytes 的量化加载
         # "load_format": "bitsandbytes",
     }

weclone/utils/config.py

@@ -58,6 +58,9 @@ def create_config_by_arg_type(arg_type: str, wc_config: WcConfig) -> BaseModel:
         config_dict = {**common_config, **wc_config.infer_args.model_dump()}
         return WCInferConfig(**config_dict)
 
+    elif arg_type == "vllm":
+        return wc_config.vllm_args
+
     elif arg_type == "test_model":
         return wc_config.test_model_args
 

weclone/utils/config_models.py

@@ -164,8 +164,12 @@ class InferArgs(BaseModel):
     max_length: int = Field(..., description="最大生成长度")
 
 
+class VllmArgs(BaseModel):
+    gpu_memory_utilization: float = Field(default=0.9, description="vllm GPU内存利用率")
+
+
 class TestModelArgs(BaseModel):
-    test_data_path: str = Field("dataset/test_data.json", description="测试数据路径")
+    test_data_path: str = Field(default="dataset/test_data.json", description="测试数据路径")
 
 
 class WcConfig(BaseModel):
@@ -175,7 +179,8 @@ class WcConfig(BaseModel):
     make_dataset_args: MakeDatasetArgs = Field(..., description="数据处理参数")
     train_sft_args: TrainSftArgs = Field(..., description="SFT微调参数")
     infer_args: InferArgs = Field(..., description="推理参数")
-    test_model_args: TestModelArgs = TestModelArgs(test_data_path="dataset/test_data.json")
+    vllm_args: VllmArgs = Field(VllmArgs())
+    test_model_args: TestModelArgs = Field(TestModelArgs())
 
 
 class WCInferConfig(CommonArgs, InferArgs):