高级 TGI 服务器配置

Jetstream Pytorch 和 Pytorch XLA 后端

Jetstream Pytorch 是一个高度优化的 Pytorch 引擎，用于在 Cloud TPU 上服务 LLM。如果依赖项可用，此引擎默认被选中。

我们建议将 Jetstream 与 TGI 结合使用以获得最佳性能。如果出于某种原因您想使用 Pytorch/XLA 后端，您可以设置环境变量 JETSTREAM_PT_DISABLE=1。

有关更多信息，请参阅我们关于 jetstream 和 pytorch XLA 之间区别 的讨论。

量化

使用 Jetstream Pytorch 引擎时，可以启用量化以减少内存占用并提高吞吐量。要启用量化，请设置环境变量 QUANTIZATION=1。例如，在 2x4 TPU v5e（每个芯片 16GB * 8 = 每个 pod 128 GB）上，您可以服务多达 70B 参数的模型，例如 Llama 3.3-70B。量化在权重加载时即时进行 int8 格式。与任何量化选项一样，您可以预期模型精度会略有下降。如果未启用量化选项，模型将以 bf16 格式提供。

如何解决内存需求

如果您遇到 Backend(NotEnoughMemory(2048))，这里有一些有助于减少 TGI 内存使用量的解决方案

Optimum-TPU 特定参数

• -e QUANTIZATION=1: 启用量化。这将使内存需求减少近一半
• -e MAX_BATCH_SIZE=n: 您可以手动减小批处理大小

TGI 特定参数

• --max-input-length：最大输入序列长度
• --max-total-tokens: 最大组合输入和输出 token 数
• --max-batch-prefill-tokens：批处理的最大令牌数
• --max-batch-total-tokens: 批处理中的最大总 token 数

要减少内存使用，您可以尝试减小 --max-input-length、--max-total-tokens、--max-batch-prefill-tokens 和 --max-batch-total-tokens 的值。

分片

分片由 TGI 服务器自动完成，因此您的模型会使用所有可用的 TPU。我们进行张量并行，因此层会自动分割到所有可用的 TPU 中。但是，TGI 路由器将只看到一个分片。

有关张量并行的更多信息，请参见此处 https://huggingface.co/docs/text-generation-inference/conceptual/tensor_parallelism。

理解配置

关键参数说明

必需参数

• --shm-size 16GB: 增加默认的共享内存分配。
• --privileged: 访问 TPU 所需。
• --net host: 使用主机网络模式。这些是运行 TPU 容器所必需的，以便容器能够正确访问 TPU 硬件。

可选参数

• -v ~/hf_data:/data: 模型存储的卷挂载，这使您无需在每次启动时重新下载模型权重。您可以使用任何您喜欢的文件夹，只要它映射到 /data 即可。
• -e SKIP_WARMUP=1: 禁用预热以进行快速测试（不建议用于生产）。这些是 TGI 和 optimum-TPU 用于配置服务器行为的参数。

下一步

请查看 TGI 文档 以获取更多 TGI 服务器配置选项。