Pelatihan Cloud TPU v5e

Cloud TPU v5e adalah akselerator AI generasi terbaru dari Google Cloud. Dengan jejak 256 chip yang lebih kecil per Pod, v5e dioptimalkan agar menjadi produk bernilai tertinggi untuk pelatihan, fine-tuning, dan inferensi Jaringan Neural Konvolusional (CNN) dan transformator, teks-ke-gambar.

Konsep, arsitektur sistem, dan konfigurasi Cloud TPU v5e

Jika Anda baru menggunakan Cloud TPU, lihat beranda dokumentasi TPU.

Konsep umum Cloud TPU (misalnya, slice, host, chip, dan TensorCore), serta arsitektur sistem Cloud TPU dijelaskan di halaman Arsitektur Sistem Cloud TPU.

Setiap versi Cloud TPU memerlukan jenis akselerator khusus untuk pelatihan dan inferensi. Jenis akselerator ini dijelaskan di bagian konfigurasi v5e.

Inferensi

Inferensi adalah proses penggunaan model terlatih untuk membuat prediksi pada data baru. ID ini digunakan oleh proses penayangan.

Irisan

Slice mewakili kumpulan chip yang semuanya berada di dalam Pod yang sama dan terhubung melalui interkoneksi antar-chip berkecepatan tinggi (ICI). v5e memiliki bentuk slice 2D. Lihat tabel di bagian konfigurasi v5e untuk mengetahui bentuk slice yang didukung.

Bentuk chip dan topologi chip juga mengacu pada bentuk irisan.

Inferensi

Penayangan adalah proses men-deploy model machine learning terlatih ke lingkungan produksi tempat model tersebut dapat digunakan untuk membuat prediksi atau keputusan. Latensi dan ketersediaan tingkat layanan penting untuk penayangan.

Host tunggal versus multi host

Host adalah komputer fisik (CPU) yang menjalankan VM. Host dapat menjalankan beberapa VM sekaligus.

Slice yang menggunakan kurang dari 8 chip menggunakan paling banyak satu host. Slice yang lebih besar dari 8 chip memiliki akses ke lebih dari satu host dan dapat menjalankan pelatihan terdistribusi menggunakan beberapa host. Lihat halaman Arsitektur Sistem TPU untuk mengetahui detail tentang slice dan chip.

v5e mendukung pelatihan multi-host dan inferensi multi-host (menggunakan SAX).

VM TPU

Virtual machine yang menjalankan Linux yang memiliki akses ke TPU yang mendasarinya. Untuk TPU v5e, setiap VM TPU memiliki akses langsung ke 1, 4, atau 8 chip, bergantung pada jenis akselerator yang ditentukan pengguna. VM TPU juga dikenal sebagai pekerja.

Pekerja

Lihat TPU VM.

Mulai

Untuk informasi tentang hardware TPU v5e, lihat Arsitektur Sistem.

Mengamankan kapasitas

Hubungi Bagian Penjualan Cloud untuk mulai menggunakan Cloud TPU v5e untuk workload AI Anda.

Menyiapkan Project Google Cloud

1. Login ke Akun Google Anda. Jika Anda belum melakukannya, daftar untuk membuat akun baru.
2. Di Konsol Google Cloud, pilih atau buat project Google Cloud dari halaman pemilih project.

3. Penyiapan penagihan diperlukan untuk semua penggunaan Google Cloud, jadi pastikan penagihan diaktifkan untuk project Anda.

4. Instal gcloud alpha components.

5. Jika Anda pengguna TPU yang menggunakan kembali komponen gcloud alpha yang ada, perbarui komponen ini untuk memastikan bahwa perintah dan flag yang relevan didukung:

   gcloud components update
   

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

   gcloud services enable tpu.googleapis.com
   

7. Aktifkan akun layanan TPU.

   Akun layanan memungkinkan layanan Cloud TPU mengakses layanan Google Cloud lainnya. Akun layanan yang dikelola pengguna adalah praktik Google Cloud yang direkomendasikan. Ikuti panduan ini untuk membuat dan memberikan peran. Peran berikut diperlukan:

   • TPU Admin
   • Storage Admin: Diperlukan untuk mengakses Cloud Storage
   • Penulis Log: Diperlukan untuk menulis log dengan Logging API
   • Penulis Metrik Monitoring: Diperlukan untuk menulis metrik ke Cloud Monitoring

