透過 JetStream 在 GKE 上使用 TPU 提供 Gemma

注意： 現在建議使用 vLLM，在 GKE 的 TPU 上提供 LLM。如要開始使用，請參閱「透過 vLLM 在 GKE 上使用 TPU Trillium 提供 LLM」一文。

本教學課程說明如何使用 Google Kubernetes Engine (GKE) 的張量處理單元 (TPU)，提供 Gemma 大型語言模型 (LLM) 服務。您將預先建構的容器連同 JetStream 和 MaxText 部署至 GKE。您也可以設定 GKE，在執行階段從 Cloud Storage 載入 Gemma 7B 權重。

本教學課程的適用對象為機器學習 (ML) 工程師、平台管理員和操作員，以及有興趣使用 Kubernetes 容器自動化調度管理功能提供 LLM 服務的資料和 AI 專家。如要進一步瞭解我們在Google Cloud 內容中提及的常見角色和範例工作，請參閱「常見的 GKE 使用者角色和工作」。

閱讀本頁面之前，請先熟悉下列概念：

• Autopilot 模式和 Standard 模式
• 目前可用的 TPU 版本 (Cloud TPU 系統架構)
• GKE 中的 TPU

背景

本節說明本教學課程中使用的主要技術。

Gemma

Gemma 是一組開放授權的輕量級生成式人工智慧 (AI) 模型，這些 AI 模型可在應用程式、硬體、行動裝置或代管服務中執行。您可以使用 Gemma 模型生成文字，也可以調整這些模型來執行特定工作。

詳情請參閱 Gemma 說明文件。

TPU

TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)，用於加速機器學習和 AI 模型，這些模型是使用 TensorFlow、PyTorch 和 JAX 等架構建構而成。

本教學課程涵蓋 Gemma 7B 模型服務。GKE 會在單一主機 TPUv5e 節點上部署模型，並根據模型需求設定 TPU 拓撲，以低延遲方式提供提示。

JetStream

JetStream 是 Google 開發的開放原始碼推論服務架構，JetStream 可在 TPU 和 GPU 上執行高效能、高處理量和記憶體最佳化推論作業。這項技術提供進階效能最佳化功能，包括持續批次處理和量化技術，可協助您部署 LLM。JetStream 可讓 PyTorch/XLA 和 JAX TPU 服務達到最佳效能。

如要進一步瞭解這些最佳化作業，請參閱 JetStream PyTorch 和 JetStream MaxText 專案存放區。

MaxText

MaxText 是高效能、可擴充且可調整的 JAX LLM 實作項目，以 Flax、Orbax 和 Optax 等開放原始碼 JAX 程式庫為基礎建構而成。MaxText 的僅解碼器 LLM 實作是以 Python 編寫，它大量運用 XLA 編譯器，無須建構自訂核心即可達到高效能。

如要進一步瞭解 MaxText 支援的最新模型和參數大小，請參閱 MaxtText 專案存放區。