Menjalankan SFT untuk Gemma 4 31B di TPU v6e multi-host

Tutorial ini menunjukkan cara menjalankan supervised fine-tuning (SFT) pada cluster Tensor Processing Unit (TPU) v6e multi-host menggunakan MaxText dan Cluster Toolkit. Anda menggunakan Cluster Toolkit untuk menjalankan workload pelatihan multi-host dan mengekspor hasilnya kembali ke format Hugging Face untuk inferensi.

Ringkasan workload

Menjalankan workload SFT ini di Cloud TPU dan GKE melibatkan komponen berikut:

  • Model: Gemma 4 31B (google/gemma-4-31b-it), model open-weights yang canggih. Dalam tutorial ini, bobot model didownload dari Hugging Face, dikonversi ke format yang kompatibel dengan MaxText, dan di-deploy di cluster GKE dengan infrastruktur Cloud TPU v6e (v6e-32).
  • Metode pelatihan: SFT menggunakan MaxText. Workload ini menjalankan pipeline fine-tuning yang skalabel di Cloud TPU v6e menggunakan JAX. Model ini mengoptimalkan parameter model selama 1.000 langkah menggunakan pelatih SFT bawaan MaxText (maxtext.trainers.post_train.sft.train_sft) dan set data penyesuaian instruksi.
  • Set data dan template: UltraChat 200k (HuggingFaceH4/ultrachat_200k), set data dialog multi-turn berskala besar dan berkualitas tinggi. Workload menggunakan set data ini untuk menyesuaikan model dan meningkatkan kemampuan percakapan serta mengikuti petunjuk.
  • Infrastruktur dan framework: GKE, yang dikelola menggunakan Cluster Toolkit, di Cloud TPU v6e (v6e-32, dengan topologi 4x8). Framework-nya adalah MaxText (tpu_post_training:0.2.4) dengan backend JAX.

Tujuan

  • Instal Cluster Toolkit dan dependensinya.
  • Men-deploy cluster Cluster Toolkit.
  • Mengonversi model Hugging Face ke format MaxText.
  • Menjalankan workload pelatihan SFT di TPU.
  • Mengonversi kembali model yang telah disesuaikan ke format Hugging Face untuk inferensi.

Biaya

Dalam dokumen ini, Anda akan menggunakan komponen Google Cloudyang dapat ditagih berikut:

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga.

Pengguna Google Cloud baru mungkin memenuhi syarat untuk mendapatkan uji coba gratis.

Setelah menyelesaikan tugas yang dijelaskan dalam dokumen ini, Anda dapat menghindari penagihan berkelanjutan dengan menghapus resource yang Anda buat. Untuk mengetahui informasi selengkapnya, lihat Pembersihan.

Sebelum memulai

  • Anda memerlukan token akses Hugging Face untuk menggunakan tutorial ini. Anda dapat mendaftar untuk mendapatkan akun gratis di Hugging Face. Setelah Anda memiliki akun, buat token akses:

    1. Di halaman Welcome to Hugging Face, klik avatar akun Anda, lalu pilih Token akses.
    2. Di halaman Token akses, klik Create new token.
    3. Pilih jenis token Read dan masukkan nama untuk token Anda.
    4. Token akses Anda akan ditampilkan. Simpan token di tempat yang aman.

Untuk mendapatkan izin yang Anda perlukan untuk menyelesaikan tutorial ini, minta administrator untuk memberi Anda peran IAM berikut di project Anda:

Untuk mengetahui informasi selengkapnya tentang pemberian peran, lihat Mengelola akses ke project, folder, dan organisasi.

Anda mungkin juga bisa mendapatkan izin yang diperlukan melalui peran khusus atau peran bawaan lainnya.

Menyiapkan variabel lingkungan Anda

Siapkan variabel lingkungan Anda:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_post_training:0.2.4"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

Ganti kode berikut:

  • YOUR_PROJECT_ID: ID project Google Cloud Anda.
  • YOUR_REGION: region tempat pemesanan TPU Anda berada dan tempat Anda men-deploy cluster.
  • YOUR_ZONE: zona tempat pemesanan TPU Anda berada dan tempat Anda men-deploy cluster.
  • YOUR_CLUSTER_NAME: nama cluster GKE Anda. Nama harus diawali dengan huruf kecil, hanya berisi huruf kecil, angka, dan tanda hubung, serta terdiri dari 20 karakter atau kurang, sehingga ID akun layanan yang didasarkan pada hal ini tetap berada dalam batas 30 karakter.
  • YOUR_BUCKET_NAME: nama unik global untuk bucket Cloud Storage.
  • YOUR_RESERVATION_NAME: nama pemesanan TPU v6e Anda.
  • YOUR_HF_TOKEN: token akses Hugging Face Anda.

