Jalankan pelatihan RL multi-host untuk Qwen3-30b-a3b pada TPU v6e

Tutorial ini menunjukkan cara menjalankan pelatihan reinforcement learning (RL) multi-host di cluster v6e-32 Tensor Processing Unit (TPU) menggunakan MaxText dan Cluster Toolkit. Anda menggunakan Cluster Toolkit untuk menjalankan beban kerja pelatihan multi-host dan mengekspor hasilnya kembali ke format Hugging Face untuk penayangan.

Tujuan

• Instal Cluster Toolkit dan dependensinya.
• Deploy cluster Cluster Toolkit.
• Mengonversi model Hugging Face ke format MaxText.
• Jalankan workload pelatihan RL di cluster TPU v6e.
• 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
• Compute Engine

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.

Penting: Untuk menyelesaikan tutorial ini, Anda harus memiliki akses ke TPU v6e-32. Tutorial ini mengasumsikan bahwa Anda sudah memiliki reservasi untuk TPU ini di projectGoogle Cloud Anda untuk memastikan bahwa Anda memiliki kapasitas yang tersedia saat Anda membutuhkannya. Untuk mempelajari lebih lanjut opsi Anda dalam memesan kapasitas Cloud TPU untuk tanggal dan waktu mendatang, lihat Tentang pemesanan Cloud TPU.

Sebelum memulai

• Anda memerlukan token akses Hugging Face untuk menggunakan tutorial ini. Anda dapat mendaftar untuk mendapatkan 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 dan masukkan nama untuk token Anda.
  4. Token akses Anda akan ditampilkan. Simpan token di tempat yang aman.

• Di situs Hugging Face, setujui perjanjian lisensi untuk model yang ingin Anda latih. Tutorial ini menggunakan model qwen3-30b-a3b.

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

• TPU Admin (roles/tpu.admin)
• Kubernetes Engine Admin (roles/container.admin)
• Compute Admin (roles/compute.admin)
• Storage Admin (roles/storage.admin)
• Service Account User (roles/iam.serviceAccountUser)
• Service Account Admin (roles/iam.serviceAccountAdmin)
• Project IAM Admin (roles/resourcemanager.projectIamAdmin)
• Service Usage Admin (roles/serviceusage.serviceUsageAdmin)

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

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 COMPUTE_TYPE="ct6e-standard-4t"
export TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-30b-a3b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

Ganti kode berikut:

• YOUR_PROJECT_ID: ID Google Cloud project Anda.
• YOUR_REGION: region tempat Anda ingin men-deploy cluster.
• YOUR_ZONE: zona tempat Anda ingin men-deploy cluster.
• YOUR_CLUSTER_NAME: nama cluster Google Kubernetes Engine Anda (maksimal 20 karakter).
• YOUR_BUCKET_NAME: nama unik secara global untuk bucket Cloud Storage.
• YOUR_RESERVATION_NAME: nama pemesanan Anda.
• YOUR_HF_TOKEN: token akses Hugging Face Anda.

Menginstal dependensi Cluster Toolkit

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

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

Menginstal Cluster Toolkit

Instal paket bawaan untuk Cluster Toolkit di direktori kerja Anda saat ini dengan mengikuti petunjuk di Menginstal Cluster Toolkit.

Misalnya, Anda dapat mendownload dan mengekstrak paket ke direktori kerja saat ini sebagai berikut:

wget -qO- "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:-v1.103.0}/gcluster_bundle_linux_amd64.tgz" | tar -xz

Mengekstrak paket ke direktori kerja Anda saat ini akan memberikan biner gcluster dan cetak biru examples/ yang Anda gunakan dalam langkah-langkah berikutnya.

Membuat cluster Cluster Toolkit

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. Salin blueprint Cluster Toolkit ke direktori kerja Anda saat ini:

   cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml .

3. Secara default, akun layanan node pool cluster Anda tidak memiliki izin yang diperlukan untuk menulis ke bucket Cloud Storage Anda. Agar akun layanan node pool dapat menulis ke bucket Cloud Storage Anda, Anda harus memberinya peran Storage Admin. Untuk memberikan peran ini, edit file gke-tpu-v6e-advanced.yaml dengan memperbarui modul service-account bernama node_pool_service_account:

   - id: node_pool_service_account
     source: modules/project/service-account
     settings:
       name: gke-np-sa
       project_roles:
       - logging.logWriter
       - monitoring.metricWriter
       - monitoring.viewer
       - stackdriver.resourceMetadata.writer
       - storage.admin
       - artifactregistry.reader

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

   ./gcluster deploy gke-tpu-v6e-advanced.yaml \
       --vars project_id="${PROJECT}" \
       --vars deployment_name="${CLUSTER_NAME}" \
       --vars region="${REGION}" \
       --vars zone="${ZONE}" \
       --vars num_slices="${CLUSTER_NODEPOOL_COUNT}" \
       --vars tpu_topology="${TOPOLOGY}" \
       --vars authorized_cidr="0.0.0.0/0" \
       --vars reservation="${RESERVATION:-}" \
       -l IGNORE --auto-approve -w

