Menjalankan workload TPU di container Docker

Container Docker mempermudah konfigurasi aplikasi dengan menggabungkan kode dan semua dependensi yang diperlukan dalam satu paket yang dapat didistribusikan. Anda dapat menjalankan container Docker dalam VM TPU untuk menyederhanakan konfigurasi dan berbagi aplikasi Cloud TPU. Dokumen ini menjelaskan cara menyiapkan container Docker untuk setiap framework ML yang didukung oleh Cloud TPU.

Penting: Anda dapat menjalankan maksimum 256 pod Docker di VM TPU.

Melatih model PyTorch di container Docker

Perangkat TPU

1. Buat VM Cloud TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. Hubungkan ke VM TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

3. Pastikan pengguna Anda telah diberi peran Artifact Registry Reader. Google Cloud Untuk mengetahui informasi selengkapnya, lihat Memberikan peran Artifact Registry.

4. Mulai container di VM TPU menggunakan image PyTorch/XLA nightly

   sudo docker run --net=host -ti --rm --name your-container-name --privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
   bash

   Catatan: Setelah menjalankan perintah ini, command prompt Anda akan berubah yang menunjukkan bahwa terminal terhubung ke container yang sedang berjalan.

5. Mengonfigurasi runtime TPU

   Ada dua opsi runtime PyTorch/XLA: PJRT dan XRT. Sebaiknya Anda menggunakan PJRT kecuali jika Anda memiliki alasan untuk menggunakan XRT. Untuk mempelajari lebih lanjut berbagai konfigurasi runtime, lihat dokumentasi runtime PJRT.

   PJRT

   export PJRT_DEVICE=TPU

   XRT

   export XRT_TPU_CONFIG="localservice;0;localhost:51011"

6. Buat clone repositori PyTorch XLA

   git clone --recursive https://github.com/pytorch/xla.git

7. Latih ResNet50

   python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource.

1. Ketik exit untuk keluar dari container Docker
2. Ketik exit untuk keluar dari VM TPU

3. Hapus VM TPU

   gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Slice TPU

Saat menjalankan kode PyTorch pada slice TPU, Anda harus menjalankan kode pada semua pekerja TPU secara bersamaan. Salah satu cara untuk melakukannya adalah dengan menggunakan perintah gcloud compute tpus tpu-vm ssh dengan flag --worker=all dan --command. Prosedur berikut menunjukkan cara membuat image Docker untuk mempermudah penyiapan setiap pekerja TPU.

1. Membuat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v4-32 \
   --version=tpu-ubuntu2204-base

2. Tambahkan pengguna saat ini ke grup Docker

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=us-central2-b \
   --worker=all \
   --command='sudo usermod -a -G docker $USER'

3. Buat clone repositori PyTorch XLA

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="git clone --recursive https://github.com/pytorch/xla.git"

4. Jalankan skrip pelatihan di container pada semua pekerja TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="docker run --rm --privileged --net=host  -v ~/xla:/xla -e PJRT_DEVICE=TPU us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 python /xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

   Flag perintah Docker:

   • --rm menghapus container setelah prosesnya berakhir.
   • --privileged mengekspos perangkat TPU ke container.
   • --net=host mengikat semua port container ke VM TPU untuk memungkinkan komunikasi antara host di pod.
   • -e menetapkan variabel lingkungan.

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource.

Hapus VM TPU menggunakan perintah berikut:

gcloud compute tpus tpu-vm delete your-tpu-name \
--zone=us-central2-b

Melatih model JAX di container Docker

Perangkat TPU

1. Buat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. Hubungkan ke VM TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a

3. Mulai daemon Docker di VM TPU

   sudo systemctl start docker

4. Memulai container Docker

   sudo docker run --net=host -ti --rm --name your-container-name \
   --privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
   bash

5. Menginstal JAX

   pip install jax[tpu]

6. Menginstal FLAX

   pip install --upgrade clu
   git clone https://github.com/google/flax.git
   pip install --user -e flax

7. Instal paket tensorflow dan tensorflow-dataset

   pip install tensorflow
   pip install tensorflow-datasets

8. Jalankan skrip pelatihan FLAX MNIST

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

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource.

1. Ketik exit untuk keluar dari container Docker
2. Ketik exit untuk keluar dari VM TPU

3. Hapus VM TPU

   gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Slice TPU

Saat menjalankan kode JAX pada slice TPU, Anda harus menjalankan kode JAX di semua pekerja TPU secara bersamaan. Salah satu cara untuk melakukannya adalah dengan menggunakan perintah gcloud compute tpus tpu-vm ssh dengan flag --worker=all dan --command. Prosedur berikut menunjukkan cara membuat image Docker untuk mempermudah penyiapan setiap pekerja TPU.

1. Buat file bernama Dockerfile di direktori saat ini dan tempelkan teks berikut

   FROM python:3.10
   RUN pip install jax[tpu]
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   RUN pip install tensorflow
   RUN pip install tensorflow-datasets
   WORKDIR ./flax/examples/mnist

2. Siapkan Artifact Registry

   gcloud artifacts repositories create your-repo \
   --repository-format=docker \
   --location=europe-west4 --description="Docker repository" \
   --project=your-project
   
   gcloud artifacts repositories list \
   --project=your-project
   
   gcloud auth configure-docker europe-west4-docker.pkg.dev

3. Membangun gambar Docker

   docker build -t your-image-name .

4. Tambahkan tag ke image Docker Anda sebelum mengirimkannya ke Artifact Registry. Untuk mengetahui informasi selengkapnya tentang cara menggunakan Artifact Registry, lihat Bekerja dengan image container.

   docker tag your-image-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

5. Mengirim image Docker Anda ke Artifact Registry

   docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

6. Membuat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

7. Tarik image Docker dari Artifact Registry di semua pekerja TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command='sudo usermod -a -G docker ${USER}'

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker pull europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag"

8. Jalankan container di semua pekerja TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag bash"

9. Jalankan skrip pelatihan di semua pekerja TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5"

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource.

1. Menonaktifkan container di semua pekerja

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker kill your-container-name"

2. Hapus VM TPU

   gcloud compute tpus tpu-vm delete your-tpu-name \
   --zone=europe-west4-a

Langkah berikutnya

• Tutorial Cloud TPU
• Mengelola TPU
• Arsitektur Cloud TPU
• Menjalankan kode JAX pada slice TPU