Hugging Face 文本生成推理現在支援 AWS Inferentia2

TL;DR

Hugging Face 已在 AWS Inferentia2 上透過 Amazon SageMaker 正式推出 Text Generation Inference（TGI），提供一條無需 GPU、具備生產等級的路徑，以低延遲與高併發度服務大型語言模型（LLM）。

什麼是 Text Generation Inference（TGI）？

TGI 是為 LLM 專門打造的推理伺服器，提供以下功能：

  • 張量平行化與持續批次處理，以實現高吞吐量。
  • 最佳化支援流行的開源模型，如 Llama、Mistral 等。
  • 具備生產級部署，已被 Grammarly、Uber 與 Deutsche Telekom 等公司使用。

為何 Inferentia2 重要

AWS Inferentia2 晶片提供比 GPU 更低成本的高效能推理。透過將 TGI 與 Inferentia2 及 SageMaker 整合，Hugging Face 提供以下優勢：

  • 無縫、受管的部署體驗。
  • 相容於支援 HuggingChat、OpenAssistant 與 Hugging Face 無伺服器端點的相同後端技術。
  • 為偏好 AWS 專屬硬體的客戶提供替代的運算堆疊。

在 Inferentia2 上部署 Zephyr 7B – 步驟指南

此部落格文章說明如何在 ml.inf2.8xlarge 執行個體上部署 Zephyr 7B 模型（Mistral‑7B‑v0.1 的 DPO 微調版本）。步驟完整且任何 SageMaker 使用者皆可重現。

1. 設定開發環境

pip install transformers "sagemaker>=2.206.0" --upgrade --quiet

此腳本會取得 IAM 執行角色、建立 SageMaker 工作階段，並印出角色 ARN 與區域。

2. 取得 TGI NeuronX 容器映像

from sagemaker.huggingface import get_huggingface_llm_image_uri
llm_image = get_huggingface_llm_image_uri(
    "huggingface-neuronx",
    version="0.0.20"
)
print(f"llm image uri: {llm_image}")

撰寫本文時，最新的 DLC 版本尚未透過輔助函式公開，故使用原始的 ECR URI。

3. 編譯或取得 Inferentia2 的快取模型

Inferentia2 不支援動態形狀，因此序列長度與批次大小必須在編譯時固定。Hugging Face 提供 neuron model cache，內含預先編譯的設定（例如 Mistral‑7B、Zephyr‑7B）。若缺少所需的設定，使用者可透過 Optimum CLI 進行編譯：

optimum-cli export neuron -m HuggingFaceH4/zephyr-7b-beta \
    --batch_size 4 --sequence_length 2048 \
    --num_cores 2 --auto_cast_type bf16 ./zephyr-7b-beta-neuron

編譯完成的產物會推送至 Hub，路徑為 aws-neuron/zephyr-7b-seqlen-2048-bs-4-cores-2。

4. 定義端點設定並部署

TGI NeuronX 容器的關鍵環境變數包括：

  • HF_MODEL_ID, HF_NUM_CORES, HF_BATCH_SIZE, HF_SEQUENCE_LENGTH, HF_AUTO_CAST_TYPE
  • MAX_BATCH_SIZE, MAX_INPUT_LENGTH, MAX_TOTAL_TOKENS
from sagemaker.huggingface import HuggingFaceModel
config = {
    "HF_MODEL_ID": "HuggingFaceH4/zephyr-7b-beta",
    "HF_NUM_CORES": "2",
    "HF_BATCH_SIZE": "4",
    "HF_SEQUENCE_LENGTH": "2048",
    "HF_AUTO_CAST_TYPE": "bf16",
    "MAX_BATCH_SIZE": "4",
    "MAX_INPUT_LENGTH": "1512",
    "MAX_TOTAL_TOKENS": "2048",
}
llm_model = HuggingFaceModel(role=role, image_uri=llm_image, env=config)
llm = llm_model.deploy(initial_instance_count=1, instance_type="ml.inf2.8xlarge", container_startup_health_check_timeout=1800)

部署通常需要 10–15 分鐘。

5. 執行推理並與 Zephyr 7B 對話

模型使用聊天模板。Tokenizer 的 apply_chat_template 方法會將 OpenAI 風格的訊息字典轉換為所需的提示格式。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("aws-neuron/zephyr-7b-seqlen-2048-bs-4-cores-2")
messages = [
    {"role": "system", "content": "You are the AWS expert"},
    {"role": "user", "content": "Can you tell me an interesting fact about AWS?"},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
payload = {
    "do_sample": True,
    "top_p": 0.6,
    "temperature": 0.9,
    "top_k": 50,
    "max_new_tokens": 256,
    "repetition_penalty": 1.03,
    "return_full_text": False,
    "stop": ["</s>"]
}
chat = llm.predict({"inputs": prompt, "parameters": payload})
print(chat[0]["generated_text"][len(prompt):])

回應證實模型已在 Inferentia2 上正常運作。

6. 清理資源

llm.delete_model()
llm.delete_endpoint()

影響與未來工作

  • 具成本效益的擴展：相較於 GPU 叢集，Inferentia2 提供更廉價的運算選項，以支援高吞吐量的 LLM 服務。
  • 模型覆蓋：目前的快取包含 Llama、Mistral 與 Zephyr；Hugging Face 計畫擴充支援的架構。
  • 編譯工作流程：使用者可依賴預先快取的建置，或使用 Optimum 編譯自訂設定（約需 45 分鐘）。
  • 路線圖：持續的工作旨在支援更多模型、改進快取系統，並簡化編譯流程。

如有問題，請於 Twitter 或 LinkedIn 聯絡 Philipp Schmid。

Sources