Halaman ini diterjemahkan oleh Cloud Translation API.

Melatih model menggunakan TPU7x (Ironwood)

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. Gunakan Google Cloud CLI untuk membuat instance cluster GKE secara manual guna melakukan penyesuaian atau perluasan yang tepat pada 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 proof-of-concept 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 GitHub XPK.

Instal XPK dan dependensi

Instal XPK. Ikuti petunjuk di repositori GitHub XPK.
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
```
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 RESERVATION_NAME=YOUR_RESERVATION_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. Untuk Pratinjau, hanya us-central1-c yang didukung.
- YOUR_CLUSTER_NAME: Nama cluster baru.
- YOUR_ACCELERATOR_TYPE: Versi dan topologi TPU. Contoh, tpu7x-4x4x8. Untuk mengetahui daftar topologi yang didukung, lihat Konfigurasi yang didukung.
- YOUR_RESERVATION_NAME: Nama pemesanan Anda. Untuk pemesanan bersama, gunakan projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME.
- YOUR_BUCKET_NAME: Nama bucket Cloud Storage Anda, yang akan menjadi direktori output untuk pelatihan model.

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

Ikuti petunjuk di bagian Mengonfigurasi MTU untuk mengoptimalkan konfigurasi jaringan Anda.
Isi variabel ${CLUSTER_ARGUMENTS}, yang akan Anda gunakan dalam perintah xpk cluster create:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"
```
Buat cluster GKE dengan node pool TPU7x menggunakan perintah xpk cluster create:
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
Menyetel tanda --cluster-cpu-machine-type ke n1-standard-8 (atau lebih besar) memastikan bahwa kumpulan node default memiliki CPU yang memadai 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.

Tambahkan pengecualian pemeliharaan untuk mencegah upgrade cluster:

gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"

Ganti kode berikut:

EXCLUSION_START_TIME: Waktu mulai yang Anda pilih untuk pengecualian pemeliharaan dalam format YYYY-MM-DDTHH:MM:SSZ.
EXCLUSION_END_TIME: Waktu berakhir yang Anda pilih untuk pengecualian pemeliharaan dalam format YYYY-MM-DDTHH:MM:SSZ.

Membangun atau mengupload image Docker MaxText

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

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

Setelah perintah berhasil dieksekusi, 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.