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
Buat file bernama
Dockerfile
di direktori saat ini dan tempelkan teks berikutFROM 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
Membuat bucket Cloud Storage
gcloud storage buckets create gs://your-bucket-name --location=europe-west4
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
Menyalin Dockerfile ke VM TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
SSH ke VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Membangun gambar Docker
sudo docker build -t your-image-name .
Memulai container Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
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
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.
- Ketik
exit
untuk keluar dari container Docker - Ketik
exit
untuk keluar dari VM TPU - Menghapus VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Pod TPU
Buat file bernama
Dockerfile
di direktori saat ini dan tempelkan teks berikutFROM 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
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
Menyalin Dockerfile ke VM TPU
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
SSH ke VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Membangun gambar Docker
sudo docker build -t your-image-name .
Memulai container Docker
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
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.
- Ketik
exit
untuk keluar dari container Docker - Ketik
exit
untuk keluar dari VM TPU - Menghapus VM TPU
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Melatih model PyTorch dalam container Docker
Perangkat TPU
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
SSH ke VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
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
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"
Meng-clone repo PyTorch XLA
git clone --recursive https://github.com/pytorch/xla.git
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.
- Ketik
exit
untuk keluar dari container Docker - Ketik
exit
untuk keluar dari VM TPU - 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.
Membuat VM TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=us-central2-b \ --accelerator-type=v4-32 \ --version=tpu-ubuntu2204-base
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"
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
Membuat VM TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
SSH ke VM TPU
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Memulai daemon Docker di VM TPU
sudo systemctl start docker
Memulai container Docker
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
Menginstal JAX
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Menginstal FLAX
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
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.
- Ketik
exit
untuk keluar dari container Docker - Ketik
exit
untuk keluar dari VM TPU 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.
Buat file bernama
Dockerfile
di direktori saat ini dan tempelkan teks berikutFROM 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
Membangun gambar Docker
docker build -t your-image-name .
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
Mengirim image Docker ke Artifact Registry
docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
Membuat VM TPU
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type==v2-8 \ --version=tpu-ubuntu2204-base
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"
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"
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.
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"
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
- Tutorial Cloud TPU
- Mengelola TPU
- Arsitektur Sistem Cloud TPU
- Menjalankan kode TensorFlow pada slice Pod TPU
- Menjalankan kode JAX pada slice Pod TPU