TPU v6e에서 Gemma 4 26B에 대한 멀티 호스트 강화 학습 훈련을 실행합니다.

이 튜토리얼에서는 MaxText 및 Cluster Toolkit을 사용하여 Tensor Processing Unit (TPU) v6e-64 클러스터에서 다중 호스트 강화 학습 (RL) 학습을 실행하는 방법을 보여줍니다. Cluster Toolkit을 사용하여 다중 호스트 학습 워크로드를 실행하고 제공을 위해 결과를 Hugging Face 형식 으로 다시 내보냅니다.

목표

• Cluster Toolkit 및 종속 항목을 설치합니다.
• MaxText 및 종속 항목을 설치합니다.
• Cluster Toolkit 클러스터를 배포합니다.
• Hugging Face 모델을 MaxText 형식으로 변환합니다.
• TPU v6e 클러스터에서 RL 학습 워크로드를 실행합니다.
• 제공을 위해 미세 조정된 모델을 Hugging Face 형식으로 다시 변환합니다.

비용

이 문서에서는 비용이 청구될 수 있는 구성요소를 사용합니다 Google Cloud.

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine

프로젝트 사용량을 기준으로 예상 비용을 산출하려면 가격 계산기를 사용하세요.

신규 Google Cloud 사용자는 무료 체험판을 사용할 수 있습니다.

이 문서에 설명된 태스크를 완료했으면 만든 리소스를 삭제하여 청구가 계속되는 것을 방지할 수 있습니다. 자세한 내용은 정리를 참조하세요.

중요: 이 튜토리얼을 완료하려면 v6e-64 TPU에 액세스할 수 있어야 합니다. 이 튜토리얼에서는 필요할 때 용량을 확보할 수 있도록 프로젝트에서 이러한 TPU를 이미 예약했다고 가정합니다. Google Cloud 향후 날짜와 시간에 Cloud TPU 용량을 예약하는 옵션에 대해 자세히 알아보려면 Cloud TPU 예약 정보를 참고하세요.

시작하기 전에

이 튜토리얼을 사용하려면 Hugging Face 액세스 토큰이 필요합니다. Hugging Face에서 무료 계정을 등록할 수 있습니다. 계정이 있으면 액세스 토큰을 생성합니다.

1. Hugging Face 시작 페이지에서 계정 아바타를 클릭하고 액세스 토큰을 선택합니다.
2. 액세스 토큰 페이지에서 새 토큰 만들기 를 클릭합니다.
3. 읽기 토큰 유형을 선택하고 토큰 이름을 입력합니다.
4. 액세스 토큰이 표시됩니다. 토큰을 안전한 곳에 저장합니다.

• Hugging Face 웹사이트에서 학습할 모델의 라이선스 계약에 동의합니다. 이 튜토리얼에서는 gemma4-26b 모델을 사용합니다.

이 튜토리얼을 완료하는 데 필요한 권한을 얻으려면 관리자에게 프로젝트에 대한 다음 IAM 역할을 부여해 달라고 요청하세요.

• 이 튜토리얼을 완료하려면 다음 안내를 따르세요.
  • TPU 관리자 (roles/tpu.admin)
  • 서비스 계정 사용자 (roles/iam.serviceAccountUser)
  • Compute 편집자 (roles/compute.editor)
• MaxText 컨테이너 이미지를 준비하려면 다음 안내를 따르세요.
  • Cloud Build 편집자 (roles/cloudbuild.builds.editor)
  • Artifact Registry 관리자 (roles/artifactregistry.admin)
  • 스토리지 관리자 (roles/storage.admin)
  • 서비스 사용량 관리자 (roles/serviceusage.serviceUsageAdmin)

역할 부여에 대한 자세한 내용은 프로젝트, 폴더, 조직에 대한 액세스 관리를 참조하세요.

커스텀 역할이나 다른 사전 정의된 역할을 통해 필요한 권한을 얻을 수도 있습니다.

환경 변수 설정

환경 변수 설정:

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"

다음을 바꿉니다.

