fix #2782 #2798

2024-03-12 15:53:29 +08:00
parent c901aa63ff
commit 07f9b754a7
3 changed files with 19 additions and 2 deletions
--- a/src/llmtuner/chat/vllm_engine.py
+++ b/src/llmtuner/chat/vllm_engine.py
@@ -31,9 +31,11 @@ class VllmEngine(BaseEngine):
            model=model_args.model_name_or_path,
            trust_remote_code=True,
            max_model_len=model_args.vllm_maxlen,
-            tensor_parallel_size=get_device_count(),
+            tensor_parallel_size=get_device_count() or 1,
+            gpu_memory_utilization=model_args.vllm_gpu_util,
            disable_log_stats=True,
            disable_log_requests=True,
+            enforce_eager=model_args.vllm_enforce_eager,
        )
        self.model = AsyncLLMEngine.from_engine_args(engine_args)
        self.tokenizer = load_tokenizer(model_args)