TPU v6e에서 Qwen3-30b-a3b에 대한 멀티 호스트 RL 학습을 실행합니다.

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

목표

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

비용

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

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

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

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

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

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

시작하기 전에

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

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

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

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

• TPU 관리자 (roles/tpu.admin)
• Kubernetes Engine 관리자(roles/container.admin)
• Compute 관리자 (roles/compute.admin)
• 스토리지 관리자 (roles/storage.admin)
• 서비스 계정 사용자(roles/iam.serviceAccountUser)
• 서비스 계정 관리자 (roles/iam.serviceAccountAdmin)
• 프로젝트 IAM 관리자 (roles/resourcemanager.projectIamAdmin)
• 서비스 사용량 관리자(roles/serviceusage.serviceUsageAdmin)

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

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

환경 변수 설정

환경 변수 설정:

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"

다음을 바꿉니다.

• YOUR_PROJECT_ID: Google Cloud 프로젝트의 ID입니다.
• YOUR_REGION: 클러스터를 배포할 리전입니다.
• YOUR_ZONE: 클러스터를 배포할 영역입니다.
• YOUR_CLUSTER_NAME: Google Kubernetes Engine 클러스터의 이름 (최대 20자)입니다.
• 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의 사전 빌드 번들을 설치합니다.

예를 들어 다음과 같이 번들을 다운로드하여 현재 작업 디렉터리에 추출할 수 있습니다.

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

번들을 현재 작업 디렉터리에 추출하면 다음 단계에서 사용할 gcluster 바이너리와 examples/ 블루프린트가 제공됩니다.

Cluster Toolkit 클러스터 만들기

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

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

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

2. Cluster Toolkit 청사진을 현재 작업 디렉터리에 복사합니다.

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

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
       - artifactregistry.reader

4. gcluster deploy 명령어를 사용하여 청사진 gke-tpu-v6e-advanced.yaml을 사용하고 --vars 플래그를 사용하여 필수 변수를 전달하여 Cluster Toolkit 클러스터를 배포합니다.

   ./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. Container Registry 인증을 구성하고 GKE 서비스 계정에 스토리지 관리자 역할(roles/storage.admin)을 부여합니다.

   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

모델을 MaxText 형식으로 변환

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

1. 후속 명령어를 간소화하려면 gcluster job config 명령어를 사용하여 기본 프로젝트, 클러스터, 위치를 구성합니다.

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

2. gcluster job submit 명령어를 사용하여 Hugging Face 형식의 모델을 MaxText 형식으로 변환하고 Cloud Storage 버킷에 저장합니다.

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

4. 변환된 모델 파일을 Cloud Storage 버킷에서 사용할 수 있는지 확인합니다.

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

학습 워크로드 시작

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

• 학습 작업의 상태를 확인하려면 다음 단계를 따르세요.

  # 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

  참고: GKE에서 완료된 JobSet은 기본적으로 1시간 후에 자동으로 정리됩니다 (gcluster job submit의 --gke-ttl-after-finished=1h 설정에 따라 관리됨). 워크로드가 더 이상 ./gcluster job list의 출력에 나열되지 않는 경우 이 정리 작업은 예상되는 동작이며 작업이 실패했음을 의미하지 않습니다. 모든 체크포인트와 모델 가중치는 Cloud Storage 버킷에 안전하게 유지됩니다.

• 학습 체크포인트가 Cloud Storage 버킷에 생성되었는지 확인하려면 다음 단계를 따르세요.

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

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

학습 워크로드가 완료되면 모델을 다시 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"

• 변환 작업의 상태를 확인하려면 다음 단계를 따르세요.

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

• 학습된 Hugging Face 모델 가중치와 구성 파일이 Cloud Storage 버킷에 있는지 확인하려면 다음 단계를 따르세요.

  gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/"

삭제

추가 요금이 청구되지 않도록 하려면 gcluster destroy 명령을 사용하여 이 튜토리얼에서 만든 리소스를 삭제하세요.

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

다음 단계

• Cloud TPU에 대한 자세한 내용은 Cloud TPU 소개를 참고하세요.
• v6e-32 TPU의 아키텍처 및 구성에 관한 자세한 내용은 TPU v6e를 참고하세요.
• Cluster Toolkit에 대한 자세한 내용은 Cluster Toolkit 개요를 참고하세요.