Melatih model menggunakan TPU7x (Ironwood)

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

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

Untuk contoh lainnya yang dioptimalkan untuk TPU7x, lihat Resep Pelatihan untuk Ironwood TPU 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 lingkungan GKE produksi yang ada secara presisi. 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. Alat ini dirancang untuk praktisi ML agar dapat menyediakan TPU dan menjalankan tugas pelatihan tanpa memerlukan keahlian Kubernetes yang mendalam. Gunakan XPK untuk membuat cluster GKE dengan cepat dan menjalankan workload 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 Kapasitas Penuh, yang memberi Anda akses penuh ke semua kapasitas yang dicadangkan (tanpa penahanan) dan visibilitas penuh ke topologi hardware TPU, status penggunaan, dan status kesehatan. Untuk mengetahui informasi selengkapnya, lihat Ringkasan mode Kapasitas Penuh.

Men-deploy workload pelatihan dengan MaxText dan XPK

Gunakan Accelerated Processing Kit (XPK) untuk membuat cluster GKE untuk pembuktian konsep dan pengujian.

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 yang diaktifkan.
  • Dapatkan akses ke TPU7x. Untuk mengetahui informasi selengkapnya, hubungi tim akun Anda.
  • Pastikan akun yang Anda gunakan dengan XPK memiliki peran yang tercantum di repositori GitHub XPK.

Menginstal 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 Anda. Google Cloud
    • YOUR_ZONE: Zona tempat cluster akan dibuat.
    • 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 tidak 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 dan NIC tunggal

Pilih salah satu opsi berikut untuk membuat cluster. Sebaiknya gunakan jaringan kustom dengan MTU 8.896 untuk performa 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.

Mem-build atau mengupload image Docker MaxText

Anda dapat mem-build image Docker secara lokal menggunakan skrip yang disediakan oleh MaxText atau menggunakan image yang telah di-build sebelumnya.

Mem-build 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 perintah berhasil dieksekusi, Anda akan melihat image bernama maxtext_base_image yang dibuat secara lokal. Anda dapat menggunakan image lokal secara langsung dalam perintah workload xpk.

Mengupload image (opsional)

Setelah mem-build 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 perintah ini berhasil dieksekusi, Anda akan melihat image MaxText di gcr.io dengan nama gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

Menentukan 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 flag --enable-debug-logs untuk 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 workload harus unik dalam cluster. Dalam contoh ini, $(date +%H%M) ditambahkan ke nama workload untuk memastikan keunikan.

Langkah berikutnya