• YOUR_PROJECT_ID: 프로젝트 Google Cloud 의 ID입니다.
• YOUR_REGION: 클러스터를 배포할 리전입니다.
• YOUR_ZONE: 클러스터를 배포할 영역입니다.
• YOUR_CLUSTER_NAME: Google Kubernetes Engine 클러스터의 이름입니다.
• YOUR_REPOSITORY_NAME: MaxText 이미지의 Artifact Registry 저장소 이름입니다.
• YOUR_BUCKET_NAME: Cloud Storage 버킷의 전역적으로 고유한 이름입니다.
• YOUR_RESERVATION_NAME: 예약의 이름입니다.
• YOUR_HF_TOKEN: Hugging Face 액세스 토큰입니다.

Cluster Toolkit 종속 항목 설치

Linux 또는 macOS 클라이언트 또는 워크스테이션에서 이 튜토리얼을 완료하려면 종속 항목 설치에서 Cluster Toolkit 문서의 관련 단계를 따르세요.

Cloud Shell을 사용하는 경우 이 섹션을 건너뛸 수 있습니다.

Cluster Toolkit 설치

Cluster Toolkit 설치의 안내에 따라 Cluster Toolkit의 사전 빌드된 번들을 설치합니다.

MaxText 컨테이너 이미지 준비

필요한 종속 항목 설치를 포함하여 MaxText 컨테이너 이미지를 준비하려면 다음 단계를 완료하세요.

1. Cloud Storage 버킷을 만듭니다.

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

2. Artifact Registry 저장소를 만듭니다.

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

3. 다음 콘텐츠와 함께 파일 이름 cloudbuild.yaml을 사용하여 저장소의 루트 디렉터리에 파일을 만듭니다.

   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. Cloud Build를 사용하여 MaxText Docker 이미지를 빌드합니다.

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

Cluster Toolkit 클러스터 만들기

64개의 v6e TPU 칩으로 Cluster Toolkit 클러스터를 만들고 배포하려면 다음 단계를 완료하세요.

1. gke.gcsfuse.profileUser라는 커스텀 Identity and Access Management (IAM) 역할을 만듭니다.

   # 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. Cloud Storage 버킷을 만듭니다.

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

3. 기본적으로 클러스터 노드 풀 서비스 계정에는 Cloud Storage 버킷에 쓸 수 있는 데 필요한 권한이 없습니다. 노드 풀 서비스 계정이 Cloud Storage 버킷에 쓸 수 있도록 하려면 Storage Admin 역할을 부여해야 합니다. 이 역할을 부여하려면 node_pool_service_account라는 service-account 모듈 을 업데이트하여 gke-tpu-v6e-advanced.yaml 파일을 수정합니다.

   - 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. 기본 IPv6 및 머신 유형 설정을 재정의하는 맞춤 설정을 gke-tpu-v6e-cluster 블록에 적용합니다.

   - 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. gke-tpu-v6e-advanced.yaml 청사진을 사용하고 --vars 플래그를 사용하여 필요한 변수를 전달하여 Cluster Toolkit 클러스터를 배포합니다.

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

모델을 MaxText 형식으로 변환

MaxText 형식으로 모델을 학습시키려면 Hugging Face 형식에서 MaxText 형식으로 변환해야 합니다.

1. Cluster Toolkit 클러스터 만들기를 완료한 후 Docker를 구성합니다.

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

2. 후속 명령어를 간소화하려면 기본 프로젝트, 클러스터, 위치를 구성합니다.

   # Configure gcluster Defaults
   ./gcluster job config set project ${PROJECT}
   ./gcluster job config set cluster ${CLUSTER_NAME}
   ./gcluster job config set location ${REGION}

3. 모델을 Hugging Face 형식에서 MaxText 형식으로 변환하고 Cloud Storage 버킷에 저장합니다.

   ./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. 변환 작업 상태를 확인합니다.

   # 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}

학습 워크로드 시작

변환 프로세스가 완료되면 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'}}\""

학습 작업 상태를 확인합니다.

# 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}

학습된 모델을 Hugging Face 형식으로 다시 변환

학습 워크로드가 완료되면 모델을 Hugging Face 형식으로 다시 변환합니다.

./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"

변환 작업 상태를 확인합니다.

# 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/

정리

추가 요금이 발생하지 않도록 하려면 이 튜토리얼 중에 만든 리소스를 삭제하세요.

./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}

다음 단계

• Cloud TPU에 대한 자세한 내용은 Cloud TPU 소개를 참조하세요.
• v6e-64 TPU의 아키텍처 및 구성 세부정보는 TPU v6e를 참조하세요.
• Cluster Toolkit에 대한 자세한 내용은 Cluster Toolkit 개요를 참조하세요.