vLLM 的 optimum-neuron 插件

optimum-neuron 包包含一个 vLLM 插件，该插件注册了一个名为 “optimum-neuron” 的 vLLM 平台，专门用于简化 Hugging Face Hub 上的模型部署到 AWS Trainium 和 Inferentia 的过程。

该平台支持两种操作模式：

它可以直接从 Hub 对预导出的 Neuron 模型进行推理，
但它也允许使用缓存的构件直接简化部署普通模型，而无需重新编译。

备注

仅缓存了给定模型所有可能配置中的一个相关子集，
你可以使用 optimum-cli 获取每个模型的所有缓存配置。
要部署未在 Hugging Face Hub 上缓存的模型，你需要事先导出它们。

设置

使用 optimum-neuron vLLM 平台最简单的方法是使用 Hugging Face Neuron 深度学习 AMI 启动一个 Amazon EC2 实例。

注意：optimum-neuron 平台尚不支持 Trn2 实例。

启动实例后，请按照连接到您的实例中的说明连接到该实例。
进入实例后，通过运行以下命令激活预安装的 optimum-neuron 虚拟环境：

source /opt/aws_neuronx_venv_pytorch_2_7/bin/activate

离线推理示例

测试模型最简单的方法是使用 Python API：

from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(model="unsloth/Llama-3.2-1B-Instruct",
          max_num_seqs=4,
          max_model_len=4096,
          tensor_parallel_size=2,
          device="neuron")

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

在线推理示例

你还可以启动一个与 OpenAI 兼容的推理服务器。

python -m vllm.entrypoints.openai.api_server \
    --model="unsloth/Llama-3.2-1B-Instruct" \
    --max-num-seqs=4 \
    --max-model-len=4096 \
    --tensor-parallel-size=2 \
    --port=8080 \
    --device "neuron"

使用以下命令测试模型：

curl 127.0.0.1:8080/v1/completions \
    -H 'Content-Type: application/json' \
    -X POST \
    -d '{"prompt":"One of my fondest memory is", "temperature": 0.8, "max_tokens":128}'