5. Konfigurasi autentikasi Container Registry dan berikan peran Storage Admin (roles/storage.admin) ke akun layanan GKE Anda:

   gcloud auth configure-docker gcr.io --quiet
   gcloud auth configure-docker "${REGION}-docker.pkg.dev" --quiet
   gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet
   gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet

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 qwen-hf-to-mt \
     --num-slices 1 \
     --image "${CLOUD_IMAGE_NAME}" \
     --compute-type "${COMPUTE_TYPE}" \
     --topology "${TOPOLOGY}" \
     --await-job-completion \
     --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
     python3 -m maxtext.checkpoint_conversion.to_maxtext \
     model_name=${MODEL_NAME} \
     hf_access_token=${HF_TOKEN} \
     --hf_model_path='Qwen/Qwen3-30B-A3B-Instruct-2507' \
     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"

3. Gunakan perintah gcluster job logsuntuk memeriksa status tugas konversi:

   # 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 qwen-hf-to-mt --main-only -f

4. Pastikan file model yang dikonversi tersedia di bucket Cloud Storage Anda:

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

Mulai workload pelatihan

Setelah proses konversi selesai, mulai workload pelatihan RL:

./gcluster job submit \
  --name="qwen-rl" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET}/pathways/" \
  --gke-ttl-after-finished="24h" \
  --restarts=0 \
  --env="GRPC_DNS_RESOLVER=native" \
  --env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --command="(echo 190G > /sys/fs/cgroup/memory.max || echo 190G > /sys/fs/cgroup/memory/memory.limit_in_bytes) 2>/dev/null || true && \
    export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, tunix.generate.vllm_sampler as vs; p = pathlib.Path(vs.__file__); p.write_text(p.read_text().replace('reshard_chunk_size: Optional[int] = None', 'reshard_chunk_size: Optional[int] = 4').replace('reshard_chunk_size=self.config.reshard_chunk_size', 'reshard_chunk_size=4'))\" && \
    python3 -c \"import pathlib, re; p = pathlib.Path('/deps/src/maxtext/trainers/post_train/rl/utils_rl.py'); p.write_text(re.sub('optax[.]adamw[(][^)]+[)]', 'optax.adafactor(learning_rate=learning_rate)', p.read_text()))\" && \
    JAX_PLATFORMS=proxy,cpu ENABLE_PATHWAYS_PERSISTENCE=1 \
    python3 -m maxtext.trainers.post_train.rl.train_rl \
    run_name=rl \
    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} \
    data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json \
    num_batches=50 \
    num_test_batches=0 \
    batch_size=8 \
    train_micro_batch_size=8 \
    max_target_length=512 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=1 \
    rollout_expert_parallelism=4 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='Qwen/Qwen3-30B-A3B-Instruct-2507' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.55 \
    remat_policy=full \
    async_scheduling=False \
    allow_split_physical_axes=true \
    ragged_gather_reduce_fallback=True \
    enable_dp_attention=False \
    decode_sampling_temperature=0.8 \
    decode_sampling_top_k=50 \
    decode_sampling_nucleus_p=0.95 \
    learning_rate=2e-5 \
    learning_rate_schedule_steps=100 \
    rl.num_generations=4 \
    rl.reshard_chunk_size=4 \
    debug=True \
    vllm_hf_overrides='{\"architectures\": [\"MaxTextForCausalLM\"]}' \
    vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'model_call_mode': 'inference', 'enable_dp_attention': false, 'allow_split_physical_axes': true, 'use_ragged_sort': false, 'ragged_gather_reduce_fallback': true, 'prefuse_moe_weights': true, 'weight_dtype': 'bfloat16'}}\""

• Untuk memeriksa status tugas pelatihan:

  # Use the list command to check status
  ./gcluster job list
  
  # Ensure kubectl credentials are configured
  gcloud container clusters get-credentials "${CLUSTER_NAME}" \
      --location="${REGION}" \
      --project="${PROJECT}"
  
  # Check progress of the job
  kubectl logs -f \
      -l jobset.sigs.k8s.io/replicatedjob-name=pathways-head \
      -c workload-container

  Catatan: JobSet yang telah selesai di GKE akan otomatis dibersihkan setelah satu jam secara default (diatur oleh setelan --gke-ttl-after-finished=1h di gcluster job submit). Jika workload Anda tidak lagi tercantum dalam output ./gcluster job list, pembersihan ini adalah perilaku yang diharapkan dan tidak berarti tugas gagal. Semua checkpoint dan bobot model tetap disimpan dengan aman di bucket Cloud Storage Anda.

• Untuk memverifikasi bahwa checkpoint pelatihan dibuat di bucket Cloud Storage Anda:

  gcloud storage ls "gs://${GCS_BUCKET}/${MODEL_NAME}/trained/rl/checkpoints/actor/"

Mengonversi kembali model terlatih ke format Hugging Face

Setelah beban kerja pelatihan selesai, konversi kembali model ke format Hugging Face:

./gcluster job submit \
  --name="qwen-mt-to-hf" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --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/rl/checkpoints/actor/50/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"

• Untuk memeriksa status tugas konversi:

  # 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 qwen-mt-to-hf --main-only -f
  # The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

• Untuk memverifikasi bahwa file konfigurasi dan bobot model Hugging Face terlatih ada 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}"
gcloud storage rm -r "gs://${GCS_BUCKET}"

# To delete the local deployment folder and copied blueprint
rm -rf .ghpc "${CLUSTER_NAME}" gke-tpu-v6e-advanced.yaml

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.
• Untuk mengetahui informasi selengkapnya tentang Cluster Toolkit, lihat Ringkasan Cluster Toolkit.