8. Konfigurasikan project dan zona.

   Project ID Anda adalah nama project Anda yang ditampilkan di Cloud Console.

   export PROJECT_ID=your-project-id
   export ZONE=us-west4-a
   
   gcloud compute tpus tpu-vm service-identity create --zone=${ZONE}
   
   gcloud auth login
   gcloud config set project ${PROJECT_ID}
   gcloud config set compute/zone ${ZONE}

Menyediakan lingkungan Cloud TPU

Praktik terbaiknya adalah menyediakan Cloud TPU v5es sebagai resource yang diantrekan menggunakan perintah queued-resource create. Namun, Anda juga dapat menggunakan Create Node API (gcloud alpha compute tpus tpu-vm create) untuk menyediakan Cloud TPU v5es.

Membuat variabel lingkungan

Tetapkan variabel lingkungan yang diperlukan untuk pembuatan TPU.

Ganti variabel (berwarna merah) dalam daftar berikut dengan nilai yang akan Anda gunakan untuk tugas pelatihan atau inferensi.

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

Variabel	Deskripsi
PROJECT_ID	Nama Project Google Cloud
ACCELERATOR_TYPE	Lihat bagian Jenis Akselerator untuk jenis akselerator yang didukung.
ZONA	Semua kapasitas dalam us-west4-a.
RUNTIME_VERSION	Gunakan v2-alpha-tpuv5-lite.
SERVICE_ACCOUNT	Alamat ini adalah alamat akun layanan yang dapat Anda temukan di Konsol Google Cloud -> IAM -> Akun Layanan. Misalnya: [e-mail]
TPU_NAME	ID teks TPU yang ditetapkan pengguna yang dibuat saat permintaan resource dalam antrean dialokasikan.
QUEUED_RESOURCE_ID	ID teks yang ditetapkan pengguna dari permintaan aset dalam antrean. Lihat dokumen resource yang diantrekan untuk informasi tentang resource yang diantrekan.
QUOTA_TYPE	Dapat berupa reserved atau best-effort. Jika tidak satu pun yang ditentukan, QUOTA_TYPE akan ditetapkan secara default ke on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.
VALID_UNTIL_DURATION	Durasi saat permintaan valid. Lihat resource dalam antrean untuk mengetahui informasi tentang durasi valid yang berbeda.

Membuat resource TPU

gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

Jika resource yang diantrekan berhasil dibuat, status dalam kolom response akan menjadi WAITING_FOR_RESOURCES atau FAILED. Jika resource yang diantrekan berada dalam status WAITING_FOR_RESOURCES, artinya resource yang diantrekan telah lulus validasi awal dan sedang menunggu kapasitas. Setelah kapasitas tersedia, permintaan akan ditransisikan ke PROVISIONING. Berada dalam status WAITING_FOR_RESOURCES tidak berarti sudah pasti bahwa Anda akan mendapatkan kuota yang dialokasikan dan mungkin perlu waktu untuk berubah dari status WAITING_FOR_RESOURCES ke status ACTIVE. Jika resource yang diantrekan berada dalam status FAILED, alasan kegagalan akan ada di output. Permintaan akan berakhir jika permintaan tidak diisi di --valid-until-duration, dan statusnya menjadi "FAILED".

Anda akan dapat mengakses VM TPU menggunakan SSH setelah resource yang diantrekan berada dalam status ACTIVE.

Gunakan perintah [list](/tpu/docs/managing-tpus-tpu-vm) atau [describe](/tpu/docs/managing-tpus-tpu-vm) untuk membuat kueri status resource dalam antrean.

gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

state menunjukkan status resource dalam antrean. Status ditentukan sebagai:

Status	Deskripsi
WAITING_FOR_RESOURCES	Perintah create resource yang diantrekan telah diterima dan akan memulai penyediaan, segera setelah kapasitas tersedia.
PROVISIONING	TPU Slice sedang disediakan.
ACTIVE	Semua TPU disediakan dan siap digunakan. Jika diberikan, skrip startup akan mulai mengeksekusi di semua TPU, saat status resource yang diantrekan beralih ke ACTIVE.
FAILED	Slice tidak dapat disediakan.
SUSPENDING	Satu atau beberapa slice sedang dihapus.
SUSPENDED	Semua slice yang mendasarinya akan dihapus, tetapi resource yang diantrekan tetap utuh, hingga dihapus secara eksplisit. Saat ini, resource dalam antrean yang ditangguhkan tidak dapat dilanjutkan dan harus dihapus.
DELETING	Resource yang diantrekan sedang dihapus.

