Google은 AI 기술을 사용하여 콘텐츠를 사용자의 기본 언어로 번역합니다. AI 번역에는 오류가 있을 수 있습니다.

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

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

워크로드 요약

Cloud TPU 및 GKE에서 이 RL 사후 학습 워크로드를 실행하려면 다음 구성요소가 필요합니다.

모델: Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it)는 고급 개방형 가중치 명령 조정 모델입니다. 이 튜토리얼에서는 Hugging Face에서 모델 가중치를 다운로드하고 MaxText 호환 형식으로 변환한 후 Cloud TPU v6e 인프라를 사용하여 GKE 클러스터에 배포합니다.
강화 학습 (RL) 알고리즘: MaxText 및 Pathways를 사용하는 그룹 상대 정책 최적화 (GRPO) 워크로드는 vLLM 샘플러 백엔드와 페어링된 MaxText의 기본 제공 RL 트레이너(maxtext.trainers.post_train.rl.train_rl)를 사용하여 확장 가능한 학습 후 파이프라인을 실행합니다. GRPO는 프롬프트당 출력 그룹을 생성하고 별도의 비평가 모델 없이 상대적 보상을 평가하여 정책을 최적화합니다.
데이터 세트 및 템플릿: 대규모 수학 명령 조정 데이터 세트인 OpenMathInstruct-2 (openmathinstruct2_rl.json)입니다. 워크로드는 RL 출시 단계에서 구조화된 채팅 템플릿을 사용하여 다단계 수학적 추론 기능을 향상합니다.
인프라 및 프레임워크: Cloud TPU v6e (ct6e-standard-4t, 4x8 토폴로지가 있는 v6e-32 슬라이스)에서 Cluster Toolkit을 사용하여 관리되는 GKE. 프레임워크는 Pathways 및 JAX 백엔드가 있는 MaxText(tpu_post_training:0.2.4)입니다.

목표

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

비용

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

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

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

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

시작하기 전에

이 튜토리얼을 사용하려면 Hugging Face 액세스 토큰이 필요합니다. Hugging Face에서 무료 계정에 가입할 수 있습니다. 계정이 있으면 액세스 토큰을 생성합니다.
1. Hugging Face에 오신 것을 환영합니다 페이지에서 계정 아바타를 클릭하고 액세스 토큰을 선택합니다.
2. 액세스 토큰 페이지에서 새 토큰 만들기를 클릭합니다.
3. 읽기 토큰 유형을 선택하고 토큰 이름을 입력합니다.
4. 액세스 토큰이 표시됩니다. 토큰을 안전한 장소에 저장합니다.

Hugging Face의 google/gemma-4-26b-a4b-it 모델 카드에서 모델의 라이선스 계약에 동의합니다.

튜토리얼을 완료하는 데 필요한 권한을 얻으려면 관리자에게 프로젝트에 대한 다음 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 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"

다음을 바꿉니다.

YOUR_PROJECT_ID: Google Cloud 프로젝트의 ID입니다.
YOUR_REGION: TPU 예약이 있고 클러스터를 배포하는 리전입니다.
YOUR_ZONE: TPU 예약이 있고 클러스터를 배포하는 영역입니다.
YOUR_CLUSTER_NAME: GKE 클러스터의 이름입니다. 이름은 소문자로 시작해야 하며 소문자, 숫자, 하이픈만 포함해야 하고 20자(영문 기준) 이하여야 합니다.
YOUR_BUCKET_NAME: Cloud Storage 버킷의 전역 고유 이름입니다.
YOUR_RESERVATION_NAME: TPU v6e 예약 이름입니다.
YOUR_HF_TOKEN: Hugging Face 액세스 토큰입니다.

Cluster Toolkit 종속 항목 설치

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

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

Cluster Toolkit 설치

CLUSTER_TOOLKIT_VERSION (v1.103.0)에 정의된 특정 버전의 Cluster Toolkit에 대해 사전 빌드된 번들을 다운로드하고 현재 작업 디렉터리에 압축을 풉니다.

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

환경 인증

배포 스크립트를 실행하기 전에 Google Cloud CLI를 인증하고 애플리케이션 기본 사용자 인증 정보 (ADC)를 설정해야 합니다. 이렇게 하면 Cluster Toolkit이 사용자를 대신하여 리소스를 프로비저닝하고 이미지를 Artifact Registry에 푸시할 수 있습니다.

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

Cluster Toolkit 클러스터 만들기

32개의 v6e TPU 칩이 있는 Cluster Toolkit 클러스터를 만들고 배포하려면 다음 단계를 완료합니다.

Cloud Storage 버킷을 만듭니다.

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

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

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

모델을 MaxText 형식으로 변환

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

후속 명령어를 간소화하려면 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:?}"

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

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

이 명령어에는 --await-job-completion 플래그가 포함되어 있습니다. 즉, 변환이 완료될 때까지 현재 터미널에서 대기합니다.

선택사항: 작업이 실행되는 동안 변환 로그를 스트리밍하려면 gcluster를 추출한 작업 디렉터리에서 별도의 터미널을 열고 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

원래 터미널에서 변환된 모델 파일을 Cloud Storage 버킷에서 사용할 수 있는지 확인합니다.
```
gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"
```

학습 워크로드 시작

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

상태를 확인하고 조정자 포드가 준비될 때까지 기다린 후 학습 작업의 로그를 스트리밍합니다.
```
# 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
```
제출 직후 GKE가 Pathways 조정자 CPU 노드 풀 (cpu-np)을 확장하는 동안 ./gcluster job list 명령어의 출력에 Suspended가 잠시 표시될 수 있습니다. 작업이 시작되면 0단계(Pre RL Training: ... mean_reward=...)에서 초기 사전 학습 평가를 출력하고, 초기 XLA 컴파일 중에 몇 분 동안 로그 출력을 일시중지한 후 50단계 동안 학습을 진행하여 단계 번호 없이 각 단계에서 Reshard finished 및 출시 세대를 출력합니다.

다음 단계로 진행하기 전에 로그에 Post RL Training: ... mean_reward=...이 출력되고 ./gcluster job list 명령어의 출력에 gemma4-training 상태가 Succeeded로 표시될 때까지 기다립니다.

참고: GKE에서 완료된 JobSet는 구성된 TTL (수명)이 만료되면 자동으로 삭제됩니다. 이전 학습 작업의 경우 24시간(--gke-ttl-after-finished="24h")입니다. 완료된 워크로드는 TTL이 만료된 후 ./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="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"

선택사항: 이전 --await-job-completion 명령어가 첫 번째 터미널에서 대기하는 동안 실시간 변환 로그를 모니터링하려면 gcluster를 추출한 작업 디렉터리에서 별도의 터미널을 열고 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/

원래 터미널에서 학습된 Hugging Face 모델 가중치와 구성 파일이 Cloud Storage 버킷에 있는지 확인합니다.
```
gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"
```

삭제

추가 요금이 발생하지 않도록 하려면 gcluster destroy 명령어를 사용하여 이 튜토리얼에서 만든 리소스를 삭제합니다.

./gcluster destroy "${CLUSTER_NAME:?}" --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

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

다음 단계

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