Menjalankan kode PyTorch pada slice Pod TPU
PyTorch/XLA mengharuskan semua VM TPU dapat mengakses data dan kode model. Anda dapat menggunakan skrip startup untuk mendownload software yang diperlukan untuk mendistribusikan data model ke semua VM TPU.
Jika menghubungkan VM TPU ke Virtual Private Cloud (VPC), Anda harus menambahkan aturan firewall dalam project agar traffic masuk ke port 8470 - 8479. Untuk mengetahui informasi lebih lanjut tentang cara menambahkan aturan firewall, lihat Menggunakan aturan firewall
Menyiapkan lingkungan Anda
Di Cloud Shell, jalankan perintah berikut untuk memastikan Anda menjalankan versi
gcloud
saat ini:$ gcloud components update
Jika Anda perlu menginstal
gcloud
, gunakan perintah berikut:$ sudo apt install -y google-cloud-sdk
Buat beberapa variabel lingkungan:
$ export PROJECT_ID=project-id $ export TPU_NAME=tpu-name $ export ZONE=us-central2-b $ export RUNTIME_VERSION=tpu-ubuntu2204-base $ export ACCELERATOR_TYPE=v4-32
Membuat VM TPU
$ gcloud compute tpus tpu-vm create ${TPU_NAME} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --accelerator-type=${ACCELERATOR_TYPE} \ --version ${RUNTIME_VERSION}
Mengonfigurasi dan menjalankan skrip pelatihan
Tambahkan sertifikat SSH ke project Anda:
ssh-add ~/.ssh/google_compute_engine
Menginstal PyTorch/XLA di semua pekerja VM TPU
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --worker=all --command=" pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"
Clone XLA di semua pekerja VM TPU
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --worker=all --command="git clone -b r2.2 https://github.com/pytorch/xla.git"
Menjalankan skrip pelatihan pada semua pekerja
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --worker=all \ --command="PJRT_DEVICE=TPU python3 ~/xla/test/test_train_mp_imagenet.py \ --fake_data \ --model=resnet50 \ --num_epochs=1 2>&1 | tee ~/logs.txt"
Pelatihan ini memerlukan waktu sekitar 5 menit. Setelah selesai, Anda akan melihat pesan yang mirip dengan ini:
Epoch 1 test end 23:49:15, Accuracy=100.00 10.164.0.11 [0] Max Accuracy: 100.00%
Pembersihan
Setelah selesai menggunakan VM TPU, ikuti langkah-langkah berikut untuk membersihkan resource Anda.
Putuskan koneksi dari Compute Engine:
(vm)$ exit
Verifikasi bahwa 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 europe-west4-a