高階 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 伺服器配置選項。