diff --git a/vllm/run.sh b/vllm/run.sh index 97a52fc6..e19f38d1 100644 --- a/vllm/run.sh +++ b/vllm/run.sh @@ -1,9 +1,14 @@ MODEL=Meta-Llama-3-8B-Instruct +if [ -z "$1" ]; then + echo "Error: No argument provided. Please specify a model name." + exit 1 +fi + if [ "$1" = "llama" ]; then - CUDA_VISIBLE_DEVICES=0,1 python -m vllm.entrypoints.openai.api_server \ + CUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server \ --model="/hdd/zwj/models/meta-llama/$MODEL" \ - --tensor-parallel-size 2 \ + --tensor-parallel-size 1 \ --trust-remote-code \ --device auto \ --gpu-memory-utilization 0.98 \ @@ -28,4 +33,4 @@ fi # kv-cache-dtype kv 缓存存储的数据类型。如果为“auto”,则将使用模型默认的数据类型。CUDA 11.8及以上版本 支持 fp8 (=fp8_e4m3) 和 fp8_e5m2。ROCm (AMD GPU) 支持 fp8 (=fp8_e4m3) # served-model-name 对外提供的API中的模型名称 # host 监听的网络地址,0.0.0.0表示所有网卡的所有IP,127.0.0.1表示仅限本机 -# port API服务的端口 \ No newline at end of file +# port API服务的端口 diff --git a/vllm/test_vllm.py b/vllm/test_vllm.py new file mode 100644 index 00000000..582ae272 --- /dev/null +++ b/vllm/test_vllm.py @@ -0,0 +1,18 @@ +from openai import OpenAI +# Set OpenAI's API key and API base to use vLLM's API server. +openai_api_key = "EMPTY" +openai_api_base = "http://localhost:8080/v1" + +client = OpenAI( + api_key=openai_api_key, + base_url=openai_api_base, +) + +chat_response = client.chat.completions.create( + model="llama", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Tell me a joke."}, + ] +) +print("Chat response:", chat_response) \ No newline at end of file