排解 JAX - TPU 問題
本指南提供 JAX 疑難排解資訊的指標，協助您找出並解決在 Cloud TPU 上訓練 JAX 模型時可能遇到的問題。
如需 Cloud TPU 的一般入門指南，請參閱快速入門導覽課程：建立 TPU 執行個體。
一般 JAX 問題
如果您在開發訓練模型或使用 JAX 訓練時遇到問題，請參閱 JAX 常見問題。
如要進一步瞭解使用 JAX 編寫訓練應用程式時可能遇到的程式設計錯誤，請參閱「JAX 錯誤」。
剖析 JAX 效能
您可以使用「剖析 JAX 效能」一文所述的工具，瞭解 TPU 資源的使用情況。
排解記憶體問題
您可以使用 JAX 裝置記憶體剖析器監控記憶體用量，但無法直接管理記憶體用量。
JAX 裝置記憶體分析器可用於：
- 瞭解特定時間點哪些陣列和可執行檔位於 TPU 記憶體中，或
- 追蹤記憶體流失問題。
您無法指定特定作業的 TPU 記憶體配置方式。 如要進一步瞭解 JAX 專屬的 TPU 效能問題，請參閱「使用 JAX 的 TPU 效能注意事項」。
排解 TPU 問題
下列各節說明如何解決在 TPU 上執行 JAX 程式時可能遇到的一些常見問題。
如何確認 TPU 是否正在執行？
只要 JAX 沒有列印「No GPU/TPU found, falling back to CPU.」，所有項目都會在 TPU 上執行。
如要確認 TPU 是否已啟用，請查看 jax.devices()，您應該會看到顯示多個 TPU 裝置，或是使用 assert jax.devices()[0].platform == 'tpu' 以程式輔助方式驗證。
RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available.
您看到 RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No
TPU Platform available 錯誤訊息，或在 TPU VM 上發現下列 /tmp/tpu_logs/tpu_driver.WARNING 警告：
W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:VERSION
這項錯誤可能表示您執行的 TPU VM 版本有誤。確認您執行的是目前的 JAX 執行階段版本，然後重試。
排解 TPU 和 GKE 問題
如要協助排解問題，請在 GKE 工作負載資訊清單中啟用詳細記錄功能，然後將記錄檔提供給 GKE 支援團隊。
TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0
以下各節說明與 TPU 和 GKE 設定相關的錯誤訊息和解決方式。
no endpoints available for service 'jobset-webhook-service'
這表示 Jobset 安裝有誤。檢查 jobset-controller-manager deployment Kubernetes Pod 是否正在執行中。詳情請參閱 JobSet 疑難排解說明文件。
TPU 初始化失敗：無法連線
請確認 GKE 節點版本為 1.30.4-gke.1348000 以上版本 (不支援 GKE 1.31)。