Menjalankan penghitungan di VM Cloud TPU menggunakan JAX

Dokumen ini memberikan pengantar singkat tentang cara menggunakan JAX dan Cloud TPU.

Sebelum mengikuti panduan memulai ini, Anda harus membuat akun Google Cloud Platform, menginstal Google Cloud CLI, dan mengonfigurasi perintah gcloud. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan akun dan project Cloud TPU.

Menginstal Google Cloud CLI.

Google Cloud CLI berisi alat dan library untuk berinteraksi dengan produk dan layanan Google Cloud. Untuk mengetahui informasi selengkapnya, lihat Menginstal Google Cloud CLI.

Mengonfigurasi perintah gcloud

Jalankan perintah berikut untuk mengonfigurasi gcloud agar dapat menggunakan project Google Cloud Anda dan menginstal komponen yang diperlukan untuk pratinjau VM TPU.

  $ gcloud config set account your-email-account
  $ gcloud config set project your-project-id

Mengaktifkan Cloud TPU API

  1. Aktifkan Cloud TPU API menggunakan perintah gcloud berikut di Cloud Shell. (Anda juga dapat mengaktifkannya dari Konsol Google Cloud).

    $ gcloud services enable tpu.googleapis.com

  2. Jalankan perintah berikut untuk membuat identitas layanan.

    $ gcloud beta services identity create --service tpu.googleapis.com

Membuat VM Cloud TPU dengan gcloud

Dengan VM Cloud TPU, model dan kode Anda berjalan langsung di mesin host TPU. Anda dapat menggunakan SSH langsung ke host TPU. Anda dapat menjalankan kode arbitrer, menginstal paket, melihat log, dan men-debug kode langsung di Host TPU.

  1. Buat VM TPU dengan menjalankan perintah berikut dari Cloud Shell atau terminal komputer tempat Google Cloud CLI diinstal.

    (vm)$ gcloud compute tpus tpu-vm create tpu-name \
--zone=us-central1-a \
--accelerator-type=v3-8 \
--version=tpu-ubuntu2204-base

    Kolom wajib diisi

    zone
    Zona tempat Anda berencana membuat Cloud TPU.
    accelerator-type
    Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
    version
    Versi software Cloud TPU. Untuk semua jenis TPU, gunakan tpu-ubuntu2204-base.

Menghubungkan ke VM Cloud TPU

Lakukan SSH ke VM TPU menggunakan perintah berikut: 

$ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

Kolom wajib diisi

tpu_name
Nama VM TPU yang Anda hubungkan.
zone
Zona tempat Anda membuat Cloud TPU.

Menginstal JAX di VM Cloud TPU

(vm)$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Pemeriksaan sistem

Verifikasi bahwa JAX dapat mengakses TPU dan dapat menjalankan operasi dasar:

Mulai penafsir Python 3:

(vm)$ python3
 
>>> import jax

Menampilkan jumlah core TPU yang tersedia:

>>> jax.device_count()

Jumlah core TPU ditampilkan. Jika Anda menggunakan TPU v4, nilai ini harus 4. Jika Anda menggunakan TPU v2 atau v3, nilai ini harus 8.

Lakukan penghitungan sederhana:

>>> jax.numpy.add(1, 1)

Hasil penambahan numpy ditampilkan:

Output dari perintah: 

Array(2, dtype=int32, weak_type=true)

Keluar dari penafsiran Python:

>>> exit()

Menjalankan kode JAX di VM TPU

Sekarang Anda dapat menjalankan kode JAX yang diinginkan. Contoh flax adalah tempat yang tepat untuk memulai menjalankan model ML standar di JAX. Misalnya, untuk melatih jaringan konvolusi MNIST dasar:

  1. Menginstal dependensi contoh Flax

    (vm)$ pip install --upgrade clu
(vm)$ pip install tensorflow
(vm)$ pip install tensorflow_datasets

  2. Menginstal FLAX

    (vm)$ git clone https://github.com/google/flax.git
