Pemecahan 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 guna mulai menggunakan Cloud TPU, lihat panduan memulai JAX.

Catatan: Jika Anda tidak dapat menyelesaikan masalah menggunakan panduan ini, lihat Mendapatkan Dukungan untuk bantuan lebih lanjut.

Masalah umum JAX

Jika Anda mengalami masalah saat mengembangkan model pelatihan atau 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 resource TPU Anda digunakan menggunakan alat yang dijelaskan dalam performa Profiling JAX.

Memecahkan masalah memori

Catatan: Memori TPU (memori bandwidth tinggi digunakan bersama di setiap chip TPU. Hal ini memungkinkan koordinasi yang efisien antara dua TensorCore di chip.

Anda dapat memantau penggunaan memori dengan JAX Device Memory Profiler, tetapi Anda tidak dapat langsung mengelola cara penggunaannya.

Memory Profiler Perangkat dapat digunakan untuk:

• Mencari tahu array dan file yang dapat dieksekusi yang ada di memori TPU pada waktu tertentu, atau
• Lacak kebocoran memori.

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

Memecahkan masalah TPU

Bagaimana cara memastikan TPU dapat berjalan?

Detail

Semuanya akan dijalankan di TPU selama JAX tidak mencetak "Tidak ada GPU/TPU yang ditemukan, kembali ke 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: Tidak dapat menginisialisasi backend 'tpu': UNAVAILABLE: Tidak ada Platform TPU yang tersedia.

Detail

Pesan error runtime ini dan/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 TPU VM yang salah.

Pastikan Anda menjalankan versi runtime JAX saat ini, lalu coba lagi.