Menjalankan aplikasi Cloud TPU di container Docker

Penampung Docker memudahkan konfigurasi aplikasi dengan menggabungkan kode Anda 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 penampung Docker untuk setiap framework ML yang didukung oleh Cloud TPU.

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

Melatih model TensorFlow di container Docker

Perangkat TPU

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

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models

2. Membuat bucket Cloud Storage

   gcloud storage buckets create gs://your-bucket-name --location=europe-west4

3. Membuat VM TPU

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

4. Menyalin Dockerfile ke VM TPU

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

5. SSH ke VM TPU

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

6. Membangun gambar Docker

   sudo docker build -t your-image-name .

7. Memulai container Docker

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

8. Menetapkan variabel lingkungan

   export STORAGE_BUCKET=gs://your-bucket-name
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x

9. Melatih ResNet

   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="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"

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. Menghapus VM TPU

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

Pod TPU

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

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models

2. Membuat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v3-32 \
   --version=tpu-vm-tf-2.18.0-pod-pjrt

3. Menyalin Dockerfile ke VM TPU

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

4. SSH ke VM TPU

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

5. Membangun gambar Docker

   sudo docker build -t your-image-name .

6. Memulai container Docker

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

7. Melatih ResNet

   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="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"

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. Menghapus VM TPU

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

Melatih model PyTorch dalam container Docker

Perangkat TPU

1. Membuat 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. SSH ke VM TPU

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

3. Mulai container di VM TPU menggunakan image PyTorch/XLA harian.

   sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash

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

4. 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 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"

5. Meng-clone repo PyTorch XLA

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

6. Melatih 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. Menghapus VM TPU

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

Pod TPU

Saat menjalankan kode PyTorch di Pod TPU, Anda harus menjalankan kode di semua pekerja TPU secara bersamaan. Salah satu cara 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 memudahkan 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. Menambahkan 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. Jalankan skrip pelatihan dalam penampung di 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  -e PJRT_DEVICE=TPU gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm python /pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

   Flag perintah Docker:

   • --rm menghapus penampung setelah prosesnya berakhir.
   • --privileged mengekspos perangkat TPU ke penampung.
   • --net=host mengikat semua port penampung 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. Membuat VM TPU

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

2. SSH ke VM TPU

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

3. Memulai daemon Docker di VM TPU

   sudo systemctl start docker

4. Memulai container Docker

   sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash

5. Menginstal JAX

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

6. Menginstal FLAX

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

7. Menjalankan 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. Menghapus VM TPU

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

Pod TPU

Saat menjalankan kode JAX di Pod TPU, Anda harus menjalankan kode JAX di semua pekerja TPU secara bersamaan. Salah satu cara 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.8
   RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   WORKDIR ./flax/examples/mnist

2. Membangun gambar Docker

   docker build -t your-image-name .

3. Tambahkan tag ke image Docker sebelum mengirimnya ke Artifact Registry. Untuk mengetahui informasi selengkapnya tentang cara menggunakan Artifact Registry, lihat Menggunakan image container.

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

4. Mengirim image Docker ke Artifact Registry

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

5. Membuat VM TPU

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

6. Ambil 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"

7. Jalankan penampung 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:your-tag bash"

8. 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. Nonaktifkan penampung 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 menggunakan perintah berikut:

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

Melatih model JAX dalam penampung Docker menggunakan JAX Stable Stack

Penting: JAX Stable Stack masih dalam fase eksperimental.

Anda dapat mem-build image Docker MaxText dan MaxDiffusion menggunakan image dasar JAX Stable Stack.

JAX Stable Stack menyediakan lingkungan yang konsisten untuk MaxText dan MaxDiffusion dengan menggabungkan JAX dengan paket inti seperti orbax, flax, dan optax, serta libtpu.so. Library ini diuji untuk memastikan kompatibilitas dan memberikan fondasi yang stabil untuk mem-build dan menjalankan MaxText dan MaxDiffusion. Hal ini menghilangkan potensi konflik karena versi paket yang tidak kompatibel.

JAX Stable Stack menyertakan libtpu.so yang dirilis sepenuhnya dan memenuhi syarat, library inti yang mendorong kompilasi, eksekusi, dan konfigurasi jaringan ICI program TPU. Rilis libtpu menggantikan build harian yang sebelumnya digunakan oleh JAX, dan memastikan fungsi komputasi XLA yang konsisten di TPU dengan pengujian kualifikasi tingkat PJRT di IR HLO/StableHLO.

Untuk mem-build image Docker MaxText dan MaxDiffusion dengan JAX Stable Stack, saat Anda menjalankan skrip docker_build_dependency_image.sh, tetapkan variabel MODE ke stable_stack dan tetapkan variabel BASEIMAGE ke image dasar yang ingin Anda gunakan.

Contoh berikut menentukan us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1 sebagai gambar dasar:

   bash docker_build_dependency_image.sh MODE=stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1
   

Untuk daftar image dasar JAX Stable Stack yang tersedia, lihat image JAX Stable Stack di Artifact Registry.

Langkah selanjutnya

• Tutorial Cloud TPU
• Mengelola TPU
• Arsitektur Sistem Cloud TPU
• Menjalankan kode TensorFlow pada slice Pod TPU
• Menjalankan kode JAX pada slice Pod TPU