Menjalankan pelatihan reinforcement learning multi-host untuk Gemma 4 26B di TPU v6e

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

Ringkasan workload

Menjalankan workload pasca-pelatihan RL ini di Cloud TPU dan GKE melibatkan komponen berikut:

  • Model: Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it), model instruction-tuned 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.
  • Algoritma Reinforcement Learning (RL): Group Relative Policy Optimization (GRPO) menggunakan MaxText dan Pathways. Workload menjalankan pipeline pasca-pelatihan yang skalabel menggunakan pelatih RL bawaan MaxText (maxtext.trainers.post_train.rl.train_rl) yang dipasangkan dengan backend sampler vLLM. GRPO mengoptimalkan kebijakan dengan menghasilkan sekelompok output per perintah dan mengevaluasi reward relatif tanpa memerlukan model kritik terpisah.
  • Set data dan template: OpenMathInstruct-2 (openmathinstruct2_rl.json), set data penyesuaian instruksi matematika berskala besar. Workload ini menggunakan template chat terstruktur selama fase peluncuran RL untuk meningkatkan kemampuan penalaran matematika multi-langkah.
  • Infrastruktur dan framework: GKE, yang dikelola menggunakan Cluster Toolkit, di Cloud TPU v6e (slice ct6e-standard-4t, v6e-32 dengan topologi 4x8). Framework-nya adalah MaxText (tpu_post_training:0.2.4) dengan Pathways dan backend JAX.

Tujuan

  • Instal Cluster Toolkit dan dependensinya.
  • Men-deploy cluster Cluster Toolkit.
  • Mengonversi model Hugging Face ke format MaxText.
  • Menjalankan workload pelatihan RL di cluster TPU v6e.
  • Mengonversi model yang telah disesuaikan kembali 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 COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-32"
export TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

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.
  • 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.

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

Sebelum menjalankan skrip deployment, Anda harus mengautentikasi Google Cloud CLI dan menetapkan Kredensial Default Aplikasi (ADC). Hal ini memastikan bahwa Cluster Toolkit dapat menyediakan resource dan mengirim image ke Artifact Registry atas nama Anda.

gcloud auth login
gcloud auth application-default login --no-launch-browser
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. Gunakan perintah gcluster deploy untuk men-deploy cluster Cluster Toolkit menggunakan blueprint gke-tpu-v6e-advanced.yaml dan meneruskan variabel yang diperlukan menggunakan flag --vars:

    TMP_DIR=$(mktemp -d)
BLUEPRINT_FILE="${TMP_DIR}/gke-tpu-v6e-advanced.yaml"

cp ./examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml "${BLUEPRINT_FILE}"

# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" "${BLUEPRINT_FILE}"
# 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)]' "${BLUEPRINT_FILE}"
# Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" "${BLUEPRINT_FILE}"

./gcluster deploy "${BLUEPRINT_FILE}" \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=${CLUSTER_NODEPOOL_COUNT:?},tpu_topology=${TOPOLOGY:?},authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    -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="gemma4-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='google/gemma-4-26b-a4b-it' \
    base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME}/max-text-format/ \
    scan_layers=False \
    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 gemma4-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 pelatihan RL:

    ./gcluster job submit \
  --name="gemma4-training" \
  --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="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    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} \
    scan_layers=False \
    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 \
    batch_size=4 \
    train_micro_batch_size=2 \
    max_target_length=640 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=4 \
    rollout_expert_parallelism=1 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='google/gemma-4-26b-a4b-it' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.7 \
    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'}}\""

  2. Periksa status, tunggu hingga Pod koordinator siap, dan streaming log tugas pelatihan:

    # Use the list command to check status (initially shows "Suspended" briefly while the CPU node pool scales up)
./gcluster job list

# Ensure kubectl is configured
gcloud container clusters get-credentials "${CLUSTER_NAME:?}" \
  --location="${REGION:?}" \
  --project="${PROJECT:?}"

# Wait for the Pathways head pod to be ready before streaming logs
kubectl wait --for=condition=Ready pod \
    -l job-name=gemma4-training-pathways-head-0 \
    --timeout=600s

# Stream training logs (runs for 50 steps; wait for "Post RL Training: ... mean_reward=..." and job status Succeeded before converting to Hugging Face)
kubectl logs -f \
    -l job-name=gemma4-training-pathways-head-0 \
    -c workload-container

    Segera setelah pengiriman, output dari perintah ./gcluster job list mungkin akan menampilkan Suspended secara singkat saat GKE meningkatkan skala node pool CPU koordinator Pathways (cpu-np). Saat tugas dimulai, tugas tersebut akan menghasilkan evaluasi pra-pelatihan awal pada langkah 0 (Pre RL Training: ... mean_reward=...), menjeda output log selama beberapa menit selama kompilasi XLA awal, lalu melatih selama 50 langkah, mencetak Reshard finished dan pembuatan peluncuran di setiap langkah tanpa nomor langkah.

    Sebelum melanjutkan ke langkah berikutnya, tunggu hingga log menghasilkan output Post RL Training: ... mean_reward=... dan output dari perintah ./gcluster job list menunjukkan status gemma4-training sebagai Succeeded.

  3. Pastikan 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

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

    ./gcluster job submit \
  --name="gemma4-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=False \
      use_multimodal=False \
      weight_dtype=bfloat16 \
      --override_model_architecture"

  2. Opsional: Untuk memantau log konversi langsung saat perintah --await-job-completion sebelumnya menunggu di terminal pertama Anda, buka terminal terpisah di direktori kerja tempat Anda mengekstrak gcluster, dan gunakan 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 gemma4-mt-to-hf --main-only -f

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

  3. Di terminal asli, pastikan bobot model Hugging Face dan file konfigurasi yang telah dilatih 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:?}" --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

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

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.