Dokumen ini memberikan pengantar singkat tentang cara menggunakan Pax di TPU host tunggal (v2-8, v3-8, v4-8).
Pax adalah framework untuk mengonfigurasi dan menjalankan eksperimen machine learning di atas JAX. Pax berfokus pada menyederhanakan ML dalam skala besar dengan berbagi komponen infrastruktur dengan framework ML yang ada dan memanfaatkan library pemodelan Praxis untuk modularitas.
Tujuan
- Menyiapkan resource TPU untuk pelatihan
- Menginstal Pax di TPU host tunggal
- Melatih model SPMD berbasis transformer menggunakan Pax
Sebelum memulai
Jalankan perintah berikut untuk mengonfigurasi gcloud agar dapat menggunakan project Cloud TPU dan menginstal komponen yang diperlukan untuk melatih model yang menjalankan Pax pada TPU host tunggal.
Menginstal Google Cloud CLI.
Google Cloud CLI berisi alat dan library untuk berinteraksi dengan produk dan layanan Google Cloud CLI. Jika Anda belum pernah menginstalnya, instal sekarang dengan mengikuti petunjuk dalam Menginstal Google Cloud CLI.
Mengonfigurasi perintah gcloud
(Jalankan gcloud auth list untuk melihat akun Anda yang tersedia).
$ gcloud config set account account
$ gcloud config set project project-id
Mengaktifkan Cloud TPU API
Aktifkan Cloud TPU API menggunakan perintah gcloud berikut di Cloud Shell.
(Anda juga dapat mengaktifkannya dari Google Cloud Console).
$ gcloud services enable tpu.googleapis.com
Jalankan perintah berikut untuk membuat identitas layanan (akun layanan).
$ gcloud beta services identity create --service tpu.googleapis.com
Membuat VM TPU
Dengan VM Cloud TPU, model dan kode Anda berjalan langsung di VM TPU. Anda menerapkan SSH langsung ke VM TPU. Anda dapat menjalankan kode arbitrer, menginstal paket, melihat log, dan men-debug kode langsung di VM TPU.
Buat VM TPU dengan menjalankan perintah berikut dari Cloud Shell atau terminal komputer tempat Google Cloud CLI terinstal.
Tetapkan zone berdasarkan ketersediaan dalam kontrak Anda,
referensikan Region dan Zona TPU
jika diperlukan.
Tetapkan variabel accelerator-type ke v2-8, v3-8, atau v4-8.
Tetapkan variabel version ke tpu-vm-base untuk versi TPU v2 dan v3, atau tpu-vm-v4-base untuk TPU v4.
$ gcloud compute tpus tpu-vm create tpu-name \
--zone zone \
--accelerator-type accelerator-type \
--version version
Hubungkan ke VM Google Cloud TPU Anda
Jalankan SSH ke VM TPU Anda dengan menggunakan perintah berikut:
$ gcloud compute tpus tpu-vm ssh tpu-name --zone zone
Saat Anda login ke VM, perintah shell Anda akan berubah dari
username@projectname menjadi username@vm-name:
Menginstal Pax di VM Google Cloud TPU
Instal Pax, JAX, dan libtpu di VM TPU menggunakan perintah berikut:
(vm)$ python3 -m pip install -U pip \
python3 -m pip install paxml jax[tpu]
-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Pemeriksaan sistem
Uji apakah semuanya telah diinstal dengan benar dengan memastikan JAX melihat inti TPU:
(vm)$ python3 -c "import jax; print(jax.device_count())"
Jumlah inti TPU ditampilkan, seharusnya 8 jika Anda menggunakan v2-8 atau v3-8, atau 4 jika Anda menggunakan v4-8.
Menjalankan kode Pax di VM TPU
Sekarang Anda dapat menjalankan kode Pax yang diinginkan. Contoh lm_cloud adalah tempat yang tepat untuk mulai menjalankan model di Pax. Misalnya, perintah berikut melatih model bahasa SPMD berbasis transformator parameter 2B pada data sintetis.
Perintah berikut menunjukkan output pelatihan untuk model bahasa SPMD. Mobil ini melatih 300 langkah dalam waktu sekitar 20 menit.
(vm)$ python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir
Pada slice v4-8, output-nya harus mencakup:
Kerugian dan waktu langkah
Tensor ringkasan pada step=step_# loss = loss
tensor ringkasan pada step=step_# Langkah per detik x
Pembersihan
Agar tidak dikenakan biaya pada akun Google Cloud Anda untuk resource yang digunakan dalam tutorial ini, hapus project yang berisi resource tersebut, atau simpan project dan hapus setiap resource-nya.
Setelah selesai menggunakan VM TPU, ikuti langkah-langkah berikut untuk membersihkan resource Anda.
Putuskan koneksi dari instance Compute Engine jika Anda belum melakukannya:
(vm)$ exit
Hapus Cloud TPU Anda.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
Langkah selanjutnya
Untuk mengetahui informasi selengkapnya tentang Cloud TPU, lihat: