멀티 호스트 TPU v6e에서 Gemma 4 31B의 SFT 실행

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

워크로드 요약

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

• 모델: 고급 가중치 공개 모델인 Gemma 4 31B (google/gemma-4-31b-it)입니다. 이 튜토리얼에서는 Hugging Face에서 모델 가중치를 다운로드하고 MaxText 호환 형식으로 변환한 후 Cloud TPU v6e (v6e-32) 인프라를 사용하여 GKE 클러스터에 배포합니다.
• 학습 방법: MaxText를 사용한 SFT 워크로드는 JAX를 사용하여 Cloud TPU v6e에서 확장 가능한 미세 조정 파이프라인을 실행합니다. MaxText의 기본 제공 SFT 트레이너(maxtext.trainers.post_train.sft.train_sft)와 안내 조정 데이터 세트를 사용하여 1,000단계에 걸쳐 모델 파라미터를 최적화합니다.
• 데이터 세트 및 템플릿: 대규모 고품질 멀티턴 대화 데이터 세트인 UltraChat 200k (HuggingFaceH4/ultrachat_200k)입니다. 워크로드는 이 데이터 세트를 사용하여 모델을 미세 조정하고 대화 및 명령 추종 기능을 개선합니다.
• 인프라 및 프레임워크: Cloud TPU v6e (v6e-32, 4x8 토폴로지 사용)에서 Cluster Toolkit을 사용하여 관리되는 GKE 프레임워크는 JAX 백엔드가 있는 MaxText (tpu_post_training:0.2.4)입니다.

목표

• Cluster Toolkit 및 종속 항목을 설치합니다.
• Cluster Toolkit 클러스터를 배포합니다.
• Hugging Face 모델을 MaxText 형식으로 변환합니다.
• TPU에서 SFT 학습 워크로드를 실행합니다.
• 파인 튜닝된 모델을 서빙을 위해 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의 google/gemma-4-31b-it 모델 카드에서 모델의 라이선스 계약에 동의합니다.

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

• 이 튜토리얼을 완료하려면 다음이 필요합니다.
  • TPU 관리자 (roles/tpu.admin)
  • 서비스 계정 사용자(roles/iam.serviceAccountUser)
  • Compute 관리자 (roles/compute.admin)
  • 스토리지 관리자 (roles/storage.admin)
  • Kubernetes Engine 관리자(roles/container.admin)
  • 서비스 계정 관리자 (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 TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

다음을 바꿉니다.

• YOUR_PROJECT_ID: Google Cloud프로젝트의 ID
• YOUR_REGION: TPU 예약이 있고 클러스터를 배포하는 리전입니다.
• YOUR_ZONE: TPU 예약이 있고 클러스터를 배포하는 영역입니다.
• YOUR_BUCKET_NAME: Cloud Storage 버킷의 전역 고유 이름입니다.
• YOUR_RESERVATION_NAME: TPU v6e 예약 이름입니다.
• YOUR_HF_TOKEN: Hugging Face 액세스 토큰입니다.

기본적으로 스크립트는 MODEL_NAME를 gemma4-31b (google/gemma-4-31b-it의 MaxText 모델 이름)로, CLUSTER_NAME를 gke-tpu-v6e로 설정합니다. CLUSTER_NAME를 맞춤설정하는 경우 이름은 소문자로 시작해야 하고 소문자, 숫자, 하이픈만 포함해야 하며 생성된 서비스 계정 ID가 30자(영문 기준) 글자 수 제한을 넘지 않도록 20자(영문 기준) 글자 수 제한 이하여야 합니다.

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

참고: Cluster Toolkit 설치 문서에서는 최신 출시 버전을 cluster-toolkit 디렉터리에 다운로드하는 방법을 설명합니다. 하지만 청사진 변수는 Cluster Toolkit 버전에 따라 다를 수 있으므로 이 튜토리얼에서는 최신 출시 버전이 아닌 CLUSTER_TOOLKIT_VERSION (v1.103.0)에 지정된 정확한 버전을 설치해야 합니다. 번들을 추출하면 다음 단계에서 사용하는 gcluster 바이너리와 examples/ 청사진이 제공됩니다. 번들을 cluster-toolkit 디렉터리에 추출한 경우 후속 명령어를 실행하기 전에 해당 디렉터리 (cd cluster-toolkit)로 변경해야 합니다.

환경 인증

Cluster Toolkit에서 리소스를 프로비저닝하고 이미지를 Artifact Registry에 푸시할 수 있도록 Google Cloud CLI를 인증하고 애플리케이션 기본 사용자 인증 정보를 설정합니다.

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
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 클러스터를 만들고 배포하려면 다음 단계를 완료합니다.

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

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

2. 청사진 구성을 스테이징하고 배포 청사진 예시를 격리된 작업 디렉터리에 복사합니다.

   mkdir -p tmp
   cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

3. tmp/gke-tpu-v6e-advanced.yaml를 업데이트하여 클러스터 노드 풀 서비스 계정 (gke-np-sa)에 Cloud Storage 버킷에 체크포인트를 작성할 수 있는 스토리지 관리자 역할 (roles/storage.admin)을 부여하고, 시스템 노드 풀 머신 유형을 e2-standard-8로 전환하고, 리전별 머신 유형 재고 부족을 방지하기 위해 시스템 노드 풀을 대상 영역에 고정합니다.

   # Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
   sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
   # Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
   sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
   # 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)]' tmp/gke-tpu-v6e-advanced.yaml

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

   ./gcluster deploy tmp/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 \
       -l IGNORE \
       --auto-approve -w

모델을 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 hf-to-mt \
       --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"

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

3. 선택사항: 작업이 실행되는 동안 변환 로그를 스트리밍하려면 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 hf-to-mt --main-only -f

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

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

학습 워크로드 시작

1. 변환 프로세스가 완료되면 SFT 워크로드를 시작합니다.

   ./gcluster job submit --name sft \
       --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"

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

2. 선택사항: 작업이 실행되는 동안 실시간 학습 로그를 스트리밍하려면 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 sft --main-only -f

   초기 XLA 컴파일이 완료되면 조정자는 1, 000단계 실행의 각 단계에 대한 출력 진행 상황을 completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ... 형식으로 로깅합니다. 학습 손실은 5단계에서 약 1.57에서 1,000단계에서 약 0.82로 감소합니다. 다음 섹션으로 진행하기 전에 1,000단계가 완료되고 ./gcluster job list 명령어의 출력에 sft 상태가 Succeeded로 표시될 때까지 기다립니다.

   참고: GKE에서 완료된 JobSet는 구성된 TTL (수명)이 만료되면 자동으로 삭제됩니다. 기본값은 1시간 (--gke-ttl-after-finished=1h)입니다. 완료된 워크로드는 TTL이 만료된 후 ./gcluster job list 명령어 출력에 나열되지 않습니다. 하지만 모든 체크포인트와 모델 가중치는 Cloud Storage 버킷에 저장됩니다.

3. 원래 터미널에서 SFT 학습 체크포인트가 Cloud Storage 버킷에 생성되었는지 확인합니다.

   gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

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

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

   ./gcluster job submit --name mt-to-hf \
       --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 \
           --override_model_architecture"

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

2. 선택사항: 작업이 실행되는 동안 실시간 변환 로그를 스트리밍하려면 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 mt-to-hf --main-only -f
   
   # The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ 

3. 원래 터미널에서 최종 Hugging Face Safetensors 가중치가 Cloud Storage 버킷에 생성되었는지 확인합니다.

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

삭제

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

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

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

다음 단계

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