compatible with vllm==0.3.* (#396)

Author: Jintao-Huang

swift/llm/utils/vllm_utils.py

@@ -8,6 +8,7 @@
 from modelscope import GenerationConfig, snapshot_download
 from torch import dtype as Dtype
 from tqdm import tqdm
+from transformers import PreTrainedTokenizerBase
 from vllm import (AsyncEngineArgs, AsyncLLMEngine, EngineArgs, LLMEngine,
                   SamplingParams)
 
@@ -20,6 +21,13 @@
 logger = get_logger()
 
 
+def _get_vllm_tokenizer(vllm_engine: LLMEngine) -> PreTrainedTokenizerBase:
+    tokenizer = vllm_engine.tokenizer
+    if not isinstance(tokenizer, PreTrainedTokenizerBase):
+        tokenizer = tokenizer.tokenizer
+    return tokenizer
+
+
 def get_vllm_engine(model_type: str,
                     torch_dtype: Optional[Dtype] = None,
                     *,
@@ -89,7 +97,11 @@ def get_vllm_engine(model_type: str,
     llm_engine.engine_args = engine_args
     llm_engine.model_dir = model_dir
     llm_engine.model_type = model_type
-    llm_engine.tokenizer = tokenizer
+    if isinstance(llm_engine.tokenizer, PreTrainedTokenizerBase):
+        llm_engine.tokenizer = tokenizer
+    else:
+        # compatible with vllm==0.3.*
+        llm_engine.tokenizer.tokenizer = tokenizer
     generation_config_path = os.path.join(model_dir, 'generation_config.json')
     if os.path.isfile(generation_config_path):
         generation_config = GenerationConfig.from_pretrained(model_dir)
@@ -330,7 +342,7 @@ def prepare_vllm_engine_template(
         max_model_len=args.max_model_len,
         use_async=use_async,
         **kwargs)
-    tokenizer = llm_engine.tokenizer
+    tokenizer = _get_vllm_tokenizer(llm_engine)
     if use_async:
         model_config = asyncio.run(llm_engine.get_model_config())
         llm_engine.model_config = model_config

tests/llm/test_run.py

@@ -208,7 +208,7 @@ def test_self_cognition(self):
             return
         for dataset in [[], [DatasetName.alpaca_zh, DatasetName.alpaca_en]]:
             sft_args = SftArguments(
-                model_type=ModelType.qwen_7b_chat,
+                model_type=ModelType.qwen1half_1_8b_chat_int4,
                 dataset=dataset,  # no dataset
                 train_dataset_sample=100,
                 dtype='fp16',