Jetstream Pytorch 和 PyTorch XLA 之間的區別

本指南向 optimum-tpu 使用者解釋 Jetstream Pytorch 和 PyTorch XLA 之間的區別，因為它們是 TGI 中可用的兩種後端。

JetStream PyTorch 是一個基於 PyTorch XLA 構建的高效能推理引擎。它在 TPU 上執行大型語言模型 (LLM) 時針對吞吐量和記憶體效率進行了最佳化。

特性	Jetstream Pytorch	PyTorch XLA
訓練	❌	✅
服務	✅	✅
效能	更高的服務效能	標準效能
靈活性	僅限於服務	完整的 PyTorch 生態系統
用例	生產推理	開發和訓練
整合	針對部署最佳化	標準 PyTorch 工作流程

注意：預設情況下，optimum-tpu 使用 PyTorch XLA 進行訓練，使用 Jetstream Pytorch 進行服務。

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

您可以找到更多關於以下內容的資訊

• PyTorch XLA：https://pytorch.org/xla/ 和 https://github.com/pytorch/xla
• Jetstream Pytorch：https://github.com/AI-Hypercomputer/jetstream-pytorch