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}'