text-generation-inference Neuron 后端部署指南在 AWS Trainium 与 Inferentia 上运行 TGI【免费下载链接】text-generation-inferenceLarge Language Model Text Generation Inference项目地址: https://gitcode.com/GitHub_Trending/te/text-generation-inference本指南以 TGI 仓库的 Neuron 后端文档docs/source/backends/neuron.md为主线系统讲解如何在 AWS Trainium 1 与 Inferentia 2 芯片上部署 text-generation-inference 服务。你将掌握从 Hugging Face Hub 直接拉起服务、使用 Neuron Model Cache 免导出运行标准模型、部署本地已导出 Neuron 模型以及正确配置 batch size、输入长度等静态维度参数的完整实战方案。一、Neuron 后端概述与硬件支持Neuron 后端是 TGI 针对 AWS Trainium 与 Inferentia 系列芯片专门实现的推理后端目标是把 TGI 的完整服务能力HTTP 路由、动态批处理调度与 AWS Neuron SDK 的编译执行能力结合起来。当前仓库支持的硬件目标如下Trainium 1trn1 / trn1n 实例Inferentia 2inf2 实例。从仓库结构看Neuron 后端由三部分组成见 backends/neuron/README.mdAWS Neuron SDKneuronx-cc、torch-neuronx、neuronx-distributed、libneuronxla 等运行时与编译器TGI v2 的 launcher 与 router负责 HTTP 入口、请求队列、连续批处理调度Neuron 专用推理服务器位于 backends/neuron/server 下的 Python gRPC 服务负责模型的 prefill/decode 与 token 选择。这种三层结构决定了其部署形态用户通过 Docker 镜像启动容器容器内的 entrypoint 先执行 tgi_entry_point.py 推导出正确的服务环境变量再启动text-generation-launcher拉起 router 与 Python 推理服务器。支持的功能特性Neuron 后端支持 TGI 的基础功能连续批处理continuous batching新请求可以在解码过程中随时插入无需等整个批次结束Token 流式输出token streaming通过/generate_stream逐 token 返回结果贪婪搜索与多项式采样复用 transformers 的生成策略transformers 生成策略文档。从源码看这些能力分别由 generator.py 中的NeuronGenerator类实现prefill()负责新请求的接入与 slot 分配decode()负责对已 prefill 的批次逐 token 解码而 token 选择逻辑贪婪或采样通过TokenSelector来自optimum.neuron.generation完成。仓库测试 test_continuous_batching.py 专门验证了请求 1 解码到一半时插入请求 2两者最终输出一致的连续批处理行为可以作为功能正确性的参考依据。二、从 Hugging Face Hub 部署服务模型卡方式最简单的部署路径是直接利用模型卡片上提供的部署入口打开目标模型页面点击右侧的Deploy按钮选择部署服务支持Inference Endpoints与SageMaker硬件类型选择AWS Trainium Inferentia按页面提示完成后续配置。这种方式由托管平台负责 Neuron 设备挂载、镜像拉取等基础设施细节适合快速验证或生产托管场景。三、在专用主机上部署服务Docker 方式在自有 EC2 实例上部署时直接运行 Neuron 版的 TGI 容器即可。容器启动命令分为两组参数docker run system_parameters ghcr.io/huggingface/text-generation-inference:VERSION-neuron service_parameterssystem parameters负责宿主机与容器之间的端口、数据卷、Neuron 设备映射service parameters透传给text-generation-launcher即 TGI 的服务参数如--model-id、--max-batch-size等。注意VERSION需替换为具体的 TGI 版本号例如3.3.5-neuron。部署前需要明确一点Neuron 模型是预编译产物。Neuron 后端支持两种模型来源直接部署已导出为 Neuron 格式的模型借助Neuron Model Cache在启动时自动导出你自己的普通模型。3.1 通用系统参数system parameters无论采用哪种模型来源以下系统级配置都强烈建议遵循共享数据卷挂载到/data容器内/data是模型缓存目录挂载共享卷可以显著加速后续服务实例的启动避免重复下载/导出正确暴露 Neuron 设备每个 Neuron 设备包含2 个 core因此要在 2 个 core 上部署至少需要暴露 1 个设备。生产环境推荐用--device选项显式、逐个指定设备例如--device /dev/neuron0有几个设备就重复几次快速本地测试可选用privileged模式启动容器以一次性访问全部 Neuron 设备不推荐生产使用HF_TOKEN访问 gated 仓库如 Llama 系列时需要通过-e HF_TOKEN${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_parameters3.2 使用 Hub 上的标准模型推荐借助 Neuron Model Cache我们维护了一个覆盖主流架构与常用部署参数的Neuron Model Cache位于aws-neuron/optimum-neuron-cache。因此即便模型没有预先导出为 Neuron 格式也可以快速拉起服务但需满足两个前提启动时必须指定导出参数或使用默认参数目标模型的配置必须已在缓存中。以下示例展示了如何从 Hub 标准模型直接部署meta-llama/Meta-Llama-3-8B8 个 core、fp16 精度export HF_TOKENYOUR_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_TYPEfp16 \ -e HF_NUM_CORES8 \ 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关于HF_AUTO_CAST_TYPE与HF_NUM_CORES这两个环境变量从 tgi_env.py 可以看出它们是 Neuron 后端服务器侧的关键配置HF_NUM_CORES指定编译/运行所需的 Neuron core 数量即张量并行度tp_degreeHF_AUTO_CAST_TYPE指定模型权重与激活的自动转换精度类型如fp16。这两个变量连同 router 侧的MAX_BATCH_SIZE、MAX_TOTAL_TOKENS、MAX_INPUT_TOKENS、MAX_BATCH_PREFILL_TOKENS共同构成 Neuron 后端完整的环境变量集合。parse_cmdline_and_set_env()会把命令行参数转换为对应的环境变量而当模型配置不完整时model.py 的get_export_kwargs_from_env()会从这些环境变量中提取导出参数batch_size、sequence_length、num_cores、auto_cast_type配合optimum-neuron在启动时完成模型导出。3.3 使用已导出到本地路径的模型如果已经用optimum-cli把模型导出为 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注意此时无需指定任何服务参数所有参数都会从模型导出配置中推导出来但必须暴露足够数量的设备以匹配导出阶段指定的 core 数量。这一推导逻辑的实现位于 tgi_env.py 的neuron_config_to_env()它会读取模型的 NeuronConfig含batch_size、sequence_length、tp_degree、auto_cast_type将其写入环境变量文件MAX_BATCH_SIZE、MAX_TOTAL_TOKENS、HF_NUM_CORES、HF_AUTO_CAST_TYPE若未显式指定则自动令MAX_INPUT_TOKENS sequence_length // 2、MAX_BATCH_PREFILL_TOKENS batch_size × MAX_INPUT_TOKENS并写入同一文件。该文件随后由 tgi-entrypoint.sh 通过source加载再启动text-generation-launcher。3.4 使用 Hub 上的 Neuron 模型若模型已以 Neuron 格式推送到 Hugging Face Hub即可在组织内直接共享并部署无需任何导出步骤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从 model.py 的fetch_model()流程看服务端会先尝试读取模型的NeuronConfig如果能读到说明是 Neuron 模型直接snapshot_download拉取忽略*.bin权重文件读不到则走 Neuron Model Cache 的兼容条目查找逻辑找不到兼容缓存时报错并提示去请求缓存或自行导出。3.5 兼容性校验机制源码视角Neuron 后端在加载模型前会做严格的环境兼容性检查见 tgi_env.py 的check_env_and_neuron_config_compatibility()主要校验项包括校验项规则core 数量模型tp_degree不得超过本机可用 core 数available_cores编译器版本本地neuronxcc版本必须与模型编译时版本一致Hub 缓存条目强制校验MAX_BATCH_SIZE不得超过 Neuron 配置中的batch_sizeMAX_TOTAL_TOKENS不得超过 Neuron 配置中的sequence_lengthHF_NUM_CORES不得超过 Neuron 配置中的tp_degreeHF_AUTO_CAST_TYPE必须与模型配置中的auto_cast_type/torch_dtype一致MAX_INPUT_TOKENS必须小于模型序列长度这解释了为什么文档反复强调服务参数必须与模型导出配置对齐——任何一项不匹配都会导致启动失败或加载报错。四、服务参数选择静态输入维度的约束查看可用服务参数列表可以在不带任何服务参数的情况下运行docker run ghcr.io/huggingface/text-generation-inference:VERSION-neuron --help推理端点的配置本质上是吞吐量与延迟之间的权衡并行处理的请求越多吞吐越高但延迟也会随之增加。4.1 Neuron 模型的静态维度特性与 GPU 后端不同Neuron 模型的输入维度是静态的[batch_size, max_length]——这是由编译期固定下来的无法在运行时动态调整。这给以下参数带来了硬性约束--max-batch-size必须等于模型编译时的batch size--max-input-length必须小于模型编译时的max_length--max-total-tokens必须等于模型编译时的max_length按单请求计。此外虽然不是严格必需但为了高效的 prefill预填充阶段建议将--max-batch-prefill-tokens设置为batch_size × max-input-length。这些约束在源码中有直接印证NeuronGenerator.warmup()generator.py会校验请求批次大小不得超过模型neuron_config.batch_size否则抛出错误并提示请确保编译模型的 batch_size 与传入 TGI 的 batch_size 一致通常通过环境变量MAX_BATCH_SIZE设置。而prefill()在分配 slot 时如果空 slot 数量不足以容纳新请求也会要求把max_batch_size与静态 batch size 对齐。从 tgi_env.py 的tgi_router_env_vars可以看到这四个参数在 router 侧对应的环境变量分别是服务参数环境变量--max-batch-sizeMAX_BATCH_SIZE--max-total-tokensMAX_TOTAL_TOKENS--max-input-tokensMAX_INPUT_TOKENS兼容旧名MAX_INPUT_LENGTH--max-batch-prefill-tokensMAX_BATCH_PREFILL_TOKENSparse_cmdline_and_set_env()的逻辑是命令行参数优先仅当命令行未指定且环境变量已存在时才使用环境变量值并在设置完成后统一写入os.environ保证 router 与 server 拿到一致的值。4.2 如何选择合适的 batch size如前所述Neuron 模型静态 batch size 会直接影响端点的延迟与吞吐因此选择时必须同时考虑性能目标与内存容量两个维度内存约束是硬性下限必须在实例可用的总设备内存内装下指定batch_size的模型。每个 Neuron core 提供16GB 显存每个设备有2 个 core即每个设备共 32GB。例如inf2.24xlarge这类实例的总可用内存由实例规格决定选型前请先核对型号规格性能权衡更大的 batch size 提高吞吐但增加单请求延迟更小的 batch size 反之。建议结合自身场景在线低延迟 or 离线高吞吐在实测中调优。五、查询服务/generate 与 /generate_stream服务启动后可以通过 TGI 的标准 HTTP 路由查询模型非流式生成/generatecurl 127.0.0.1:8080/generate \ -X POST \ -d {inputs:What is Deep Learning?,parameters:{max_new_tokens:20}} \ -H Content-Type: application/json流式生成/generate_streamcurl 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 地址与端口。流式输出的底层实现在 generator.py 的Slot._decode_next_tokens()中它会以 token 偏移为游标只解码新增 token 对应的文本片段并处理 UTF-8 多字节字符可能被截断的情况检测到不完整的字节序列时暂不返回等待下一个 token 补齐从而保证中文、emoji 等内容的流式输出不会出现乱码。仓库测试 test_generator_slot.py 用空格文本、中文、emoji 三类用例验证了这一解码逻辑。六、服务端架构与请求处理流程源码视角6.1 启动链路Neuron 容器的启动链路清晰可循Dockerfile.neuron、tgi-entrypoint.sh容器ENTRYPOINT指向 tgi-entrypoint.sh脚本调用 tgi_entry_point.py解析命令行参数与环境变量若所有tgi_env_vars均已设置则直接跳过否则读取模型的 NeuronConfig 或查找 Hub 缓存中的兼容配置通过neuron_config_to_env()把推导结果写入临时环境变量文件脚本source该环境变量文件exec text-generation-launcher $启动 TGI launcherrouterlauncher 拉起 Python 推理服务器server.py以 gRPC 协议在unix://uds_path-0上提供服务。6.2 gRPC 服务与请求处理server.py 中的TextGenerationService实现了 TGI 的 gRPC 协议proto 定义见 proto/generate.proto 与 proto/v3/generate.proto核心 RPC 包括Info/Health服务元信息与健康检查Prefill新请求的预填充首次前向生成第一个 tokenDecode基于 KV cache 逐 token 解码FilterBatch/ClearCache批次过滤与缓存清理Warmup启动时验证硬件能否支撑目标负载返回模型支持的最大 token 数。所有异常统一由 interceptor.py 的ExceptionInterceptor捕获并转为 gRPC INTERNAL 状态码返回。NeuronGenerator内部用固定数量的Slot数量等于静态batch_size表示批次槽位EMPTY空闲、READY就绪、PAUSE暂停KV cache 仍会填充prefill 时把新请求分配到空闲 slotdecode 期间通过暂停/恢复机制实现连续批处理。七、构建自定义 Neuron 镜像可选若需定制镜像例如修改 Neuron SDK 版本或预装依赖可以使用仓库提供的构建入口make -C backends/neuron image该目标backends/neuron/Makefile会基于 Dockerfile.neuron 构建text-generation-inference:VERSION-neuron镜像并额外打上latest-neuron标签。镜像内部固定安装了aws-neuronx-dkms、torch-neuronx、neuronx-cc、optimum-neuron等 Neuron 软件栈并在base阶段安装 CPU 版 PyTorch 以避免引入 CUDA 依赖。八、部署注意事项小结版本匹配容器 tag 使用ghcr.io/huggingface/text-generation-inference:VERSION-neuron请按实际发布版本替换VERSION设备数量每个 Neuron 设备 2 个 core暴露设备数 × 2 必须 ≥ 模型tp_degree生产环境用--device显式暴露不要用privileged参数对齐--max-batch-size、--max-total-tokens必须与模型导出配置完全一致--max-input-length必须小于序列长度本地 Neuron 模型会自动推导参数无需手动指定数据卷务必把/data挂载为持久化卷避免重复下载与导出gated 模型记得通过-e HF_TOKEN传入有效的 Hugging Face 令牌内存预算batch size 的选择受限于每 core 16GB每设备 2 core的总显存容量选型前先核对实例规格。至此你已经可以从零开始在 AWS Trainium / Inferentia 实例上通过 Docker 部署 Neuron 后端的 TGI 服务并正确配置静态维度参数、完成流式与非流式请求的调用。更深入的架构背景可参阅 docs/source/architecture.mdGPU 与 Trainium 的实例化部署差异可对照 docs/source/installation.md 与 docs/source/installation_inferentia.md。【免费下载链接】text-generation-inferenceLarge Language Model Text Generation Inference项目地址: https://gitcode.com/GitHub_Trending/te/text-generation-inference创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考