Menjalankan workload TPU di container Docker

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.

Melatih model PyTorch di container Docker

PyTorch di Cloud TPU menjelaskan library PyTorch untuk TPU dan menyediakan link ke dokumentasi untuk mendapatkan paket Python dan image Docker. Anda dapat menggunakan TorchTPU secara langsung (lihat Mulai menggunakan GKE untuk image container rilis publik) atau menggunakan library referensi TorchTitan untuk membangun dan menjalankan workload pelatihan dalam container.

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 Google Cloud Anda. Gunakan project yang sudah ada atau buat project baru.
  • TPU_NAME: Nama TPU.
  • ZONE: Zona tempat membuat VM TPU. 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, bersihkan 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 pada slice TPU, Anda harus menjalankan kode JAX di semua TPU worker 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 worker TPU.

  1. Buat file bernama Dockerfile di direktori Anda 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. 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 gambar Docker

    docker build -t your-image-name .

  4. Tambahkan tag ke image Docker Anda sebelum mengirimnya 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 TPU worker

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

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

    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, bersihkan resource.

  1. Menonaktifkan container di semua worker

    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