Menghubungkan ke VM TPU menggunakan SSH

Bagian berikut ini menjelaskan cara menginstal biner di setiap VM TPU di slice TPU dan menjalankan kode. Dalam konteks ini, VM TPU juga dikenal sebagai pekerja.

Lihat bagian Jenis VM untuk menentukan jumlah VM yang akan dimiliki slice Anda.

Untuk menginstal biner atau menjalankan kode, hubungkan ke VM TPU Anda menggunakan perintah ssh tpu-vm.

gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Untuk mengakses VM TPU tertentu dengan SSH, gunakan tanda --worker yang mengikuti indeks berbasis 0:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} --worker=1

Jika memiliki bentuk slice yang lebih dari 8 chip, Anda akan memiliki beberapa VM dalam satu slice. Dalam hal ini, gunakan flag --worker=all untuk menjalankan penginstalan pada semua VM TPU tanpa harus terhubung ke setiap VM secara terpisah. Contoh:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Kelola

Semua perintah yang dapat Anda gunakan untuk mengelola VM TPU dijelaskan dalam Mengelola TPU.

Penyiapan Framework

Bagian ini menjelaskan proses penyiapan umum untuk pelatihan model kustom menggunakan JAX atau PyTorch dengan TPU v5e. Dukungan TensorFlow tersedia di versi runtime TPU tpu-vm-tf-2.15.0-pjrt dan tpu-vm-tf-2.15.0-pod-pjrt.

Untuk petunjuk penyiapan inferensi, lihat Pengantar inferensi v5e.

Penyiapan untuk JAX

Jika memiliki bentuk slice yang lebih dari 8 chip, Anda akan memiliki beberapa VM dalam satu slice. Dalam hal ini, Anda harus menggunakan flag --worker=all untuk menjalankan penginstalan pada semua VM TPU dalam satu langkah tanpa menggunakan SSH untuk login ke setiap VM secara terpisah:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Anda dapat menjalankan perintah berikut untuk memeriksa jumlah perangkat (output yang ditampilkan di sini dibuat dengan slice v5litepod-16). Kode ini menguji bahwa semuanya telah diinstal dengan benar dengan memastikan bahwa JAX melihat TensorCore Cloud TPU dan dapat menjalankan operasi dasar:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Outputnya akan mirip dengan berikut ini:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device\_count() menunjukkan jumlah total chip dalam bagian yang ditentukan. jax.local\_device\_count() menunjukkan jumlah chip yang dapat diakses oleh satu VM dalam slice ini.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

Outputnya akan mirip dengan berikut ini:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Coba Tutorial JAX dalam dokumen ini untuk memulai pelatihan v5e menggunakan JAX.

Penyiapan untuk PyTorch

Perlu diperhatikan bahwa v5e hanya mendukung runtime PJRT dan PyTorch 2.1+ akan menggunakan PJRT sebagai runtime default untuk semua versi TPU.

Bagian ini menjelaskan cara mulai menggunakan PJRT di v5e dengan PyTorch/XLA dengan perintah untuk semua pekerja.

Menginstal dependensi

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

Jika Anda gagal menginstal roda untuk torch/torch_xla/torchvision dan melihat error seperti pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, downgrade versi Anda dengan perintah ini:

pip3 install setuptools==62.1.0

Jalankan skrip dengan PJRT:

Catatan: Untuk model yang memiliki alokasi sering dan cukup besar, menggunakan tcmalloc dapat meningkatkan waktu pelatihan secara signifikan dibandingkan dengan implementasi malloc default, sehingga malloc default yang digunakan pada VM TPU adalah tcmalloc. Namun, bergantung pada beban kerja Anda (misalnya, dengan DLRM yang memiliki alokasi yang sangat besar untuk tabel embedding-nya) tcmalloc dapat menyebabkan pelambatan. Dalam hal ini, Anda dapat mencoba membatalkan penetapan variabel berikut menggunakan malloc default:

unset LD_PRELOAD

Berikut adalah contoh penggunaan skrip Python untuk melakukan penghitungan di VM v5e:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

Tindakan ini menghasilkan output yang mirip dengan berikut ini:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

Coba Tutorial PyTorch dalam dokumen ini untuk memulai pelatihan v5e menggunakan PyTorch.

Hapus TPU dan resource yang diantrekan di akhir sesi Anda. Untuk menghapus resource yang diantrekan, hapus slice, lalu resource yang diantrekan dengan 2 langkah:

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Kedua langkah ini juga dapat digunakan untuk menghapus permintaan resource dalam antrean yang berada dalam status FAILED.

