Image VM TPU

Saat membuat resource TPU, Anda meneruskan parameter --version atau --runtime-version yang menentukan image VM TPU. Image VM TPU berisi sistem operasi Ubuntu, Docker, dan software lain yang diperlukan untuk menjalankan kode Anda di TPU. Dokumen ini memberikan panduan tentang cara memilih image VM TPU yang sesuai saat Anda membuat Cloud TPU.

PyTorch dan JAX

Gunakan image dasar VM TPU umum berikut untuk PyTorch dan JAX, lalu instal framework yang ingin Anda gunakan.

  • tpu-ubuntu2204-base (default)
  • v2-alpha-tpuv6e (TPU v6e)
  • v2-alpha-tpuv5 (TPU v5p)
  • v2-alpha-tpuv5-lite (TPU v5e)

Lihat dokumen memulai cepat untuk PyTorch/XLA dan JAX untuk mendapatkan petunjuk penginstalan.

TensorFlow

Ada image VM TPU yang khusus untuk setiap versi TensorFlow. Versi TensorFlow berikut didukung di Cloud TPU:

  • 2.18.0
  • 2.17.0
  • 2.16.2
  • 2.16.1
  • 2.15.1
  • 2.15.0
  • 2.14.1
  • 2.14.0
  • 2.13.1
  • 2.13.0
  • 2.12.1
  • 2.12.0
  • 2.11.1
  • 2.11.0
  • 2.10.1
  • 2.10.0
  • 2.9.3
  • 2.9.1
  • 2.8.4
  • 2.8.3
  • 2.8.0
  • 2.7.4
  • 2.7.3

Untuk mengetahui informasi selengkapnya tentang versi patch TensorFlow, lihat Versi patch TensorFlow yang didukung.

Untuk TensorFlow versi 2.15.0 dan yang lebih baru, ada varian image VM TPU berdasarkan API perangkat (PJRT atau eksekutor streaming) yang Anda gunakan.

Pelatihan di v6e, v5p, dan v5e

TPU v6e, v5e, dan v5p mendukung TensorFlow 2.15.0 dan yang lebih baru. Anda menentukan image VM TPU menggunakan formulir: tpu-vm-tf-x.y.z-{pod}-pjrt dengan x adalah versi utama TensorFlow, y adalah versi minor, dan z adalah versi patch TensorFlow. Tambahkan pod setelah versi TensorFlow jika Anda menggunakan TPU multi-host. Misalnya, jika Anda menggunakan TensorFlow 2.16.0 di TPU multi-host, gunakan image VM TPU tpu-vm-tf-2.16.0-pod-pjrt. Untuk versi TensorFlow lainnya, ganti 2.16.0 dengan versi utama dan patch TensorFlow yang Anda gunakan. Jika Anda menggunakan TPU host tunggal, hapus pod.

Menayangkan di v6e dan v5e

Ada image Docker penayangan yang berisi semua persyaratan software yang diperlukan untuk penayangan dengan TensorFlow, PyTorch, dan JAX. Untuk mengetahui informasi selengkapnya, lihat pengantar inferensi Cloud TPU v5e.

TPU v4

Jika Anda menggunakan TPU v4 dan TensorFlow 2.15.0 atau yang lebih baru, ikuti petunjuk untuk pelatihan di v6e, v5p, dan v5e. Jika Anda menggunakan TensorFlow 2.10.0 atau yang lebih lama, gunakan image VM TPU khusus v4:

Versi TensorFlow Versi image VM TPU
2.10.0 tpu-vm-tf-2.10.0-v4
tpu-vm-tf-2.10.0-pod-v4
2.9.3 tpu-vm-tf-2.9.3-v4
tpu-vm-tf-2.9.3-pod-v4
2.9.2 tpu-vm-tf-2.9.2-v4
tpu-vm-tf-2.9.2-pod-v4
2.9.1 tpu-vm-tf-2.9.1-v4
tpu-vm-tf-2.9.1-pod-v4

TPU v2 dan v3

Jika Anda menggunakan TPU v2 atau v3, gunakan image VM TPU yang cocok dengan versi TensorFlow yang Anda gunakan. Misalnya, jika Anda menggunakan TensorFlow 2.14.1, gunakan image TPU tpu-vm-tf-2.14.1. Untuk versi TensorFlow lainnya, ganti 2.14.1 dengan versi TensorFlow yang Anda gunakan. Jika Anda menggunakan TPU multi-host, tambahkan pod ke akhir image TPU, misalnya tpu-vm-tf-2.14.1-pod.

Mulai TensorFlow 2.15.0, Anda juga harus menentukan API perangkat sebagai bagian dari nama image. Misalnya, jika Anda menggunakan TensorFlow 2.16.1 dengan PJRT API, gunakan image TPU tpu-vm-tf-2.16.1-pjrt. Jika Anda menggunakan API eksekutor streaming dengan versi TensorFlow yang sama, gunakan image TPU tpu-vm-tf-2.16.1-se. Versi TensorFlow yang lebih lama dari 2.15.0 hanya mendukung eksekutor streaming.

Dukungan PJRT TensorFlow

Mulai TensorFlow 2.15.0, Anda dapat menggunakan antarmuka PJRT untuk TensorFlow di TPU. PJRT menampilkan defragmentasi memori perangkat otomatis dan menyederhanakan integrasi hardware dengan framework. Untuk informasi selengkapnya tentang PJRT, lihat PJRT: Menyederhanakan Integrasi Hardware dan Framework ML.

Akselerator Fitur Dukungan PJRT Dukungan eksekutor streaming
TPU v2 - v4 Komputasi padat (tanpa API penyematan TPU) Ya Ya
TPU v2 - v4 API komputasi padat + API penyematan TPU Tidak Ya
TPU v2 - v4 tf.summary/tf.print dengan penempatan perangkat soft Tidak Ya
TPU v5e Komputasi padat (tanpa API penyematan TPU) Ya Tidak
TPU v5e API penyematan TPU T/A Tidak
TPU v5p Komputasi padat (tanpa API penyematan TPU) Ya Tidak
TPU v5p API penyematan TPU Ya Tidak

Versi libtpu

Image TensorFlow VM TPU berisi versi TensorFlow tertentu dan library libtpu yang sesuai. Jika Anda membuat image VM sendiri, gunakan versi software TensorFlow TPU berikut dan versi libtpu yang sesuai:

Versi TensorFlow Versi libtpu.so
2.18.0 1.12.0
2.17.0 1.11.0
2.16.2 1.10.1
2.16.1 1.10.1
2.15.1 1.9.0
2.15.0 1.9.0
2.14.1 1.8.1
2.14.0 1.8.0
2.13.1 1.7.1
2.13.0 1.7.0
2.12.1 1.6.1
2.12.0 1.6.0
2.11.1 1.5.1
2.11.0 1.5.0
2.10.1 1.4.1
2.10.0 1.4.0
2.9.3 1.3.2
2.9.1 1.3.0
2.8.3 1.2.3
2.8.0 1.2.0
2.7.3 1.1.2

Langkah selanjutnya