高级 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 128GB）上，您可以服务多达 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：最大组合输入和输出 tokens
• --max-batch-prefill-tokens：批处理的最大 tokens 数
• --max-batch-total-tokens：一个批次中的最大总 tokens 数

为了减少内存使用量，您可以尝试为 --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 服务器配置选项。