NeuronX 文本生成推理 (适用于 AWS Inferentia2)

文本生成推理（TGI）是一个用于部署和服务大型语言模型 (LLM) 的工具包。

它适用于 Inferentia2。

特性

支持基本的 TGI 特性

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

许可证

NeuronX TGI 在 Apache2 许可证下发布。

从 Hugging Face Hub 部署服务

部署特定模型的 NeuronX TGI 服务最简单的方法是按照模型卡片中的部署说明

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

在专用主机上部署服务

该服务只需使用两组参数运行 neuronx-tgi 容器即可启动

docker run <system_parameters> ghcr.io/huggingface/neuronx-tgi:latest <service_parameters>

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

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

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

常见系统参数

每次启动 TGI 服务时，我们强烈建议您挂载一个共享卷，并在容器中将其挂载为 /data：模型将缓存到此处，以加快服务的后续实例化速度。

另请注意，容器应能够看到足够的 Neuron 设备。实现此目的最简单的方法是以 privileged 模式启动服务，以访问所有 Neuron 设备。或者，可以使用 --device 选项显式公开每个设备。

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

以下是一个服务实例化的示例

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --privileged \
       -e HF_TOKEN=${HF_TOKEN} \
       ghcr.io/huggingface/neuronx-tgi:latest \
       <service_parameters>

如果您只想映射第一个设备，则启动命令变为

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       -e HF_TOKEN=${HF_TOKEN} \
       ghcr.io/huggingface/neuronx-tgi:latest \
       <service_parameters>

使用 🤗 Hugging Face Hub 中的标准模型（推荐）

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

如果您只想快速尝试使用尚未导出的模型的服务，则仍然可以动态导出它，但需要满足一些条件

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

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

export HF_TOKEN=<YOUR_TOKEN>
docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --privileged \
       -e HF_TOKEN=${HF_TOKEN} \
       -e HF_AUTO_CAST_TYPE="fp16" \
       -e HF_NUM_CORES=2 \
       ghcr.io/huggingface/neuronx-tgi:latest \
       --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 \
       --privileged \
       ghcr.io/huggingface/neuronx-tgi:latest \
       --model-id /data/<neuron_model_path>

注意：您无需指定任何服务参数，因为它们都将从模型导出配置中推断出来。

使用 🤗 Hugging Face Hub 中的 Neuron 模型

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

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

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --privileged \
       -e HF_TOKEN=${HF_TOKEN} \
       ghcr.io/huggingface/neuronx-tgi:latest \
       --model-id <organization>/<neuron-model>

选择服务参数

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

docker run ghcr.io/huggingface/neuronx-tgi --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 地址和端口。

AWS Trainium 和 Inferentia