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 klaster Tensor Processing Unit (TPU) v6e-64 dengan 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 disajikan.

Tujuan

  • Instal Cluster Toolkit dan dependensinya.
  • Instal MaxText dan dependensinya.
  • Deploy cluster Cluster Toolkit.
  • Konversikan model Wajah Berpelukan ke format MaxText.
  • Jalankan beban kerja pelatihan RL pada klaster 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:

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 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.
  • Di situs web Hugging Face, setujui perjanjian lisensi untuk model yang ingin Anda latih. Tutorial ini menggunakan model gemma4-26b.

Untuk mendapatkan izin yang Anda perlukan untuk menyelesaikan tutorial ini, minta administrator Anda 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.

Atur variabel lingkungan Anda

Atur variabel lingkungan Anda:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-64"
export TOPOLOGY="8x8"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
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 klaster Google Kubernetes Engine Anda.
  • YOUR_REPOSITORY_NAME: nama repositori Artifact Registry untuk gambar MaxText Anda.
  • 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.

Instal Cluster Toolkit

Instal paket pra-bangun untuk Cluster Toolkit dengan mengikuti petunjuk di Instal Cluster Toolkit.

Menyiapkan image container MaxText

Untuk mempersiapkan citra kontainer MaxText Anda, 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 Registri Artefak;

    gcloud artifacts repositories create ${REPOSITORY_NAME} \
    --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.4 --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:
  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 klaster Cluster Toolkit Anda.

Untuk membuat dan menyebarkan klaster Cluster Toolkit dengan 64v6e Untuk memasang chip TPU, selesaikan langkah-langkah berikut:

  1. Buat peran Identity and Access Management (IAM) kustom, yang diberi nama gke.gcsfuse.profileUser:

    # The GKE TPU v6e blueprint uses GCS Fuse CSI Storage Profiles which requires a custom IAM role.
# If this role is not already created in your project, you must create it before deploying.
gcloud iam roles create gke.gcsfuse.profileUser \
  --project=${PROJECT} \
  --title="GKE GCSFuse Profile User" \
  --description="Allows scanning GCS buckets for objects, retrieving bucket metadata, and creating Anywhere Caches." \
  --permissions="storage.objects.list,storage.buckets.get,storage.anywhereCaches.create,storage.anywhereCaches.get,storage.anywhereCaches.list,storage.anywhereCaches.update"

  2. Membuat bucket Cloud Storage: 

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

  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 yang 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            # Change from storage.objectViewer
    - artifactregistry.reader

  4. Terapkan pengaturan khusus ke blok gke-tpu-v6e-cluster yang menimpa pengaturan IPv6 dan tipe mesin default:

    - id: gke-tpu-v6e-cluster
  source: modules/scheduler/gke-cluster
  use: [gke-tpu-v6e-net-0, workload_service_account]
  settings:
    enable_private_ipv6_google_access: false
    system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
    system_node_pool_taints: []
    enable_private_endpoint: false # Allows access from authorized public IPs
    enable_pathways_for_tpus: $(vars.enable_pathways_for_tpus)
    enable_dataplane_v2: true
    configure_workload_identity_sa: true

  5. Deploy cluster Cluster Toolkit Anda menggunakan blueprint gke-tpu-v6e-advanced.yaml dan teruskan variabel yang diperlukan menggunakan flag --vars:

    ./gcluster deploy examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml \
    --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. Setelah selesai membuat cluster Cluster Toolkit, konfigurasikan Docker:

    # Configure docker for pulling images
gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet

  2. Untuk menyederhanakan perintah berikutnya, konfigurasikan 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}

  3. Konversikan model dari format Hugging Face ke format MaxText dan simpan di bucket Cloud Storage Anda:

    ./gcluster job submit \
  --name="gemma4-hf-to-mt" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --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=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"

  4. Periksa status pekerjaan konversi:

    # Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

Mulai workload pelatihan

Setelah proses konversi selesai, mulailah beban kerja pelatihan RL:

./gcluster job submit \
  --name="gemma4-training" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET}/pathways/" \
  --env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}'); \
      python3 -c \"import pathlib, tpu_inference.layers.common.fused_moe_gmm as f; p = pathlib.Path(f.__file__); p.write_text(p.read_text().replace('onehot_moe_permute_threshold: int = 0,', 'onehot_moe_permute_threshold: int = 100000,'))\"; \
      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} \
      num_batches=50 \
      batch_size=8 \
      rollout_tensor_parallelism=2 \
      rollout_expert_parallelism=4 \
      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 \
      hbm_utilization_vllm=0.55 \
      remat_policy=full \
      async_scheduling=False \
      allow_split_physical_axes=true \
      ragged_gather_reduce_fallback=True \
      vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \
      vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'allow_split_physical_axes': 'true', 'use_ragged_sort': 'false', 'ragged_gather_reduce_fallback': 'true', 'prefuse_moe_weights': 'true', 'weight_dtype': 'bfloat16'}}\""

Periksa status pekerjaan pelatihan:

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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-training --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

Mengonversi kembali model terlatih ke format Hugging Face

Setelah beban kerja pelatihan selesai, konversikan kembali model ke format Wajah Berpelukan:

./gcluster job submit \
  --name="gemma4-mt-to-hf" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --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"

Periksa status pekerjaan konversi:

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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

Pembersihan

Untuk menghindari biaya tambahan, hapus sumber daya yang dibuat selama tutorial ini:

./gcluster destroy ${CLUSTER_NAME} --robust
gcloud storage rm -r gs://${GCS_BUCKET}
gcloud artifacts repositories delete ${REPOSITORY_NAME} --location=${REGION} --project=${PROJECT} --quiet

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

Langkah berikutnya