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.

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

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