Beralih versi software di Cloud TPU

Ringkasan

Versi software framework yang berjalan di TPU Anda harus cocok dengan versi yang berjalan di VM lokal Anda. Versi software ini sekarang dapat dialihkan di Cloud TPU yang berjalan, tanpa menghapus dan membuat ulang TPU. Hal ini juga memungkinkan konfigurasi Cloud TPU dengan framework software versi malam tertentu. Sebaiknya pilih versi yang didukung untuk framework ini.

Penggunaan

Cara yang direkomendasikan untuk beralih versi adalah dengan menggunakan library python cloud-tpu-client.

Contoh penggunaan untuk TensorFlow.

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

Tindakan ini akan mengonfigurasi Cloud TPU agar sesuai dengan versi TensorFlow yang berjalan di VM lokal Anda, termasuk rilis resmi serta build setiap malam.

Parameter restart_type dari configure_tpu_version API menentukan perilaku mulai ulang TPU saat beralih versi. Opsinya adalah 'always' (default) dan 'ifNeeded'.

  • 'always' dapat digunakan untuk memperbaiki TPU dengan, misalnya, status UNHEALTHY_TENSORFLOW, atau yang menampilkan error Out of Memory (OOM) karena kebocoran resource dari proses sebelumnya. Jika opsi ini ditetapkan, TPU akan dimulai ulang meskipun versi framework baru tidak diinstal.

  • 'ifNeeded' dapat berguna karena tidak memulai ulang runtime jika sudah berada di versi yang tepat, sehingga tidak akan menambahkan waktu startup yang signifikan ke skrip pelatihan. Jika opsi ini ditetapkan, TPU hanya akan dimulai ulang jika tidak memiliki versi framework yang benar.

Library ini berkomunikasi langsung dengan Cloud TPU sehingga kode ini perlu dijalankan di VM dalam jaringan yang sama. Sebaiknya jalankan ini dalam kode untuk model Anda yang lain.

Opsi software tambahan

TensorFlow menyertakan string tf.__version__ yang merupakan cara paling sederhana untuk mengonfigurasi versi yang benar. Opsi software lainnya termasuk:

  • PyTorch - pytorch-1.13, pytorch-nightly-dev20220930, pytorch-nightly
  • Jax - tpu_driver, tpu_driver0.1-dev20200320, tpu_driver_nightly

Misalnya, untuk mengonfigurasi TPU agar dijalankan dengan build PyTorch terbaru setiap malam.

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()