跳至主要内容

使用 Cerebrium 部署

Cerebrium 是一个无服务器 AI 基础设施平台,可帮助公司更轻松地构建和部署基于 AI 的应用程序。他们提供冷启动时间短的无服务器 GPU,拥有超过 12 种不同的 GPU 芯片,支持自动扩缩容,且您只需为您使用的计算付费。

设置 Cerebrium

首先,我们安装并登录 Cerebrium 进行身份验证。

pip install cerebrium
cerebrium login

然后我们创建第一个项目

cerebrium init mistral-vllm

设置环境和硬件

您可以在上面使用 init 函数创建的 cerebrium.toml 文件中设置环境和硬件。在这里我们使用 Ampere A10 GPU 等。您可以在此处阅读更多内容

[cerebrium.deployment]
name = "mistral-vllm"
python_version = "3.11"
docker_base_image_url = "debian:bookworm-slim"
include = "[./*, main.py, cerebrium.toml]"
exclude = "[.*]"

[cerebrium.hardware]
cpu = 2
memory = 14.0
compute = "AMPERE_A10"
gpu_count = 1
provider = "aws"
region = "us-east-1"

[cerebrium.dependencies.pip]
sentencepiece = "latest"
torch = ">=2.0.0"
vllm = "latest"
transformers = ">=4.35.0"
accelerate = "latest"
xformers = "latest"

设置推理

在 Cerebrium 中运行代码就像编写普通 Python 代码一样,没有特殊语法。在您的 main.py 中指定以下内容

from vllm import LLM, SamplingParams
from huggingface_hub import login
from cerebrium import get_secret

# Your huggingface token (HF_AUTH_TOKEN) should be stored in your project secrets on your Cerebrium dashboard
login(token=get_secret("HF_AUTH_TOKEN"))

# Initialize the model
llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.3", dtype="bfloat16", max_model_len=20000, gpu_memory_utilization=0.9)

由于使用 Mistral 模型需要身份验证,我们需要将 Hugging Face 令牌添加到我们的Cerebrium Secrets中。请确保您添加的 Huggingface 令牌具有 WRITE 权限。首次部署时,它将下载模型并将其存储在磁盘上,因此后续调用将从磁盘加载模型。

将以下内容添加到您的 main.py 中

def run(prompt: str, temperature: float = 0.8, top_p: float = 0.75, top_k: int = 40, max_tokens: int = 256, frequency_penalty: int = 1):

sampling_params = SamplingParams(
temperature=temperature,
top_p=top_p,
top_k=top_k,
max_tokens=max_tokens,
frequency_penalty=frequency_penalty
)

outputs = llm.generate([item.prompt], sampling_params)

generated_text = []
for output in outputs:
generated_text.append(output.outputs[0].text)

return {"result": generated_text}

Cerebrium 中的每个函数都可以通过 API 端点调用。最顶层(即不在函数中)的代码只在容器首次启动时实例化,因此对于后续调用,它将简单地运行您调用的函数中定义的代码。

我们最终的 main.py 应该如下所示

from vllm import LLM, SamplingParams
from huggingface_hub import login
from cerebrium import get_secret

# Your huggingface token (HF_AUTH_TOKEN) should be stored in your project secrets on your Cerebrium dashboard
login(token=get_secret("HF_AUTH_TOKEN"))

# Initialize the model
llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.3", dtype="bfloat16", max_model_len=20000, gpu_memory_utilization=0.9)

def run(prompt: str, temperature: float = 0.8, top_p: float = 0.75, top_k: int = 40, max_tokens: int = 256, frequency_penalty: int = 1):

sampling_params = SamplingParams(
temperature=temperature,
top_p=top_p,
top_k=top_k,
max_tokens=max_tokens,
frequency_penalty=frequency_penalty
)

outputs = llm.generate([item.prompt], sampling_params)

generated_text = []
for output in outputs:
generated_text.append(output.outputs[0].text)

return {"result": generated_text}

在云端运行

cerebrium deploy

您将看到您的应用程序部署、安装 pip 包并下载模型。完成后,它将输出一个可用于调用端点的 CURL 请求。请记住,URL 结尾必须是您想调用的函数 - 在本例中为 /run。

curl --location --request POST 'https://api.cortex.cerebrium.ai/v4/p-<YOUR PROJECT ID>/mistral-vllm/run' \
--header 'Authorization: Bearer <YOUR TOKEN HERE>' \
--header 'Content-Type: application/json' \
--data-raw '{
"prompt: "What is the capital city of France?"
}'

您应该会看到类似这样的消息

{
"run_id": "nZL6mD8q66u4lHTXcqmPCc6pxxFwn95IfqQvEix0gHaOH4gkHUdz1w==",
"message": "Finished inference request with run_id: `nZL6mD8q66u4lHTXcqmPCc6pxxFwn95IfqQvEix0gHaOH4gkHUdz1w==`",
"result": {
"result": ["\nA: Paris"]
},
"status_code": 200,
"run_time_ms": 151.24988555908203
}