在 Cloud TPU 執行推論

「提供」是指將訓練好的機器學習模型部署至正式環境，以便執行推論。TPU v5e 和後續版本支援推論。提供工作負載時，首要任務是達成延遲時間服務水準目標 (SLO)。

本文說明如何在單一主機 TPU 提供模型。如果 TPU 切片有八個或更少的晶片，則有一個 TPU VM 或主機，稱為「單主機」TPU。如要瞭解如何使用 PyTorch 進行多主機推論，請參閱 vLLM-TorchTPU 和 SGLang-TorchTPU。如要瞭解如何使用 JAX 執行多主機推論，請參閱「使用 Pathways 執行多主機推論」。

設定 TPU

無論選擇哪種堆疊，都必須有具備單一主機 TPU 的 Google Cloud 專案，才能提供模型。下列步驟適用於所有程式庫和架構組合：

  1. 建立專案。您必須擁有 Google Cloud 帳戶和專案，才能使用 Cloud TPU。詳情請參閱「設定 TPU 的Google Cloud 專案」。

  2. 查看配額。請確認您有足夠的配額，可供推論使用的 TPU 核心數量。TPU v5e 的訓練和提供作業會使用不同的配額。在 Cloud TPU API 中，TPU v5e 的特定服務配額如下：

    • 隨選 v5e 資源：TPU v5 lite pod cores for serving per project per zone
    • 可先占的 v5e 資源：Preemptible TPU v5 lite pod cores for serving per project per zone

    其他 TPU 版本的訓練和提供工作負載則使用相同配額。詳情請參閱「Cloud TPU 配額」。

  3. 建立 TPU VM。請選擇有提供服務的 TPU 版本和可用區，然後建立單一主機 TPU VM。詳情請參閱「TPU 區域和可用區」和「管理 TPU 資源」。

TPU VM 執行後，請按照「選擇推論堆疊」中的程式庫和架構設定指南操作。

選擇推論堆疊

Cloud TPU 支援透過兩個開放原始碼服務程式庫進行大型語言模型 (LLM) 推論：vLLM 和 SGLang。這兩個程式庫都可在 Cloud TPU 上執行，並搭配 JAX 或 PyTorch 後端，因此無論模型是以何種框架編寫，您都能選擇最適合工作負載的程式庫。

請先選擇程式庫，然後選取與模型架構相符的後端，如下表所示：

提供程式庫 JAX 後端 PyTorch 後端
vLLM vllm-jax vllm-torchtpu
SGLang sglang-jax sglang-torchtpu

使用 vLLM 進行推論

vLLM 是開放原始碼程式庫，可快速推論及提供 LLM，以持續批次處理和 PagedAttention 為基礎。Cloud TPU 透過 JAX 模型的 vllm-jax 外掛程式和 PyTorch 模型的 vllm-torchtpu 外掛程式，與 vLLM 搭配運作。

如要查看更多使用 vLLM 在 TPU 上提供模型的範例，請參閱下列資源：

vLLM 搭配 JAX

vLLM TPU 說明文件涵蓋安裝、快速入門導覽課程、支援的模型和調整選項。詳情請參閱「在 TPU 上使用 vLLM 提供 JAX 模型」。

搭配 PyTorch 使用 vLLM

vllm-torchtpu 外掛程式會將 TorchTPU 核心和執行階段支援新增至 vLLM。詳情請參閱「在 TPU 上使用 vLLM 提供 PyTorch 模型」。

使用 SGLang 進行推論

SGLang 是開放原始碼的大型語言和多模態模型服務框架，結合高效能執行階段和結構化生成前端。在 Cloud TPU 上，SGLang 會透過 SGLang-JAX 後端 (sglang-jax) 執行 JAX 模型，並透過 SGLang-TorchTPU 後端 (sglang-torchtpu) 執行 PyTorch 模型。

SGLang 搭配 PyTorch

如要瞭解安裝、伺服器啟動和支援的設定，請參閱「在 TPU 上使用 SGLang-TorchTPU 執行階段提供 PyTorch 模型」。

SGLang 搭配 JAX

如要瞭解安裝、伺服器啟動和支援的設定，請參閱「在 TPU 上使用 SGLang-JAX 執行階段提供 JAX 模型」。

剖析推論效能

設定推論後，您可以使用分析器分析效能和 TPU 使用率。如要進一步瞭解剖析功能，請參閱下列資源：