Memantau dan Membuat Profil

Cloud TPU v5e mendukung pemantauan dan pembuatan profil menggunakan metode yang sama seperti Cloud TPU generasi sebelumnya. Anda dapat membaca Membuat profil model Anda dengan alat Cloud TPU untuk mempelajari lebih lanjut pembuatan profil dan Memantau VM Cloud TPU untuk mempelajari pemantauan lebih lanjut.

Contoh JAX/FLAX

Melatih ImageNet di v5e

Tutorial ini menjelaskan cara melatih ImageNet pada v5e menggunakan data input palsu. Jika Anda ingin menggunakan data asli, lihat file README di GitHub.

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat melakukan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Saat QueuedResource berada dalam status ACTIVE, outputnya akan mirip dengan berikut:

    state: ACTIVE
   

3. Instal JAX dan jaxlib versi terbaru:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   Catatan: Jika melihat error yang menyarankan Anda menggunakan jax>=0.4.8, Anda dapat mengabaikan pesan tersebut dengan aman.

4. Clone model ImageNet dan instal persyaratan yang sesuai:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'
   

5. Untuk menghasilkan data palsu, model tersebut memerlukan informasi tentang dimensi set data. Ini dapat dikumpulkan dari metadata {i>dataset <i}ImageNet:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} --zone=${ZONE} --worker=all --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'
   

Melatih model

Setelah semua langkah sebelumnya selesai, Anda dapat melatih model.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

Menghapus TPU dan resource yang diantrekan

Hapus TPU dan resource yang diantrekan di akhir sesi Anda.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Model FLAX Wajah Memeluk

Model Wajah Mengangkat yang diimplementasikan dalam FLAX dapat langsung difungsikan di Cloud TPU v5e. Bagian ini berisi petunjuk untuk menjalankan model populer.

Melatih ViT di Imagenette

Tutorial ini menunjukkan cara melatih model Vision Transformer (ViT) dari HuggingFace menggunakan set data imagenette Fast AI di Cloud TPU v5e.

Model ViT adalah model pertama yang berhasil melatih encoder Transformer di ImageNet dengan hasil yang sangat baik dibandingkan dengan jaringan konvolusional. Untuk mengetahui informasi selengkapnya, lihat referensi berikut:

• Ringkasan ViT
• Resource kode ViT

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat menjalankan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Saat resource yang diantrekan berada dalam status ACTIVE, outputnya akan serupa dengan berikut ini:

    state: ACTIVE
   

3. Instal JAX beserta library-nya:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Download repositori dan persyaratan penginstalan Hugging Face:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.16.1 && pip install -r examples/flax/vision/requirements.txt'
   

5. Unduh {i>dataset<i} Imagenette:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
   

Melatih model

Latih model dengan buffer yang telah dipetakan sebelumnya, berukuran 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

Menghapus TPU dan resource dalam antrean

Hapus TPU dan resource dalam antrean di akhir sesi Anda.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Hasil tolok ukur ViT

Skrip pelatihan dijalankan pada v5litepod-4, v5litepod-16, dan v5litepod-64. Tabel berikut menunjukkan throughput dengan berbagai jenis akselerator.

Jenis akselerator	V5litepod-4	V5litepod-16	V5litepod-64
Zaman	3	3	3
Ukuran tumpukan global	32	128	512
Throughput (contoh/dtk)	263,40	429,34	470,71

Train Diffusion di Pokémon

Tutorial ini menunjukkan cara melatih model Diffusion Stabil dari HuggingFace menggunakan set data Pokémon di Cloud TPU v5e.

Model Difusi Stabil adalah model teks ke gambar laten yang menghasilkan gambar realistik foto dari input teks apa pun. Untuk informasi selengkapnya, lihat referensi berikut:

• Ringkasan Difusi Stabil
• Sumber kode Diffusion Stabil

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat menjalankan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Saat resource yang diantrekan berada dalam status ACTIVE, output-nya akan serupa dengan berikut ini:

    state: ACTIVE
   

3. Instal JAX beserta library-nya.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Download repositori HuggingFace dan persyaratan penginstalan.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.16.1 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
   

Melatih model

Latih model dengan buffer yang telah dipetakan sebelumnya, berukuran 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'

Menghapus TPU dan resource yang diantrekan

Hapus TPU dan resource yang diantrekan di akhir sesi Anda.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Hasil benchmark untuk difusi

