本指南說明如何透過 PyTorch，使用 Google Kubernetes Engine (GKE) 上的張量處理單元 (TPU) 和 JetStream，提供大型語言模型 (LLM) 服務。在本指南中，您會將模型權重下載至 Cloud Storage，並使用執行 JetStream 的容器，將模型權重部署至 GKE Autopilot 或 Standard 叢集。
如果您需要在 JetStream 上部署模型時，運用 Kubernetes 功能提供的擴充性、復原能力和成本效益，這份指南是個不錯的起點。
本指南適用於使用 PyTorch 的生成式 AI 客戶、GKE 新使用者或現有使用者、機器學習工程師、MLOps (DevOps) 工程師，以及有興趣使用 Kubernetes 容器協調功能來提供 LLM 服務的平台管理員。
背景
透過 JetStream 在 GKE 上使用 TPU 提供 LLM，您就能建構完善且可用於正式環境的服務解決方案，同時享有代管型 Kubernetes 的所有優點，包括成本效益、擴充性和高可用性。本節說明本教學課程中使用的主要技術。
關於 TPU
TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)，用於加速機器學習和 AI 模型，這些模型是使用 TensorFlow、PyTorch 和 JAX 等架構建構而成。
在 GKE 中使用 TPU 之前，建議您先完成下列學習路徑：
- 如要瞭解目前可用的 TPU 版本，請參閱 Cloud TPU 系統架構。
- 瞭解 GKE 中的 TPU。
本教學課程涵蓋各種大型語言模型的服務。GKE 會在單一主機 TPUv5e 節點上部署模型，並根據模型需求設定 TPU 拓撲，以低延遲方式提供提示。
關於 JetStream
JetStream 是 Google 開發的開放原始碼推論服務架構，JetStream 可在 TPU 和 GPU 上執行高效能、高處理量和記憶體最佳化推論作業。JetStream 提供進階效能最佳化功能，包括持續批次處理、KV 快取最佳化和量化技術，可簡化 LLM 部署作業。JetStream 可讓 PyTorch/XLA 和 JAX TPU 服務達到最佳效能。
持續批次處理
持續批次處理技術會將傳入的推論要求動態分組為批次，進而縮短延遲時間並提高總處理量。
KV 快取量化
KV 快取量化會壓縮注意力機制中使用的鍵值快取，進而減少記憶體需求。
Int8 權重量化
Int8 權重量化會將模型權重的精確度從 32 位元浮點數降低至 8 位元整數，進而加快運算速度並減少記憶體用量。
如要進一步瞭解這些最佳化作業，請參閱 JetStream PyTorch 和 JetStream MaxText 專案存放區。
關於 PyTorch
PyTorch 是由 Meta 開發的開放原始碼機器學習架構，現已納入 Linux 基金會旗下。PyTorch 提供高階功能，例如張量運算和深層類神經網路。