Melatih model menggunakan TPU7x (Ironwood)

Dokumen ini menjelaskan cara menyediakan resource TPU v7x dengan Cluster Toolkit dan men-deploy workload pelatihan menggunakan MaxText (JAX) atau TorchTitan di TPU (PyTorch).

TPU7x adalah rilis pertama dalam rangkaian Ironwood,TPU generasi ketujuh Google Cloud. Generasi Ironwood dirancang untuk pelatihan dan inferensi AI berskala besar. Untuk mengetahui informasi selengkapnya, lihat TPU7x.

Untuk contoh lainnya yang dioptimalkan untuk TPU7x, lihat repositori berikut di GitHub:

Sebelum memulai

Sebelum memulai, pastikan Anda memenuhi persyaratan berikut:

Kemudian, selesaikan langkah-langkah penyiapan berikut:

  1. Tetapkan variabel lingkungan untuk project dan cluster Anda:

    export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"

    Ganti kode berikut:

    • YOUR_PROJECT_ID: Project ID Google Cloud Anda.
    • YOUR_ZONE: Zona tempat membuat cluster.
    • YOUR_CLUSTER_NAME: Nama cluster baru.
    • YOUR_BUCKET_NAME: Nama bucket Cloud Storage Anda untuk output pelatihan.

  2. Jika Anda belum memiliki bucket Cloud Storage, buatlah:

    gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

Men-deploy cluster TPU7x (Ironwood)

Deploy cluster TPU7x Google Kubernetes Engine (GKE) dengan mengikuti petunjuk di Men-deploy cluster TPU 7x GKE. Pastikan topologi node pool cluster Anda sesuai dengan persyaratan workload pelatihan, seperti 4x4x8 untuk contoh MaxText atau 2x2x4 untuk contoh TorchTitan.

Men-deploy workload pelatihan dengan MaxText

Bagian berikut menunjukkan cara membangun image container dan mengirimkan tugas pelatihan MaxText dengan Cluster Toolkit.

Membangun dan mengupload image Docker MaxText

Bangun dan kirim image Docker MaxText:

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Build the Docker image
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.8.2

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

Menentukan perintah pelatihan MaxText

Siapkan perintah untuk menjalankan skrip pelatihan Anda di dalam container Docker.

Konfigurasi model MaxText 1B melatih model bahasa dengan sekitar 1 miliar parameter. Gunakan model ini untuk menguji skala chip kecil. Performa tidak dioptimalkan.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_training \
        steps=30"

Mengirim tugas pelatihan MaxText

Kirim tugas pelatihan Anda menggunakan gcluster job submit. Untuk mengetahui detail tentang opsi pengiriman tugas, lihat Panduan Tugas Cluster Toolkit:

gcluster job submit \
    --name="maxtext-1b-$(date +%H%M)" \
    --cluster=${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --compute-type=tpu7x-standard-4t \
    --topology=4x4x8 \
    --image="gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}" \
    --command="${MAXTEXT_COMMAND}"

Nama workload harus unik dalam cluster. Contoh ini menambahkan $(date +%H%M) ke nama workload.

Men-deploy workload pelatihan dengan TorchTitan

TorchTitan menggunakan TorchTPU untuk menjalankan pelatihan PyTorch terdistribusi di Cloud TPU. Untuk mengetahui informasi selengkapnya, lihat README TorchTitan di TPU.

Langkah berikutnya