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.
  • Menyediakan cluster GKE menggunakan XPK dengan Pathways diaktifkan.
  • Mengonversi model Qwen3 14b dari format Hugging Face ke format MaxText.
  • Menjalankan workload pelatihan SFT multi-host di Cloud TPU.
  • Mengonversi model yang di-fine-tune kembali ke format Hugging Face untuk penyajian.

Biaya

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

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 akun gratis di Hugging Face. Setelah 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 Read , lalu 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: ID project Anda Google Cloud
  • 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: pemesanan 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}"

Membuat cluster Google Kubernetes Engine

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 workload berbasis CPU. Jangan jalankan tugas ini di beberapa komputer secara paralel. Perintah berikut mencakup pemeriksaan untuk memastikan 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.

Memulai workload pelatihan

Setelah proses konversi selesai, Anda dapat memulai workload fine-tuning 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 workload Anda menggunakan antarmuka command line (CLI) XPK.

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

Untuk melihat log dan penggunaan TPU, gunakan Google Cloud konsol. Anda juga dapat melihat log dengan menjalankan perintah berikut:

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

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

Mengonversi model terlatih kembali ke format Hugging Face

Setelah workload pelatihan selesai, konversi checkpoint kembali 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 di-tune 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.

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.