AWS Trainium & Inferentia 文档

最佳 Neuron 容器

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

最佳 Neuron 容器

我们为 Amazon SageMaker 提供预构建的 Optimum Neuron 容器。这些容器预装了所有 Hugging Face 库和依赖项,因此您可以立即开始使用它们。我们提供用于训练和推理的容器,以及使用 TGI 优化的文本生成容器。该表格是最新的,仅包含每个容器的最新版本。您可以在 深度学习容器发行说明 中找到旧版本

我们建议使用 sagemaker Python SDK 来检索您要使用的容器的镜像 URI。以下是一个代码片段,用于检索最新的文本生成推理容器镜像 URI

from sagemaker.huggingface import get_huggingface_llm_image_uri
 
# retrieve the llm image uri
llm_image = get_huggingface_llm_image_uri(
  "huggingface-neuronx"
)

print(f"llm image uri: {llm_image}")

可用的 Optimum Neuron 容器

类型 Optimum 版本 镜像 URI
训练 0.0.24 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training-neuronx:2.1.2-transformers4.41.1-neuronx-py310-sdk2.19.1-ubuntu20.04
推理 0.0.24 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference-neuronx:2.1.2-transformers4.41.1-neuronx-py310-sdk2.19.1-ubuntu20.04
文本生成推理 0.0.24 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.2-optimum0.0.24-neuronx-py310-ubuntu22.04

请将 763104351884 替换为正确的 AWS 账户 ID,并将 region 替换为您正在使用的 AWS 区域。