Hugging Face 文本生成推理（Text Generation Inference）现已支持 AWS Inferentia2

TL;DR

Hugging Face 已通过 Amazon SageMaker 在 AWS Inferentia2 上正式推出文本生成推理（TGI），提供一种无需 GPU、面向生产的方案，用于以低延迟和高并发提供大语言模型（LLM）服务。

什么是文本生成推理（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 对话

模型使用聊天模板。分词器的 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