Menjalankan supervised fine-tuning multi-host pada model Qwen3-14b menggunakan MaxText

Tutorial ini memberikan panduan langkah demi langkah untuk menjalankan supervised fine-tuning (SFT) pada model Qwen3-14b menggunakan MaxText di Cloud TPU. Anda akan mempelajari cara membuat image container khusus, menyediakan cluster Google Kubernetes Engine (GKE) dengan Pathways menggunakan Accelerated Processing Kit (XPK), dan menjalankan workload pelatihan multi-host.

Tujuan

• Pelajari cara membuat image container MaxText kustom yang dioptimalkan untuk pasca-pelatihan.
• Sediakan cluster GKE menggunakan XPK dengan Pathways diaktifkan.
• Konversi model Qwen3 14b dari format Hugging Face ke format MaxText.
• Jalankan workload pelatihan SFT multi-host di Cloud TPU.
• Konversi kembali model yang di-fine-tune ke format Hugging Face untuk inferensi.

Biaya

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine, for the temporary build virtual machine (VM)

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, baca bagian Pembersihan.

Sebelum memulai

• Pastikan akun pengguna atau akun layanan Anda memiliki peran berikut:
  • roles/compute.admin, untuk membuat VM build
  • roles/artifactregistry.admin, untuk mengelola repositori Docker
  • roles/storage.admin, untuk mengelola bucket data
  • roles/container.admin, untuk membuat dan mengelola cluster Google Kubernetes Engine
  • roles/iam.serviceAccountAdmin, untuk membuat akun layanan workload
  • roles/resourcemanager.projectIamAdmin, untuk menetapkan kebijakan Identity and Access Management (IAM)
  • roles/iam.serviceAccountUser, untuk bertindak sebagai akun layanan
• Instal dan lakukan inisialisasi Google Cloud CLI.
• Pastikan Anda telah menginstal Python 3.12 atau yang lebih baru di workstation Anda.

• 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 Access tokens.
  2. Di halaman Access tokens, klik Create new token.
  3. Pilih jenis token Baca dan masukkan nama untuk token Anda.
  4. Token akses Anda akan ditampilkan. Simpan token di tempat yang aman.

Menyiapkan lingkungan

Siapkan variabel lingkungan Anda dengan menjalankan skrip berikut:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
export HF_TOKEN="YOUR_HF_TOKEN"

Ganti kode berikut:

• YOUR_PROJECT_ID: Project ID Google Cloud Anda
• YOUR_REGION: region yang ingin Anda gunakan
• YOUR_ZONE: zona yang ingin Anda gunakan
• YOUR_CLUSTER_NAME: nama untuk cluster Google Kubernetes Engine Anda
• YOUR_GCS_BUCKET: nama unik untuk bucket Cloud Storage Anda
• YOUR_RESERVATION_NAME: reservasi kapasitas Anda
• YOUR_HF_TOKEN: token akses Hugging Face Anda

Menyiapkan image container MaxText

Untuk menyiapkan image container MaxText, termasuk menginstal dependensi yang diperlukan, selesaikan langkah-langkah berikut:

1. Membuat bucket Cloud Storage:

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

2. Buat repositori Artifact Registry:

   gcloud artifacts repositories create maxtext-images \
       --repository-format=docker \
       --location=$REGION \
       --project=$PROJECT \
       --description="Docker repository for MaxText images in $REGION" || true

3. Buat file di direktori root repositori Anda dengan nama file cloudbuild.yaml dan konten berikut:

   steps:
     - name: 'gcr.io/cloud-builders/docker'
       entrypoint: 'bash'
       args:
         - '-c'
         - |
           set -euo pipefail
   
           # 0. Install prerequisites (if needed)
           apt-get update && apt-get install -y curl || apk add curl || true
   
           # 1. Install uv
           curl -LsSf https://astral.sh/uv/install.sh | sh
           source $$HOME/.local/bin/env
   
           # 2. Setup Python environment and install MaxText runner
           uv venv --python 3.12 --seed maxtext_venv
           source maxtext_venv/bin/activate
           uv pip install maxtext[runner]==0.2.1 --resolution=lowest
   
           # 3. Build the Docker image (Cloud Build has Docker pre-configured)
           build_maxtext_docker_image WORKFLOW=post-training
   
           # 4. Tag the image properly
           docker tag maxtext_base_image ${_CLOUD_IMAGE_NAME}
   
   # Cloud Build automatically pushes images listed here
   images:
     - '${_CLOUD_IMAGE_NAME}'
   
   options:
     # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
     machineType: 'E2_HIGHCPU_32'

4. Gunakan Cloud Build untuk membangun image Docker MaxText Anda:

   gcloud builds submit . \
       --project=$PROJECT \
       --region=$REGION \
       --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

Buat cluster Google Kubernetes Engine Anda

Untuk menjalankan pelatihan SFT pada model Qwen3 14b, Anda memerlukan cluster Google Kubernetes Engine yang dilengkapi dengan chip TPU. Instal Accelerated Processing Kit (XPK) dan buat cluster GKE dengan dukungan Pathways.

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

Menyiapkan model untuk pelatihan

Konversi model dasar ke format MaxText menggunakan beban kerja berbasis CPU. Jangan jalankan tugas ini di beberapa komputer secara paralel. Perintah berikut mencakup pemeriksaan untuk memastikan bahwa konversi hanya berjalan di satu node TPU.

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

Melacak progres konversi model

Untuk melacak progres konversi, lakukan hal berikut:

1. Untuk mencantumkan pod yang dijadwalkan di cluster GKE Anda, jalankan perintah kubectl get pod.
2. Temukan pod bernama qwen-hf-to-mt-slice-job-0-0-HASH.
3. Untuk memeriksa output pod secara real time, jalankan perintah kubectl logs -f POD_NAME.

Mulai workload pelatihan

Setelah proses konversi selesai, Anda dapat memulai workload penyesuaian SFT menggunakan XPK.

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

Memantau workload pelatihan

Pantau status beban kerja Anda menggunakan antarmuka command line (CLI) XPK.

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

Untuk melihat log dan pemanfaatan TPU, gunakan konsolGoogle Cloud . Anda juga dapat melihat log dengan menjalankan perintah berikut:

kubectl logs -f qwen-training-pathways-head-0-0-HASH

Ganti HASH dengan hash numerik di nama pod Anda. Untuk memverifikasi nilai hash ini, jalankan perintah kubectl get pod dan periksa daftar pod yang ditampilkan.

Mengonversi kembali model terlatih ke format Hugging Face

Setelah beban kerja pelatihan selesai, konversi kembali checkpoint ke format Hugging Face.

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

Untuk melacak progres konversi, jalankan perintah kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH, dengan mengganti HASH dengan hash numerik dalam nama pod Anda.

Setelah konversi selesai, model yang disesuaikan yang disimpan di gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ siap digunakan.

Pembersihan

Untuk menghindari biaya tambahan, hapus resource yang dibuat selama tutorial ini, termasuk cluster Google Kubernetes Engine, bucket Cloud Storage, dan repositori Artifact Registry Anda.

Perhatian: Jika Anda berencana untuk terus menggunakan model yang Anda latih dalam tutorial ini, simpan versi yang dapat Anda akses sebelum menghapus bucket Cloud Storage yang berisi model tersebut.

Untuk menghapus resource yang Anda buat untuk tutorial ini, jalankan perintah berikut:

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

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.