透過 JetStream 和 Pathways 在 GKE 上使用多主機 TPU 提供 LLM

本指南說明如何使用多個節點的張量處理單元 (TPU)，在 Google Kubernetes Engine (GKE) 上提供最先進的大型語言模型 (LLM)，例如 Llama 3.1 405B。

本指南說明如何使用可攜式開放原始碼技術 (Kubernetes、JetStream、Pathways on Cloud 和 LeaderWorkerSet (LWS) API)，充分運用 GKE 的精細控制、擴充性、復原能力、可攜性和成本效益，在 GKE 上部署及提供 AI/ML 工作負載。

背景

大型語言模型越來越大，已無法再裝進單一主機 TPU 切片。如要進行 ML 推論，您可以使用 Pathways on Cloud，在 GKE 上跨多個互連的 TPU 節點，執行大規模的多主機推論。本指南將逐步說明如何使用多主機 TPU 配量佈建 GKE 叢集、使用 Pathways on Cloud 二進位檔、透過 MaxText 架構啟動 JetStream 伺服器，以及發出多主機推論要求。

透過 JetStream、MaxText 和 Pathways，在 GKE 上使用 TPU 提供 LLM，即可建構完善且可用於正式環境的服務解決方案，同時享有代管型 Kubernetes 的所有優點，包括成本效益、擴充性和高可用性。本節說明本教學課程中使用的主要技術。

關於 TPU

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

在 GKE 中使用 TPU 之前，建議您先完成下列學習路徑：

  1. 如要瞭解目前可用的 TPU 版本，請參閱 Cloud TPU 系統架構。
  2. 瞭解 GKE 中的 TPU。

本教學課程涵蓋 Llama 3.1-405B 模型服務。GKE 會在多主機 TPU v6e 節點上部署模型，並根據模型需求設定 TPU 拓撲，以低延遲服務提示。

Pathways on Cloud

Pathways 是加速器的大規模自動化調度管理層，Pathways 的設計宗旨是探索新系統和機器學習研究構想，同時維持現有模型的頂尖效能。Pathways 可讓單一 JAX 用戶端程序協調一或多個大型 TPU 切片之間的運算，簡化涵蓋數百或數千個 TPU 晶片的機器學習運算。

JetStream

JetStream 是 Google 開發的開放原始碼推論服務架構，JetStream 可在 TPU 和 GPU 上執行高效能、高處理量和記憶體最佳化推論作業。JetStream 提供進階效能最佳化功能，包括持續批次處理、KV 快取最佳化和量化技術，可簡化 LLM 部署作業。JetStream 可啟用 PyTorch/XLA 和 JAX TPU 服務，進而提升效能。

MaxText

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

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

Llama 3.1 405B

Llama 3.1 405B 是 Meta 的大型語言模型，適用於各種自然語言處理工作，包括生成文字、翻譯和回答問題。GKE 提供基礎架構，支援這類模型的分散式訓練和服務需求。

詳情請參閱 Llama 說明文件。

架構

本節說明本教學課程使用的 GKE 架構。此架構包含 GKE Standard 叢集，可佈建 TPU 並代管 JetStream 和 Pathways 元件，以部署及提供模型服務。

下圖顯示這個架構的元件：

GKE 叢集架構，其中包含 JetStream 和 Pathways 元件的多主機 TPU 節點集區。

這個架構包含下列元件：

  • GKE Standard 區域叢集。
  • 多主機 TPU 配量節點集區，用於代管 JetStream 部署作業和 Pathways 元件。
  • Pathways resource manager 會管理加速器資源，並協調使用者工作的加速器分配作業。
  • Pathways client 會與 Pathways resource manager 協調，決定編譯後的程式要放在哪裡執行。
  • Pathways worker 會在加速器機器上執行運算，並透過 IFRT Proxy 伺服器將資料傳回工作負載。
  • IFRT proxy client 會實作 OSS Interim Framework Runtime (IFRT) API，並做為工作負載和 Pathways 元件之間的通訊橋樑。
  • IFRT proxy server 會接收 IFRT proxy client 的要求，並轉送至 Pathways client，分配工作。
  • JetStream-Pathways 容器提供以 JAX 為基礎的推論伺服器，可接收推論要求，並將執行程序委派給 Pathways workers
  • Service 元件會將傳入流量分散到所有 JetStream HTTP 副本。
  • JetStream HTTP 是 HTTP 伺服器，可接受以 JetStream 必要格式做為封裝的要求，並傳送至 JetStream 的 GRPC 用戶端。