JAX のトラブルシューティング - TPU

このガイドでは、Cloud TPU で JAX モデルをトレーニングする際に発生する可能性のある問題を特定して解決するため、JAX のトラブルシューティング情報について説明します。

Cloud TPU を使い始める際の一般的なガイドについては、JAX クイックスタートをご覧ください。

JAX に関する一般的な問題

トレーニング モデルの開発または JAX でのトレーニング中に問題が発生した場合は、JAX のよくある質問をご覧ください。

JAX を使用してトレーニング アプリケーションを作成するときに発生する可能性のある一般的なプログラミング エラーについては、JAX エラーをご覧ください。

JAX パフォーマンスのプロファイルを作成する

JAX パフォーマンスのプロファイリングで説明されているツールを使用して、TPU リソースの使用状況を把握できます。

メモリに関する問題のトラブルシューティング

JAX Device Memory Profiler を使用すると、TPU メモリの使用状況を確認できます。以下に使用できます。

TPU の問題のトラブルシューティング

TPU が実行されていることを確認する方法

詳細

JAX から「GPU/TPU が見つかりません。CPU にフォールバック」と出力されない限り、すべてが TPU で実行されます。

TPU がアクティブであることを確認するには、jax.devices()(複数の TPU デバイスが表示されていることを確認します)を確認するか、プログラムで assert jax.devices()[0].platform == 'tpu' を確認します。