進階 TGI 伺服器設定

Jetstream Pytorch 與 Pytorch XLA 後端

Jetstream Pytorch 是一個經過高度優化的 Pytorch 引擎，專門用於在 Cloud TPU 上進行大型語言模型（LLM）的推論服務。如果該相依項目可用，系統將預設選用此引擎。

我們建議將 Jetstream 與 TGI 搭配使用以獲得最佳效能。若因某些原因您希望改用 Pytorch/XLA 後端，您可以設定 JETSTREAM_PT_DISABLE=1 環境變數。

欲了解更多資訊，請參閱我們關於 Jetstream 與 Pytorch XLA 差異的討論。

量化 (Quantization)

使用 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：您可以手動降低批次大小（batch size）。

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 的值。

`max-batch-prefill-tokens ≤ max-input-length * max_batch_size`。否則，系統將會出現錯誤，因為該設定不合理。如果 max-batch-prefill-tokens 設定過大，您將無法處理任何請求。

分片 (Sharding)

分片由 TGI 伺服器自動執行，因此您的模型會使用所有可用的 TPU。我們使用張量並行處理 (tensor parallelism)，因此層級會自動分配到所有可用的 TPU 上。不過，TGI 路由（router）只會看到一個分片。

關於張量並行處理的更多資訊，請參閱：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 用於設定伺服器行為的參數。

`--privileged --shm-size 16GB --net host` 是必需的，如 https://github.com/pytorch/xla 所述。

後續步驟

請查閱 TGI 文件以了解更多 TGI 伺服器設定選項。