Secara default, skrip menetapkan MODEL_NAME ke gemma4-31b (nama model MaxText untuk google/gemma-4-31b-it).

Menginstal dependensi Cluster Toolkit

Untuk menyelesaikan tutorial ini dari klien atau workstation Linux atau macOS, ikuti langkah-langkah yang relevan di bagian Menginstal dependensi dalam dokumentasi Cluster Toolkit.

Jika Anda menggunakan Cloud Shell, Anda dapat melewati bagian ini.

Menginstal Cluster Toolkit

Download dan ekstrak paket bawaan untuk versi Cluster Toolkit tertentu yang ditentukan dalam CLUSTER_TOOLKIT_VERSION (v1.103.0) ke direktori kerja Anda saat ini:

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64)              BUNDLE="gcluster_bundle_linux_amd64.tgz" ;;
  Linux-aarch64|Linux-arm64) BUNDLE="gcluster_bundle_linux_arm64.tgz" ;;
  Darwin-x86_64)             BUNDLE="gcluster_bundle_mac_amd64.tgz" ;;
  Darwin-arm64)              BUNDLE="gcluster_bundle_mac_arm64.tgz" ;;
  *) BUNDLE=""; echo "Unsupported platform: $(uname -s)-$(uname -m)" >&2 ;;
esac
mkdir -p cluster-toolkit
curl -fsSL "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:?}/${BUNDLE:?}" \
  | tar -xz -C cluster-toolkit
cd cluster-toolkit

Mengautentikasi lingkungan Anda

Untuk memastikan Cluster Toolkit dapat menyediakan resource dan mengirim image ke Artifact Registry atas nama Anda, autentikasi Google Cloud CLI dan tetapkan Kredensial Default Aplikasi Anda:

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
gcloud auth application-default set-quota-project "${PROJECT:?}"

gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION:?}-docker.pkg.dev" --quiet

Membuat cluster Cluster Toolkit Anda

Untuk membuat dan men-deploy cluster Cluster Toolkit dengan 32 chip TPU v6e, selesaikan langkah-langkah berikut:

  1. Membuat bucket Cloud Storage: 

    gcloud storage buckets create "gs://${GCS_BUCKET:?}" --project="${PROJECT:?}" --location="${REGION:?}" || true

  2. Tempatkan konfigurasi blueprint dan salin contoh blueprint deployment ke direktori kerja yang terisolasi:

    mkdir -p tmp
cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

  3. Perbarui tmp/gke-tpu-v6e-advanced.yaml untuk memberi akun layanan node pool cluster Anda (gke-np-sa) peran Storage Admin (roles/storage.admin) untuk menulis checkpoint ke bucket Cloud Storage Anda, alihkan jenis mesin node pool sistem ke e2-standard-8, dan sematkan node pool sistem ke zona target Anda untuk menghindari kehabisan stok jenis mesin regional:

    # Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
# Pin system node pool to a single zone to avoid cross-zone stockouts
sed -i '/system_node_pool_machine_type/a \      system_node_pool_zones: [$(vars.zone)]' tmp/gke-tpu-v6e-advanced.yaml

  4. Gunakan perintah gcluster deploy untuk men-deploy cluster Cluster Toolkit menggunakan blueprint tmp/gke-tpu-v6e-advanced.yaml dan meneruskan variabel yang diperlukan menggunakan flag --vars:

    ./gcluster deploy tmp/gke-tpu-v6e-advanced.yaml \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -l IGNORE \
    --auto-approve -w

Mengonversi model ke format MaxText

Untuk melatih model dalam format MaxText, Anda harus mengonversinya dari format Hugging Face ke format MaxText.

  1. Untuk menyederhanakan perintah berikutnya, gunakan perintah gcluster job config untuk mengonfigurasi project, cluster, dan lokasi default Anda:

    # Configure gcluster Defaults
