Hugging Face 的徽标

text-generation-inference 文档

用于 AWS Trainium 和 Inferentia 的 Neuron 后端

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

用于 AWS Trainium 和 Inferentia 的 Neuron 后端

Neuron 后端允许在 AWS Trainium 和 Inferentia 系列芯片上部署 TGI。

支持以下硬件目标

  • Trainium 1,
  • Inferentia 2。

特性

支持基本的 TGI 特性

  • 连续批处理，
  • 令牌流式传输，
  • 贪婪搜索和使用 transformers 的多项式采样。

从 Hugging Face Hub 部署服务

为特定模型部署 NeuronX TGI 服务的最简单方法是按照模型卡中的部署说明进行操作

  • 点击右侧的“部署”按钮，
  • 选择您的部署服务（支持“推理端点”和“SageMaker”），
  • 选择“AWS Trainium 和 Inferentia”，
  • 按照说明进行操作。

在专用主机上部署服务

只需运行带有两组参数的 text-generation-inference 容器即可启动服务

docker run <system_parameters> ghcr.io/huggingface/text-generation-inference:3.2.2-neuron <service_parameters>
  • 系统参数用于映射主机和服务之间的端口、卷和设备，
  • 服务参数将转发到 text-generation-launcher。

部署服务时，您将需要一个预编译的 Neuron 模型。Neuron TGI 后端支持两种主要的操作模式

  • 您可以将服务部署在已导出到 Neuron 的模型上，
  • 或者，您可以利用 Neuron 模型缓存来导出您自己的模型。

常用系统参数

无论何时启动 TGI 服务，我们都强烈建议您挂载一个共享卷，并将其作为容器中的 /data 挂载：这是模型将被缓存的位置，以加速服务的进一步实例化。

另请注意，应该使足够多的 neuron 设备对容器可见，因为每个 neuron 设备都有两个核心（因此在两个核心上部署时，您需要至少公开一个设备）。在生产环境中公开设备的推荐方法是显式使用 --device 选项（例如 --device /dev/neuron0），并根据需要公开的设备数量重复多次。

注意：或者，对于快速本地测试，也可以在特权模式下启动服务以访问所有 neuron 设备。

最后，如果您想访问门控存储库，您可能需要导出 HF_TOKEN。

这是一个仅公开第一个设备的服务实例化示例

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       -e HF_TOKEN=${HF_TOKEN} \
       ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron \
       <service_parameters>

使用来自 🤗 HuggingFace Hub 的标准模型（推荐）

我们在 aws-neuron/optimum-neuron-cache 下维护了一个 Neuron 模型缓存，其中包含最流行的架构和部署参数。

如果您只想快速尝试使用模型而无需先将其导出到 Neuron，那么在满足某些条件的情况下，这仍然是可能的

  • 您必须在启动服务时指定导出参数（或使用默认参数），
  • 模型配置必须被缓存。

以下代码段显示了如何从 hub 标准模型部署服务

export HF_TOKEN=<YOUR_TOKEN>
docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       --device=/dev/neuron1 \
       --device=/dev/neuron2 \
       --device=/dev/neuron3 \
       -e HF_TOKEN=${HF_TOKEN} \
       -e HF_AUTO_CAST_TYPE="fp16" \
       -e HF_NUM_CORES=8 \
       ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron \
       --model-id meta-llama/Meta-Llama-3-8B \
       --max-batch-size 1 \
       --max-input-length 3164 \
       --max-total-tokens 4096

使用导出到本地路径的模型

或者，您可以首先在本地将模型导出为 neuron 格式。

然后，您可以在共享卷中部署服务

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       --device=/dev/neuron1 \
       ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron \
       --model-id /data/<neuron_model_path>

注意：您无需指定任何服务参数，因为它们都将从模型导出配置中推断出来。但是，您必须公开足够的设备以匹配导出阶段指定的内核数。

使用来自 🤗 HuggingFace Hub 的 neuron 模型

在您的组织内部共享 neuron 模型的最简单方法是将其推送到 Hugging Face hub，这样就可以直接部署它，而无需导出。

以下代码段显示了如何从 hub neuron 模型部署服务

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       --device=/dev/neuron1 \
       -e HF_TOKEN=${HF_TOKEN} \
       ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron \
       --model-id <organization>/<neuron-model>

选择服务参数

使用以下命令列出可用的服务参数

docker run ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron --help

推理端点的配置始终是吞吐量和延迟之间的折衷：并行处理更多请求将允许更高的吞吐量，但会增加延迟。

neuron 模型具有静态输入维度 [batch_size, max_length]。

这为以下参数增加了一些限制

  • --max-batch-size 必须设置为 batch size，
  • --max-input-length 必须小于 max_length，
  • --max-total-tokens 必须设置为 max_length（它是按请求计算的）。

虽然不是绝对必要，但对于高效的预填充很重要

  • --max-batch-prefill-tokens 应设置为 batch_size * max-input-length。

选择正确的批大小

如上一段所述，neuron 模型静态批大小对端点延迟和吞吐量有直接影响。

有关优化提示，请参阅 text-generation-inference。

请注意，主要约束是在实例上可用的总设备内存（每个 neuron 核心 16GB，每个设备 2 个核心）中为指定的 batch_size 容纳模型。

查询服务

您可以使用 /generate 或 /generate_stream 路由查询模型

curl 127.0.0.1:8080/generate \
    -X POST \
    -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}' \
    -H 'Content-Type: application/json'
curl 127.0.0.1:8080/generate_stream \
    -X POST \
    -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}' \
    -H 'Content-Type: application/json'

注意：将 127.0.0.1:8080 替换为您的实际 IP 地址和端口。

< > 更新 on GitHub

←训练 Medusa Gaudi→