Skrip pelatihan berjalan pada v5litepod-4, v5litepod-16, dan v5litepod-64. Tabel berikut menunjukkan throughput.

Jenis akselerator	V5litepod-4	V5litepod-16	V5litepod-64
Langkah Kereta	1500	1500	1500
Ukuran tumpukan global	32	64	128
Throughput (contoh/dtk)	36,53	43,71	49,36

Melatih GPT2 pada {i>dataset<i} OSCAR

Tutorial ini menunjukkan cara melatih model GPT2 dari HuggingFace menggunakan set data OSCAR pada Cloud TPU v5e.

GPT2 adalah model transformer yang dilatih sebelumnya pada teks mentah tanpa pelabelan manusia. Model ini dilatih untuk memprediksi kata berikutnya dalam kalimat. Untuk informasi selengkapnya, lihat resource berikut:

• Ringkasan GPT2
• Sumber kode GPT2

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat melakukan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Saat resource yang diantrekan berada dalam status ACTIVE, outputnya akan serupa dengan berikut ini:

    state: ACTIVE
   

3. Instal JAX beserta library-nya.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Download repositori HuggingFace dan persyaratan penginstalan.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/language-modeling/requirements.txt'
   

5. Download konfigurasi untuk melatih model.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'
   

Melatih model

Latih model dengan buffer yang telah dipetakan sebelumnya, berukuran 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

Menghapus TPU dan resource dalam antrean

Hapus TPU dan resource dalam antrean di akhir sesi Anda.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Hasil benchmark untuk GPT2

Skrip pelatihan berjalan pada v5litepod-4, v5litepod-16, dan v5litepod-64. Tabel berikut menunjukkan throughput.

	V5litepod-4	V5litepod-16	V5litepod-64
Zaman	3	3	3
Ukuran tumpukan global	64	64	64
Throughput (contoh/dtk)	74,60	72,97	72,62

PyTorch/XLA

Melatih ResNet menggunakan runtime PJRT