./gcluster job config set project "${PROJECT:?}"
./gcluster job config set cluster "${CLUSTER_NAME:?}"
./gcluster job config set location "${REGION:?}"

  2. Gunakan perintah gcluster job submit untuk mengonversi model dari format Hugging Face ke format MaxText dan menyimpannya di bucket Cloud Storage Anda:

    ./gcluster job submit --name hf-to-mt \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_maxtext \
        model_name=${MODEL_NAME:?} \
        hf_access_token=${HF_TOKEN:?} \
        base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME}/max-text-format/ \
        scan_layers=True \
        use_multimodal=False \
        skip_jax_distributed_system=true \
        checkpoint_storage_use_zarr3=0 \
        checkpoint_storage_use_ocdbt=0 \
        hardware=cpu \
        --lazy_load_tensors=True"

    Perintah ini menyertakan flag --await-job-completion, yang berarti perintah tersebut menunggu di terminal Anda saat ini hingga konversi selesai.

  3. Opsional: Untuk melakukan streaming log konversi saat tugas berjalan, buka terminal terpisah di direktori kerja tempat Anda mengekstrak gcluster, dan jalankan perintah gcluster job logs:

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs hf-to-mt --main-only -f

  4. Di terminal asli, pastikan file model yang dikonversi tersedia di bucket Cloud Storage Anda:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"

Memulai workload pelatihan

  1. Setelah proses konversi selesai, mulai workload SFT:

    ./gcluster job submit --name sft \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
      python3 -m maxtext.trainers.post_train.sft.train_sft \
      run_name=sft \
      base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/ \
      model_name=${MODEL_NAME} \
      load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
      hf_access_token=${HF_TOKEN:?} \
      dataset_type=hf \
      hf_path=HuggingFaceH4/ultrachat_200k \
      per_device_batch_size=1 steps=1000 \
      profiler=xplane \
      checkpoint_storage_use_zarr3=0 \
      checkpoint_storage_use_ocdbt=0 \
      skip_jax_distributed_system=False"

    Perintah ini menyertakan flag --await-job-completion, yang berarti perintah tersebut menunggu di terminal Anda saat ini hingga tugas pelatihan selesai.

  2. Opsional: Untuk melakukan streaming log pelatihan langsung saat tugas berjalan, buka terminal terpisah di direktori kerja tempat Anda mengekstrak gcluster, dan jalankan perintah gcluster job logs:

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs sft --main-only -f

    Setelah kompilasi XLA awal selesai, koordinator akan mencatat progres output untuk setiap langkah di seluruh operasi 1.000 langkah dalam format completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ..., dengan penurunan loss pelatihan dari sekitar 1.57 pada langkah 5 menjadi sekitar 0.82 pada langkah 1000. Sebelum melanjutkan ke bagian berikutnya, tunggu hingga langkah 1000 selesai dan output dari perintah ./gcluster job list menunjukkan status sft sebagai Succeeded.

  3. Di terminal asli Anda, pastikan checkpoint pelatihan SFT dibuat di bucket Cloud Storage Anda:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

Mengonversi kembali model terlatih ke format Hugging Face

  1. Setelah workload pelatihan selesai, konversikan kembali model ke format Hugging Face:

    ./gcluster job submit --name mt-to-hf \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_huggingface \
        model_name=${MODEL_NAME:?} \
        hf_access_token=${HF_TOKEN:?} \
        load_parameters_path=gs://${GCS_BUCKET:?}/${MODEL_NAME}/trained/sft/checkpoints/1000/model_params/ \
        base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
        skip_jax_distributed_system=true \
        hardware=cpu \
        scan_layers=True \
        use_multimodal=False \
        weight_dtype=bfloat16 \
        --override_model_architecture"

    Perintah ini menyertakan flag --await-job-completion, yang berarti perintah tersebut menunggu di terminal Anda saat ini hingga konversi selesai.

  2. Opsional: Untuk melakukan streaming log konversi langsung saat tugas berjalan, buka terminal terpisah di direktori kerja tempat Anda mengekstrak gcluster, dan jalankan perintah gcluster job logs:

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs mt-to-hf --main-only -f

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  3. Di terminal asli Anda, pastikan bobot Hugging Face Safetensors akhir dibuat di bucket Cloud Storage Anda:

    gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"

Pembersihan

Agar tidak menimbulkan biaya tambahan, gunakan perintah gcluster destroy untuk menghapus resource yang dibuat selama tutorial ini:

./gcluster destroy "${CLUSTER_NAME:?}" --robust --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

# To delete the local deployment folder and staged blueprint
rm -rf .ghpc "${CLUSTER_NAME:?}" tmp

Langkah berikutnya

  • Untuk mengetahui informasi selengkapnya tentang Cloud TPU, lihat Pengantar Cloud TPU.
  • Untuk mengetahui detail arsitektur dan konfigurasi TPU v6e-32, lihat TPU v6e.