只需 24G 显存，用 vllm 跑起来 Yi-34B 中英双语大模型

这里是 只需 24G 显存，用 vllm 跑起来 Yi-34B 中英双语大模型 的笔记哦。

引言

上次介绍了用 vllm 去部署 mixtral 8x7b 的 GPTQ 量化版本，只是跑了一下 benchmark，这次玩一下 Yi-34b 这个模型的量化版本，看看怎么使用 vllm 直接暴露出来一个 API 服务。

Mixtral 这个模型虽然很强大，但是它对中文的支持并不好，尤其是对中文的知识了解的非常有限。而 Yi-34b 这个模型则是从一开始就进行的中英双语语料进行训练。

1. 目前来看，依然在 chatbot arena 名列前茅。
2. 在 huggingface 上依然有大量的微调版本出现。（我也好想去微调个试试看！）

而且 int4 版本可以在限定上下文长度的情况下部署在 24G 显存的消费级显卡上。这里我就记录下在 autodl 进行模型部署的详细步骤。

模型下载

之前有小伙伴反馈，hf 的资源其实是下不到的，这次 yi-34b 在 modelscope 上也有提供模型，我们可以通过 modelscope 进行下载。

1. 安装 modelscope

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple modelscope

2. 下载模型

from modelscope.hub.snapshot_download import snapshot_download
model_dir = snapshot_download('01ai/Yi-34B-Chat-4bits', cache_dir='autodl-tmp', revision='master', ignore_file_pattern='.bin')

执行命令，注意使用参数 ignore_file_pattern 避免下载 bin 文件。

安装 vllm

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple vllm

开启 openai 兼容的服务

python -m vllm.entrypoints.openai.api_server \
    --model /root/autodl-tmp/01ai/Yi-34B-Chat-4bits \
    --served-model-name 01ai/Yi-34B-Chat-4bits \
    --trust-remote-code \
    --max-model-len 2048 -q awq 

测试服务是否可用：

curl http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
    "model": "01ai/Yi-34B-Chat-4bits",
    "prompt": "San Francisco is a",
    "max_tokens": 7,
    "temperature": 0
}'

执行 benchmark 测试

主要跑之前要关掉上一部开启的 api server。

下载 vllm 源码

source /etc/network_turbo
git clone https://github.com/vllm-project/vllm
cd vllm/benchmarks

执行 benchmark 测试

python benchmark_throughput.py \
    --backend vllm \
    --input-len 128 --output-len 512 \
    --model /root/autodl-tmp/01ai/Yi-34B-Chat-4bits \
    -q awq --num-prompts 100 --seed 1100 \
    --trust-remote-code \
    --max-model-len 2048

使用 gradio 的 chat 组件测试 openai 兼容的 api server

依据 vllm quickstart 的代码片段创建以下的代码：

from openai import OpenAI
import gradio as gr

# Set OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"

# 创建一个 OpenAI 客户端，用于与 API 服务器进行交互
client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)

def predict(message, history):
    # 将聊天历史转换为 OpenAI 格式
    history_openai_format = [{"role": "system", "content": "你是个靠谱的 AI 助手，尽量详细的解答用户的提问。"}]
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human })
        history_openai_format.append({"role": "assistant", "content":assistant})
    history_openai_format.append({"role": "user", "content": message})

    # 创建一个聊天完成请求，并将其发送到 API 服务器
    stream = client.chat.completions.create(
        model='01ai/Yi-34B-Chat-4bits',   # 使用的模型名称
        messages= history_openai_format,  # 聊天历史
        temperature=0.8,                  # 控制生成文本的随机性
        stream=True,                      # 是否以流的形式接收响应
        extra_body={
            'repetition_penalty': 1, 
            'stop_token_ids': [7]
        }
    )

    # 从响应流中读取并返回生成的文本
    partial_message = ""
    for chunk in stream:
        partial_message += (chunk.choices[0].delta.content or "")
        yield partial_message

# 创建一个聊天界面，并启动它，share=True 让 gradio 为我们提供一个 debug 用的域名
gr.ChatInterface(predict).queue().launch(share=True)

将文件保存为 chat.py。

首先跑起来 vllm 的 openai api server:

python -m vllm.entrypoints.openai.api_server \
    --model /root/autodl-tmp/01ai/Yi-34B-Chat-4bits \
    --served-model-name 01ai/Yi-34B-Chat-4bits \
    --trust-remote-code \
    --max-model-len 2048 -q awq 

然后开启一个新的终端执行命令:

python chat.py

稍等它在终端给我们生成一个 xxxx.gradio.live 的域名，访问这个域名就可以进行测试了。

参考链接

1. 模型的下载
2. Yi-34B-Chat-4bits
3. Welcome to vLLM!
4. How to Create a Chatbot with Gradio