Merge pull request #36 from MooreThreads/update_vllm_demo_script

add tp size in json

Author: gliangMT

vllm/demo/model_config.json

@@ -0,0 +1,10 @@
+{
+    "DeepSeek-R1-Distill-Qwen-1.5B": {
+        "url": "https://www.modelscope.cn/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B.git",
+        "tensor_parallel_size": 1
+    },
+    "Qwen2.5-0.5B-Instruct": {
+        "url": "https://www.modelscope.cn/Qwen/Qwen2.5-0.5B-Instruct.git",
+        "tensor_parallel_size": 1
+    }
+}

vllm/demo/model_url.json

@@ -10,22 +10,23 @@ if [ -z "$MODEL_NAME" ]; then
     exit 1
 fi
 
-# 用 Python 解析 JSON 获取 URL
-MODEL_URL=$(python3 -c "
+read MODEL_URL TENSOR_PARALLEL_SIZE <<< $(python3 -c "
 import json
 config_file = '$CONFIG_FILE'
 model_name = '$MODEL_NAME'
 with open(config_file, 'r') as f:
     data = json.load(f)
-print(data.get(model_name, ''))
+info = data.get(model_name, {})
+print(info.get('url', ''), info.get('tensor_parallel_size', ''))
 ")
 
 if [ -z "$MODEL_URL" ]; then
     echo "× $MODEL_NAME is not supported yet, please refer to the website to try other models: https://docs.mthreads.com/mtt/mtt-doc-online/compability"
     exit 1
 fi
 
-echo "√ 找到模型 URL: $MODEL_URL"
+echo "√ Find Succeed: $MODEL_URL"
+echo "√ tensor_parallel_size: $TENSOR_PARALLEL_SIZE"
 
 # 目录和日志路径
 CURRENT_DIR=$(pwd)
@@ -52,13 +53,13 @@ fi
 
 # 权重转换
 cd "${CURRENT_DIR}/.."
-./convert_weight.sh "$MODEL_DIR/$MODEL_NAME" 1
+./convert_weight.sh "$MODEL_DIR/$MODEL_NAME" $TENSOR_PARALLEL_SIZE
 
 # 启动 vLLM 服务器
 python -m vllm.entrypoints.openai.api_server \
-        --model "$CONVERTED_MODEL_DIR/$MODEL_NAME-tp1-convert" \
+        --model "$CONVERTED_MODEL_DIR/$MODEL_NAME-tp$TENSOR_PARALLEL_SIZE-convert" \
         --trust-remote-code \
-        --tensor-parallel-size 1 \
+        --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
         -pp 1 \
         --block-size 64 \
         --max-model-len 2048 \