(vm)$ pip install --user flax

  3. Menjalankan skrip pelatihan FLAX MNIST

    (vm)$ cd flax/examples/mnist
(vm)$ python3 main.py --workdir=/tmp/mnist \
--config=configs/default.py \
--config.learning_rate=0.05 \
--config.num_epochs=5

Skrip akan mendownload set data dan memulai pelatihan. Output skrip akan terlihat seperti ini:

  0214 18:00:50.660087 140369022753856 train.py:146] epoch:  1, train_loss: 0.2421, train_accuracy: 92.97, test_loss: 0.0615, test_accuracy: 97.88
  I0214 18:00:52.015867 140369022753856 train.py:146] epoch:  2, train_loss: 0.0594, train_accuracy: 98.16, test_loss: 0.0412, test_accuracy: 98.72
  I0214 18:00:53.377511 140369022753856 train.py:146] epoch:  3, train_loss: 0.0418, train_accuracy: 98.72, test_loss: 0.0296, test_accuracy: 99.04
  I0214 18:00:54.727168 140369022753856 train.py:146] epoch:  4, train_loss: 0.0305, train_accuracy: 99.06, test_loss: 0.0257, test_accuracy: 99.15
  I0214 18:00:56.082807 140369022753856 train.py:146] epoch:  5, train_loss: 0.0252, train_accuracy: 99.20, test_loss: 0.0263, test_accuracy: 99.18

Pembersihan

Setelah selesai menggunakan VM TPU, ikuti langkah-langkah berikut untuk membersihkan resource.

  1. Putuskan koneksi dari instance Compute Engine, jika Anda belum melakukannya:

    (vm)$ exit

  2. Hapus Cloud TPU Anda.

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central1-a

  3. Pastikan resource telah dihapus dengan menjalankan perintah berikut. Pastikan TPU Anda tidak lagi tercantum. Proses penghapusan mungkin memerlukan waktu beberapa menit.

    $ gcloud compute tpus tpu-vm list \
  --zone=us-central1-a

Catatan Performa

Berikut beberapa detail penting yang sangat relevan dengan penggunaan TPU di JAX.

Padding

Salah satu penyebab paling umum untuk performa lambat di TPU adalah memperkenalkan padding yang tidak disengaja:

  • Array di Cloud TPU disusun dalam ubin. Hal ini memerlukan padding salah satu dimensi ke kelipatan 8, dan dimensi yang berbeda ke kelipatan 128.
  • Unit perkalian matriks berperforma terbaik dengan pasangan matriks besar yang meminimalkan kebutuhan padding.

dtype bfloat16

Secara default, perkalian matriks di JAX pada TPU menggunakan bfloat16 dengan akumulasi float32. Hal ini dapat dikontrol dengan argumen presisi pada panggilan fungsi jax.numpy yang relevan (matmul, dot, einsum, dll.). Pada khususnya:

  • precision=jax.lax.Precision.DEFAULT: menggunakan presisi bfloat16 campuran (tercepat)
  • precision=jax.lax.Precision.HIGH: menggunakan beberapa kartu MXU untuk mencapai presisi yang lebih tinggi
  • precision=jax.lax.Precision.HIGHEST: menggunakan lebih banyak kartu MXU untuk mencapai presisi float32 penuh

JAX juga menambahkan dtype bfloat16, yang dapat Anda gunakan untuk secara eksplisit mentransmisikan array ke bfloat16, misalnya, jax.numpy.array(x, dtype=jax.numpy.bfloat16).

Menjalankan JAX di Colab

Saat Anda menjalankan kode JAX di notebook Colab, Colab akan otomatis membuat node TPU lama. Node TPU memiliki arsitektur yang berbeda. Untuk mengetahui informasi selengkapnya, lihat Arsitektur Sistem.

Langkah selanjutnya

Untuk mengetahui informasi selengkapnya tentang Cloud TPU, lihat: