Memecahkan masalah JAX - TPU

Panduan ini memberikan petunjuk ke informasi pemecahan masalah JAX untuk membantu Anda mengidentifikasi dan menyelesaikan masalah yang mungkin Anda alami saat melatih model JAX di Cloud TPU.

Untuk panduan yang lebih umum tentang cara mulai menggunakan Cloud TPU, lihat panduan memulai JAX.

Masalah JAX umum

Jika Anda mengalami masalah saat mengembangkan model pelatihan atau melakukan pelatihan dengan JAX, lihat FAQ JAX.

Untuk error pemrograman yang lebih umum yang mungkin Anda alami saat menulis aplikasi pelatihan dengan JAX, lihat Error JAX.

Membuat profil performa JAX

Anda dapat memahami cara penggunaan resource TPU menggunakan alat yang dijelaskan dalam Membuat profil performa JAX.

Memecahkan masalah memori

Anda dapat memantau penggunaan memori dengan profiler memori perangkat JAX, tetapi Anda tidak dapat mengelola penggunaannya secara langsung.

Profiler memori perangkat JAX dapat digunakan untuk:

Anda tidak dapat menentukan cara mengalokasikan memori TPU untuk operasi tertentu. Untuk mengetahui informasi selengkapnya tentang masalah performa TPU khusus JAX, lihat Catatan Performa untuk menggunakan TPU dengan JAX.

Memecahkan masalah TPU

Bagian berikut menjelaskan cara menyelesaikan beberapa masalah umum yang mungkin Anda alami saat menjalankan program JAX di TPU.

Bagaimana cara memverifikasi bahwa TPU sedang berjalan?

Semuanya akan dijalankan di TPU selama JAX tidak mencetak "No GPU/TPU found, falling back to CPU".

Anda dapat memverifikasi bahwa TPU aktif dengan melihat jax.devices(), tempat Anda akan melihat beberapa perangkat TPU ditampilkan, atau memverifikasi secara terprogram dengan: assert jax.devices()[0].platform == 'tpu'.

RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available.

Pesan error runtime ini atau menemukan hal berikut di /tmp/tpu_logs/tpu_driver.WARNING pada VM TPU: W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:xxxxxx dapat menunjukkan bahwa Anda menjalankan versi VM TPU yang salah.

Pastikan Anda menjalankan versi runtime JAX saat ini dan coba lagi.

Memecahkan masalah TPU dan GKE

Untuk membantu pemecahan masalah, aktifkan logging panjang di manifes beban kerja GKE Anda, lalu berikan log tersebut ke dukungan GKE.

TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0

Bagian berikut menjelaskan pesan error terkait penyiapan dan penyelesaian TPU dan GKE.

tidak ada endpoint yang tersedia untuk layanan 'jobset-webhook-service'

Error ini berarti jobset tidak diinstal dengan benar. Periksa apakah Pod Kubernetes deployment jobset-controller-manager sedang berjalan. Untuk mengetahui informasi selengkapnya, lihat dokumentasi pemecahan masalah JobSet.

Inisialisasi TPU gagal: Gagal terhubung

Pastikan versi node GKE Anda adalah 1.30.4-gke.1348000 atau yang lebih baru (GKE 1.31 tidak didukung).