TPU v6e에서 Gemma 4에 지도 미세 조정 실행

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

목표

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

비용

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• 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 웹사이트에서 학습할 모델의 라이선스 계약에 동의합니다. 이 튜토리얼에서는 gemma4-31b 모델을 사용합니다.

튜토리얼을 완료하는 데 필요한 권한을 얻으려면 관리자에게 프로젝트에 대한 다음 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="gke-tpu-v6e"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-31b"
export HF_TOKEN="YOUR_HF_TOKEN"
export GCS_BUCKET="YOUR_BUCKET_NAME"

다음을 바꿉니다.

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

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.3 --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:
     # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
     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 클러스터 만들기

32개의 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 모듈을 업데이트하여 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. 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=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
       --download-dependencies \
       -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 hf-to-mt \
       --cluster ${CLUSTER_NAME} \
       --project ${PROJECT} \
       --location ${REGION} \
       --compute-type ${TPU_TYPE} \
       --num-slices 1 \
       --image ${CLOUD_IMAGE_NAME} \
       --await-job-completion \
       --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
         python3 -m maxtext.checkpoint_conversion.to_maxtext \
           model_name=${MODEL_NAME} \
           hf_access_token=${HF_TOKEN} \
           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"

변환 작업의 상태를 확인하려면 다음 명령어를 실행합니다.

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

학습 워크로드 시작

변환 프로세스가 완료되면 다음 명령어를 실행하여 SFT 워크로드를 시작할 수 있습니다.

./gcluster job submit --name sft \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --compute-type ${TPU_TYPE} \
    --num-slices 1 \
    --image ${CLOUD_IMAGE_NAME} \
    --await-job-completion \
    --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
      python3 -m maxtext.trainers.post_train.sft.train_sft \
      run_name=sft \
      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} \
      dataset_type=hf \
      hf_path=HuggingFaceH4/ultrachat_200k \
      per_device_batch_size=1 steps=1000 \
      profiler=xplane \
      checkpoint_storage_use_zarr3=0 \
      checkpoint_storage_use_ocdbt=0 \
      skip_jax_distributed_system=False"

학습 작업의 상태를 확인하려면 다음 명령어를 실행합니다.

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

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

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

./gcluster job submit --name mt-to-hf \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --compute-type ${TPU_TYPE} \
    --num-slices 1 \
    --image ${CLOUD_IMAGE_NAME} \
    --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/sft/checkpoints/1000/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"

변환 작업의 상태를 확인하려면 다음 명령어를 실행합니다.

# 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 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/ - though again, it's ~2x the size of the original...

정리

추가 비용이 발생하지 않도록 하려면 이 튜토리얼 중에 만든 리소스를 삭제합니다.

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-32 TPU의 아키텍처 및 구성 세부정보는 TPU v6e를 참조하세요.