适用于 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.3.4-neuron <service_parameters>

系统参数用于映射主机和服务之间的端口、卷和设备，
服务参数转发给 `text-generation-launcher`。

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

您可以将服务部署在已经导出为 Neuron 格式的模型上，
或者，您也可以利用 Neuron 模型缓存导出自己的模型。

通用系统参数

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

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

注意：或者，为了快速进行本地测试，也可以在 `privileged` 模式下启动服务以访问所有 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。

请注意，主要限制是能够将指定 `batch_size` 的模型适配到实例上可用的总设备内存中（每个 Neuron 核心 16GB，每个设备有 2 个核心）。

查询服务

您可以使用 `/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 地址和端口。

< > 在 GitHub 上更新

文本生成推理