Runtime TPU

Dokumen ini berisi panduan tentang memilih versi runtime yang sesuai untuk Cloud TPU Anda, bergantung pada framework dan arsitektur TPU (VM atau Node) yang Anda gunakan.

Versi runtime VM TPU

Bagian ini menjelaskan versi software TPU yang harus Anda gunakan untuk TPU dengan arsitektur VM TPU. Untuk arsitektur Node TPU, lihat Versi software TPU Node.

Versi software TPU tersedia untuk framework TensorFlow, PyTorch, dan JAX.

TensorFlow

Gunakan versi software TPU yang sesuai dengan versi TensorFlow yang digunakan untuk menulis model Anda.

Mulai TensorFlow 2.15.0, Anda juga harus menentukan runtime stream executor (SE) atau runtime PJRT. Misalnya, jika Anda menggunakan TensorFlow 2.16.1 dengan runtime PJRT, gunakan versi software TPU tpu-vm-tf-2.16.1-pjrt. Versi sebelum TensorFlow 2.15.0 hanya mendukung eksekutor streaming. Untuk mengetahui informasi selengkapnya tentang PJRT, lihat dukungan TensorFlow PJRT.

Versi software VM TensorFlow TPU yang didukung saat ini adalah:

  • tpu-vm-tf-2.16.1-pjrt
  • tpu-vm-tf-2.16.1-se
  • tpu-vm-tf-2.15.1-pjrt
  • tpu-vm-tf-2.15.1-se
  • tpu-vm-tf-2.15.0-pjrt
  • tpu-vm-tf-2.15.0-se
  • tpu-vm-tf-2.14.1
  • tpu-vm-tf-2.14.0
  • tpu-vm-tf-2.13.1
  • tpu-vm-tf-2.13.0
  • tpu-vm-tf-2.12.1
  • tpu-vm-tf-2.12.0
  • tpu-vm-tf-2.11.1
  • tpu-vm-tf-2.11.0
  • tpu-vm-tf-2.10.1
  • tpu-vm-tf-2.10.0
  • tpu-vm-tf-2.9.3
  • tpu-vm-tf-2.9.1
  • tpu-vm-tf-2.8.4
  • tpu-vm-tf-2.8.3
  • tpu-vm-tf-2.8.0
  • tpu-vm-tf-2.7.4
  • tpu-vm-tf-2.7.3

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

Dukungan PJRT TensorFlow

Mulai dari TensorFlow 2.15.0, Anda dapat menggunakan antarmuka PJRT untuk TensorFlow di TPU. PJRT menghadirkan defragmentasi memori perangkat otomatis dan menyederhanakan integrasi hardware dengan framework. Untuk informasi lebih lanjut tentang PJRT, lihat PJRT: Menyederhanakan Integrasi Hardware dan Framework ML di Blog Google Open Source.

Tidak semua fitur TPU v2, v3, dan v4 telah dimigrasikan ke runtime PJRT. Tabel berikut menjelaskan fitur yang didukung pada PJRT atau eksekutor uap.

Akselerator Fitur Didukung di PJRT Didukung di eksekutor streaming
TPU v2-v4 Komputasi padat (tanpa API embedding TPU) Ya Ya
TPU v2-v4 Dense compute API API embedding TPU Tidak Ya
TPU v2-v4 tf.summary/tf.print dengan penempatan perangkat lunak Tidak Ya
TPU v5e Komputasi padat (tanpa API embedding TPU) Ya Tidak
TPU v5e API penyematan TPU T/A - TPU v5e tidak mendukung API penyematan TPU T/A
TPU v5p Komputasi padat (tanpa API embedding TPU) Ya Tidak
TPU v5p API penyematan TPU Ya Tidak

TPU v4 dengan TensorFlow versi 2.10.0 dan yang lebih lama

Jika Anda melatih model di TPU v4 dengan TensorFlow, TensorFlow versi 2.10.0 dan yang lebih lama menggunakan versi khusus v4 yang ditampilkan dalam tabel berikut. Jika versi TensorFlow yang Anda gunakan tidak ditampilkan dalam tabel, ikuti panduan di bagian TensorFlow.

Versi TensorFlow Versi software 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

Versi Libtpu

VM TPU dibuat dengan TensorFlow dan library Libtpu terkait yang telah diinstal sebelumnya. Jika Anda membuat image VM sendiri, tentukan versi software TensorFlow TPU berikut dan versi libtpu yang sesuai:

Versi TensorFlow versi libtpu.so
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.* 1.2.0
2.7.3 1.1.2

PyTorch dan JAX

Gunakan image dasar umum berikut untuk PyTorch dan JAX, lalu instal versi framework yang Anda perlukan dengan pip.

  • tpu-ubuntu2204-base (default)
  • v2-alpha-tpuv5 (TPU v5p)

Lihat dokumen panduan memulai untuk PyTorch/XLA dan JAX untuk mengetahui petunjuk penginstalan.

Versi runtime Node TPU

Bagian ini menjelaskan versi software TPU yang harus Anda gunakan untuk TPU dengan arsitektur TPU Node. Untuk arsitektur VM TPU, lihat Versi software VM TPU.

Versi software TPU tersedia untuk framework TensorFlow, PyTorch, dan JAX.

TensorFlow

Gunakan versi software TPU yang sesuai dengan versi TensorFlow yang digunakan untuk menulis model Anda. Misalnya, jika Anda menggunakan TensorFlow 2.12.0, gunakan versi software TPU 2.12.0. Versi software TPU khusus TensorFlow adalah:

  • 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.2
  • 2.7.3

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

Saat Anda membuat TPU Node, versi terbaru TensorFlow sudah diinstal sebelumnya di TPU Node.

Langkah selanjutnya