排解 JAX - TPU 問題

本指南提供 JAX 疑難排解資訊的指標，協助您找出並解決在 Cloud TPU 上訓練 JAX 模型時可能遇到的問題。

如需 Cloud TPU 的一般入門指南，請參閱 JAX 快速入門導覽課程。

一般 JAX 問題

如果在開發訓練模型或使用 JAX 訓練時遇到問題，請參閱 JAX 常見問題。

如要進一步瞭解使用 JAX 編寫訓練應用程式時可能遇到的一般程式設計錯誤，請參閱「JAX 錯誤」。

剖析 JAX 效能

您可以使用「剖析 JAX 效能」一文所述的工具，瞭解 TPU 資源的使用情況。

排解記憶體問題

您可以使用 JAX 裝置記憶體分析器監控記憶體使用情形，但無法直接管理記憶體使用方式。

JAX 裝置記憶體分析器可用於：

您無法指定特定作業的 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.

如果 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:xxxxxx 表示您執行的 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 設定相關的錯誤訊息和解決方法。

服務「jobset-webhook-service」沒有可用端點

這表示 Jobset 安裝有誤。檢查 jobset-controller-manager 部署作業 Kubernetes Pod 是否正在執行。詳情請參閱 JobSet 疑難排解說明文件。

TPU 初始化失敗：無法連線

請確認 GKE 節點版本為 1.30.4-gke.1348000 以上版本 (不支援 GKE 1.31)。