Jetstream Pytorch 與 PyTorch XLA 之差異

本指南旨在向 optimum-tpu 使用者說明 Jetstream Pytorch 與 PyTorch XLA 之間的差異，因為它們是 TGI 中可用的兩種後端。

JetStream PyTorch 是一個建立在 PyTorch XLA 之上的高效能推論引擎。它針對在 TPU 上執行大型語言模型 (LLM) 時的輸送量和記憶體效率進行了最佳化。

特性	Jetstream Pytorch	PyTorch XLA
訓練	❌	✅
模型服務 (Serving)	✅	✅
效能	更高的服務效能	標準效能
靈活性	僅限於服務 (Inference)	完整的 PyTorch 生態系統
使用案例	生產環境推論	開發與訓練
整合	針對部署進行最佳化	標準 PyTorch 工作流程

註：預設情況下，optimum-tpu 在訓練時使用 PyTorch XLA，在服務時使用 Jetstream Pytorch。

您可以設定 optimum-tpu 以在 TGI 中使用任一版本進行服務。您可以透過在 docker 執行參數中設定 -e JETSTREAM_PT_DISABLE=1，在 TGI 中使用 Pytorch XLA 後端。

您可以從以下連結找到更多相關資訊：

• PyTorch XLA: https://pytorch.com.tw/xla/ 與 https://github.com/pytorch/xla
• Jetstream Pytorch: https://github.com/AI-Hypercomputer/jetstream-pytorch