PyTorch/XLA bermigrasi dari XRT ke PjRt dari PyTorch 2.0+. Berikut adalah petunjuk yang telah diperbarui guna menyiapkan v5e untuk workload pelatihan PyTorch/XLA.

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --{QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat melakukan SSH ke VM TPU setelah QueuedResource berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Saat resource yang diantrekan berada dalam status ACTIVE, output-nya akan serupa dengan berikut ini:

    state: ACTIVE
   

3. Menginstal dependensi khusus Torch/XLA

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         sudo apt-get update -y
         sudo apt-get install libomp5 -y
         pip3 install mkl mkl-include
         pip3 install tf-nightly tb-nightly tbp-nightly
         pip3 install numpy
         sudo apt-get install libopenblas-dev -y
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

Melatih model ResNet

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

Menghapus TPU dan resource dalam antrean

Hapus TPU dan resource dalam antrean di akhir sesi Anda.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Hasil Tolok Ukur

Tabel berikut menunjukkan throughput benchmark.

Jenis akselerator	Throughput (contoh/detik)
V5litepod-4	4.240 eks/dtk
V5litepod-16	10.810 ex/dtk
V5litepod-64	46.154 ex/dtk

Melatih GPT2 di v5e

Tutorial ini akan membahas cara menjalankan GPT2 di v5e menggunakan repositori HuggingFace pada PyTorch/XLA menggunakan set data wikitext.

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda akan dapat melakukan SSH ke VM TPU setelah QueuedResource berada dalam status ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Saat resource yang diantrekan berada dalam status ACTIVE, outputnya akan mirip dengan berikut ini:

   state: ACTIVE
   

3. Instal dependensi torch/xla.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         sudo apt-get -y update
         sudo apt install -y libopenblas-base
         pip3 install torchvision
         pip3 uninstall -y torch
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

4. Download repo dan persyaratan penginstalan HuggingFace.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         git clone https://github.com/pytorch/xla.git
         pip install --upgrade accelerate
         git clone https://github.com/huggingface/transformers.git
         cd transformers
         git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
         pip install .
         git log -1
         pip install datasets evaluate scikit-learn
         '
   

5. Download konfigurasi model terlatih.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/
         gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
   

Melatih model

Latih model 2B menggunakan ukuran tumpukan 16.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

Menghapus TPU dan resource dalam antrean

Hapus TPU dan resource dalam antrean di akhir sesi Anda.

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Hasil Tolok Ukur

Skrip pelatihan berjalan pada v5litepod-4, v5litepod-16, dan v5litepod-64. Tabel berikut menunjukkan throughput benchmark untuk berbagai jenis akselerator.

	V5litepod-4	V5litepod-16	V5litepod-64
Zaman	3	3	3
config	600 JT	2 M	16 M
Ukuran tumpukan global	64	128	256
Throughput (contoh/dtk)	66	77	31

Latih ViT di v5e

Tutorial ini akan membahas cara menjalankan VIT di v5e menggunakan repositori HuggingFace pada PyTorch/XLA pada set data cifar10.

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat melakukan SSH ke VM TPU setelah QueuedResource berada dalam status ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Saat resource yang diantrekan berada dalam status ACTIVE, outputnya akan mirip dengan berikut ini:

    state: ACTIVE
   

3. Menginstal dependensi torch/xla

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all
      --command='
         sudo apt-get update -y
         sudo apt-get install libomp5 -y
         pip3 install mkl mkl-include
         pip3 install tf-nightly tb-nightly tbp-nightly
         pip3 install numpy
         sudo apt-get install libopenblas-dev -y
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

4. Download repo dan persyaratan penginstalan HuggingFace.

      gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command="
         git clone https://github.com/suexu1025/transformers.git vittransformers; \
         cd vittransformers; \
         pip3 install .; \
         pip3 install datasets; \
         wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"
   

Melatih model

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

Menghapus TPU dan resource dalam antrean

Hapus TPU dan resource dalam antrean di akhir sesi Anda.

gcloud compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

Hasil Tolok Ukur

Tabel berikut menunjukkan throughput benchmark untuk berbagai jenis akselerator.

	V5litepod-4	V5litepod-16	V5litepod-64
Zaman	3	3	3
Ukuran tumpukan global	32	128	512
Throughput (contoh/dtk)	201	657	2.844

TensorFlow 2.x

Melatih Resnet pada satu host v5e

Tutorial ini menjelaskan cara melatih ImageNet pada v5litepod-4 atau v5litepod-8 menggunakan set data palsu. Jika Anda ingin menggunakan set data lain, lihat Menyiapkan set data.

Penyiapan

1. Buat variabel lingkungan:

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5litepod-4
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   ACCELERATOR_TYPE dapat berupa v5litepod-4 atau v5litepod-8. /

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat melakukan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE. Untuk memeriksa status resource yang diantrekan, gunakan perintah berikut:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

3. Menghubungkan ke TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

4. Menetapkan beberapa variabel lingkungan

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
   export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

5. Ubah ke direktori repositori model dan persyaratan penginstalan.

   cd ${MODELS_REPO} && git checkout r2.15.0
   pip install -r official/requirements.txt
   

Melatih model

Jalankan skrip pelatihan.

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Menghapus TPU dan resource dalam antrean

1. Menghapus TPU Anda

   gcloud compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

2. Menghapus permintaan sumber daya yang diantrekan

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

Melatih Resnet pada multi host v5e

Tutorial ini menjelaskan cara melatih ImageNet pada v5litepod-16 atau yang lebih besar menggunakan set data palsu. Jika Anda ingin menggunakan set data yang berbeda, lihat bagian Menyiapkan set data.

1. Buat variabel lingkungan:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   ACCELERATOR_TYPE dapat v5litepod-16 atau lebih besar.

2. Buat resource TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   Catatan: Flag QUOTA_TYPE dapat berupa reserved atau best-effort. Jika tidak satu pun dari opsi tersebut yang ditentukan, defaultnya adalah on-demand. Lihat quotas untuk mengetahui informasi tentang berbagai jenis kuota yang didukung oleh Cloud TPU.

   Anda dapat melakukan SSH ke VM TPU setelah resource yang diantrekan berada dalam status ACTIVE. Untuk memeriksa status resource yang diantrekan, gunakan perintah berikut:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

3. Menghubungkan ke TPU (pekerja nol) menggunakan SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

4. Menetapkan beberapa variabel lingkungan

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export TPU_LOAD_LIBRARY=0
   export TPU_NAME=your_tpu_name

5. Ubah ke direktori repositori model dan persyaratan penginstalan.

    cd $MODELS_REPO && git checkout r2.15.0
    pip install -r official/requirements.txt
   

Melatih model

Jalankan skrip pelatihan.

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

Menghapus TPU dan resource dalam antrean

1. Menghapus TPU Anda

   gcloud compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

2. Menghapus permintaan sumber daya yang diantrekan

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet