Melatih model menggunakan TPU7x (Ironwood)

Pratinjau

Produk atau fitur ini tunduk pada "Persyaratan Penawaran Pra-GA" di bagian Persyaratan Layanan Umum dari Persyaratan Khusus Layanan. Produk dan fitur pra-GA tersedia "sebagaimana adanya" dan mungkin memiliki dukungan terbatas. Untuk mengetahui informasi selengkapnya, lihat deskripsi tahap peluncuran.

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

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.

Menyediakan TPU

Anda dapat menyediakan dan mengelola TPU7x menggunakan metode berikut:

• GKE: Anda dapat menggunakan GKE untuk menyediakan dan mengelola TPU sebagai kumpulan akselerator untuk workload machine learning dalam container Anda. Gunakan Google Cloud CLI untuk membuat instance cluster GKE secara manual guna melakukan penyesuaian yang tepat atau memperluas lingkungan GKE produksi yang ada. Untuk mengetahui informasi selengkapnya, lihat Tentang TPU di GKE.
• GKE dan XPK: XPK adalah alat command line yang menyederhanakan pembuatan cluster dan eksekusi workload di GKE. Dirancang bagi praktisi ML untuk menyediakan TPU dan menjalankan tugas pelatihan tanpa memerlukan keahlian Kubernetes yang mendalam. Gunakan XPK untuk membuat cluster GKE dengan cepat dan menjalankan beban kerja untuk pembuktian konsep dan pengujian. Untuk mengetahui informasi selengkapnya, lihat repositori GitHub XPK.
• GKE dan TPU Cluster Director: TPU Cluster Director tersedia melalui reservasi mode Semua Kapasitas, yang memberi Anda akses penuh ke semua kapasitas yang dicadangkan (tanpa penundaan) dan visibilitas penuh ke topologi hardware TPU, status pemanfaatan, dan status kondisi. Untuk mengetahui informasi selengkapnya, lihat Ringkasan mode Semua Kapasitas.

Men-deploy workload pelatihan dengan MaxText dan XPK

Gunakan Accelerated Processing Kit (XPK) untuk membuat cluster GKE untuk proof-of-concept dan pengujian. XPK adalah alat command line yang dirancang untuk menyederhanakan penyediaan, pengelolaan, dan menjalankan workload machine learning.

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

Sebelum memulai

Sebelum memulai, selesaikan langkah-langkah berikut:

• Pastikan Anda memiliki Google Cloud project dengan penagihan diaktifkan.
• Mendapatkan akses ke TPU7x. Untuk informasi selengkapnya, hubungi tim akun Anda.
• Pastikan akun yang Anda gunakan dengan XPK memiliki peran yang tercantum di repositori XPK GitHub.

Instal XPK dan dependensi

1. Instal XPK. Ikuti petunjuk di repositori GitHub XPK.

2. Instal Docker menggunakan petunjuk yang diberikan oleh administrator Anda atau ikuti petunjuk penginstalan resmi. Setelah diinstal, jalankan perintah berikut untuk mengonfigurasi Docker dan menguji penginstalan:

   gcloud auth configure-docker
   sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
   docker run hello-world # Test Docker
   

3. Tetapkan variabel lingkungan berikut:

   export PROJECT_ID=YOUR_PROJECT_ID
   export ZONE=YOUR_ZONE
   export CLUSTER_NAME=YOUR_CLUSTER_NAME
   export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
   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. Untuk Pratinjau, hanya us-central1-c yang didukung.
   • YOUR_CLUSTER_NAME: Nama cluster baru.
   • YOUR_ACCELERATOR_TYPE: Versi dan topologi TPU. Misalnya, tpu7x-4x4x8. Untuk mengetahui daftar topologi yang didukung, lihat Konfigurasi yang didukung.
   • YOUR_BUCKET_NAME: Nama bucket Cloud Storage Anda, yang akan menjadi direktori output untuk pelatihan model.

4. 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
   

Membuat cluster slice tunggal NIC tunggal

Pilih salah satu opsi berikut untuk membuat cluster Anda. Sebaiknya gunakan jaringan kustom dengan MTU 8.896 untuk performa yang optimal.

Jaringan kustom

Untuk membuat jaringan kustom dengan MTU 8.896 dan menggunakannya untuk cluster Anda, ikuti langkah-langkah berikut:

1. Tetapkan variabel lingkungan untuk nama jaringan dan firewall:

   export NETWORK_NAME=NETWORK_NAME
   export NETWORK_FW_NAME=FIREWALL_NAME

   Ganti kode berikut:

   • NETWORK_NAME: Nama untuk jaringan.
   • FIREWALL_NAME: Nama untuk aturan firewall jaringan.

2. Buat jaringan kustom dengan MTU 8.896:

   gcloud compute networks create ${NETWORK_NAME} \
       --mtu=8896 \
       --project=${PROJECT_ID} \
       --subnet-mode=auto \
       --bgp-routing-mode=regional

3. Buat aturan firewall yang mengizinkan traffic TCP, ICMP, dan UDP di jaringan Anda:

   gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
       --network=${NETWORK_NAME} \
       --allow tcp,icmp,udp \
       --project=${PROJECT_ID}

4. Tetapkan variabel lingkungan untuk argumen cluster XPK guna menggunakan jaringan yang Anda buat:

   export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

5. Buat cluster XPK. Perintah berikut menyediakan kapasitas sesuai permintaan:

   xpk cluster create --cluster=${CLUSTER_NAME} \
       --cluster-cpu-machine-type=n1-standard-8 \
       --num-slices=${NUM_SLICES} \
       --tpu-type=${ACCELERATOR_TYPE} \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --on-demand \
       --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

   Untuk menggunakan kapasitas yang dicadangkan, ganti --on-demand dengan --reservation=RESERVATION_NAME. Untuk menggunakan VM Spot TPU, ganti --on-demand dengan --spot.

Jaringan default

Jika tidak memerlukan jaringan MTU tinggi, Anda dapat membuat cluster yang menggunakan jaringan VPC default. Perintah berikut menyediakan kapasitas sesuai permintaan:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

Untuk menggunakan kapasitas yang dicadangkan, ganti --on-demand dengan --reservation=RESERVATION_NAME. Untuk menggunakan VM Spot TPU, ganti --on-demand dengan --spot.

Catatan: Menetapkan flag --cluster-cpu-machine-type ke n1-standard-8 (atau lebih besar) memastikan bahwa kumpulan node default memiliki CPU yang cukup untuk pod sistem, misalnya webhook JobSet, sehingga mencegah error. Secara default, XPK menggunakan e2-standard-16. Beberapa zona hanya mendukung jenis CPU tertentu, sehingga Anda mungkin perlu beralih antara jenis n1, n2, dan e2. Jika tidak, Anda mungkin akan mengalami error kuota.

Membangun atau mengupload image Docker MaxText

Catatan: MaxText hanya mendukung Python 3.12. Bangun lingkungan virtual Anda dengan 3.12 untuk menginstal dependensi yang benar.

Anda dapat membuat image Docker secara lokal menggunakan skrip yang disediakan oleh MaxText atau menggunakan image bawaan.

Membangun secara lokal

Perintah berikut menyalin direktori lokal Anda ke dalam container:

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# 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

Setelah berhasil menjalankan perintah, Anda akan melihat image bernama maxtext_base_image yang dibuat secara lokal. Anda dapat menggunakan image lokal secara langsung di perintah workload xpk.

Upload gambar (opsional)

Setelah membangun image Docker secara lokal menggunakan petunjuk di bagian sebelumnya, Anda dapat mengupload image Docker MaxText ke dalam registry menggunakan perintah berikut:

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

Setelah berhasil menjalankan perintah ini, Anda akan melihat image MaxText di gcr.io dengan nama gcr.io/PROJECT_ID/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_xpk \
        steps=30"

Men-deploy workload pelatihan

Jalankan perintah xpk workload create untuk men-deploy tugas pelatihan Anda. Anda harus menentukan flag --base-docker-image untuk menggunakan image dasar MaxText atau menentukan flag --docker-image dan image yang ingin Anda gunakan. Anda dapat memilih untuk menyertakan tanda --enable-debug-logs guna mengaktifkan logging debug.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

Nama beban kerja harus unik dalam cluster. Dalam contoh ini, $(date +%H%M) ditambahkan ke nama beban kerja untuk memastikan keunikan.

Langkah berikutnya

• Gunakan platform Diagnostik ML Google Cloud untuk mengoptimalkan dan mendiagnosis workload Anda
• Menjalankan workload pelatihan menggunakan resep yang dioptimalkan untuk TPU7x
• Menjalankan microbenchmark TPU7x