Melatih model menggunakan TPU7x (Ironwood)

Dokumen ini menjelaskan cara menyediakan resource TPU7x dan memberikan contoh men-deploy workload pelatihan menggunakan MaxText dan Cluster Toolkit.

TPU7x adalah rilis pertama dalam keluarga 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 Resep Pelatihan untuk TPU Ironwood di GitHub.

Men-deploy workload pelatihan dengan MaxText dan Cluster Toolkit

Gunakan Cluster Toolkit untuk membuat cluster GKE yang siap produksi dan menjalankan workload pelatihan.

Bagian berikut menunjukkan cara men-deploy workload pelatihan menggunakan MaxText dan Cluster Toolkit.

Sebelum memulai

Sebelum memulai, selesaikan langkah-langkah berikut:

  1. Tetapkan variabel lingkungan berikut:

    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 cluster akan dibuat.
    • YOUR_CLUSTER_NAME: Nama cluster baru.
    • YOUR_BUCKET_NAME: Nama bucket Cloud Storage Anda, yang akan menjadi direktori output untuk pelatihan model.

  2. Jika Anda belum memiliki bucket Cloud Storage, buat bucket menggunakan perintah berikut:

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

Deploy cluster TPU7x (Ironwood)

Deploy cluster TPU7x GKE dengan mengikuti petunjuk di Men-deploy cluster TPU 7x GKE. Pastikan topologi node pool cluster Anda cocok dengan persyaratan workload pelatihan Anda (seperti 4x4x8 untuk contoh workload MaxText di bawah).

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}

Tentukan perintah pelatihan MaxText

Siapkan perintah untuk menjalankan skrip pelatihan Anda dalam container Docker.

Model MaxText 1B adalah konfigurasi dalam framework MaxText yang dirancang untuk melatih model bahasa dengan sekitar 1 miliar parameter. Gunakan model ini untuk bereksperimen dengan 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"

Men-deploy workload pelatihan

Kirim tugas pelatihan Anda menggunakan gcluster job submit. Untuk mengetahui detail 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 beban kerja harus unik dalam cluster. Dalam contoh ini, $(date +%H%M) ditambahkan ke nama beban kerja untuk memastikan keunikan.

Langkah berikutnya