Menjalankan workload TPU di container Docker

Perhatian: Halaman ini berlaku untuk Cloud TPU API, yang tidak lagi dalam pengembangan aktif development. Cloud TPU API hanya akan menerima perbaikan bug dan update keamanan. Untuk mendapatkan fitur terbaru dan dukungan untuk versi TPU terbaru, sebaiknya gunakan Compute Engine atau Google Kubernetes Engine untuk mengelola resource TPU. Untuk mengetahui informasi selengkapnya, lihat resource TPU di Compute Engine dan Tentang TPU di GKE. Cloud TPU API tidak didukung di TPU7x dan versi TPU yang lebih baru.

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

Sebelum menjalankan perintah berikut, buat variabel lingkungan berikut:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Deskripsi variabel lingkungan

• PROJECT_ID: Project ID Anda. Google Cloud Gunakan project yang ada atau buat project baru.
• TPU_NAME: Nama TPU.
• ZONE: Zona tempat VM TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
• ACCELERATOR_TYPE: Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat Versi TPU.
• RUNTIME_VERSION: Versi software Cloud TPU.

Perangkat TPU

1. Membuat VM Cloud TPU

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Menghubungkan ke VM TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

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

4. Memulai container di VM TPU menggunakan image PyTorch/XLA versi 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 gunakan 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. Meng-clone repositori PyTorch XLA

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

7. Melatih ResNet50

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

Setelah skrip pelatihan selesai, hapus 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 $TPU_NAME --zone=$ZONE
   

Slice TPU

Saat menjalankan kode PyTorch di slice TPU, Anda harus menjalankan kode di semua pekerja TPU secara bersamaan. Salah satu cara untuk melakukannya adalah 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 $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Menambahkan pengguna saat ini ke grup Docker

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --zone=$ZONE \
   --worker=all \
   --command='sudo usermod -a -G docker $USER'
   

3. Meng-clone repositori PyTorch XLA

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="git clone --recursive https://github.com/pytorch/xla.git"
   

4. Menjalankan skrip pelatihan dalam container di semua pekerja TPU

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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 dihentikan.
   • --privileged mengekspos perangkat TPU ke container.
   • --net=host mengikat semua port container ke VM TPU untuk memungkinkan komunikasi antar-host di pod.
   • -e menetapkan variabel lingkungan.

Setelah skrip pelatihan selesai, hapus resource.

Hapus VM TPU menggunakan perintah berikut:

gcloud compute tpus tpu-vm delete $TPU_NAME \
--zone=$ZONE

Melatih model JAX di container Docker

Sebelum menjalankan perintah berikut, buat variabel lingkungan berikut:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Deskripsi variabel lingkungan

• PROJECT_ID: Project ID Anda. Google Cloud Gunakan project yang ada atau buat project baru.
• TPU_NAME: Nama TPU.
• ZONE: Zona tempat VM TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
• ACCELERATOR_TYPE: Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat Versi TPU.
• RUNTIME_VERSION: Versi software Cloud TPU.

Perangkat TPU

1. Membuat VM TPU

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Menghubungkan ke VM TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

3. Memulai 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. Menginstal paket tensorflow dan tensorflow-datasets

   pip install tensorflow
   pip install tensorflow-datasets
   

8. 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, hapus 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 $TPU_NAME --zone=$ZONE
   

Slice TPU

Saat menjalankan kode JAX di slice TPU, Anda harus menjalankan kode JAX di semua pekerja TPU secara bersamaan. Salah satu cara untuk melakukannya adalah 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 tempel 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. Menyiapkan Artifact Registry

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

3. Membangun image 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 Menggunakan image container.

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

5. Mengirim image Docker ke Artifact Registry

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

6. Membuat VM TPU

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command='sudo usermod -a -G docker ${USER}'
   

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
   

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

8. Menjalankan container di semua pekerja TPU

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

9. Menjalankan skrip pelatihan di semua pekerja TPU

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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, hapus resource.

1. Matikan container di semua pekerja

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="docker kill your-container-name"
   

2. Menghapus VM TPU

   gcloud compute tpus tpu-vm delete $TPU_NAME \
   --zone=$ZONE
   

Langkah berikutnya

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