MaxText를 사용하여 Qwen3-14b 모델에서 멀티 호스트 지도 미세 조정 실행

이 튜토리얼에서는 Cloud TPU의 MaxText를 사용하여 Qwen3-14b 모델에서 지도 미세 조정 (SFT)을 실행하는 단계별 가이드를 제공합니다. 특수 컨테이너 이미지를 빌드하고, 가속 처리 키트 (XPK)를 사용하여 Pathways로 Google Kubernetes Engine (GKE) 클러스터를 프로비저닝하고, 멀티 호스트 학습 워크로드를 실행하는 방법을 알아봅니다.

목표

• 후속 학습에 최적화된 커스텀 MaxText 컨테이너 이미지를 빌드하는 방법을 알아봅니다.
• Pathways가 사용 설정된 XPK를 사용하여 GKE 클러스터를 프로비저닝합니다.
• Hugging Face 형식에서 MaxText 형식으로 Qwen3 14b 모델을 변환합니다.
• Cloud TPU에서 멀티 호스트 SFT 학습 워크로드를 실행합니다.
• 서빙을 위해 미세 조정된 모델을 Hugging Face 형식으로 다시 변환합니다.

비용

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine, for the temporary build virtual machine (VM)

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

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

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

시작하기 전에

• 사용자 계정 또는 서비스 계정에 다음 역할이 있는지 확인합니다.
  • 빌드 VM을 만드는 roles/compute.admin
  • Docker 저장소를 관리하는 roles/artifactregistry.admin
  • 데이터 버킷을 관리하는 roles/storage.admin
  • Google Kubernetes Engine 클러스터를 만들고 관리하는 roles/container.admin
  • 워크로드 서비스 계정을 만드는 roles/iam.serviceAccountAdmin
  • Identity and Access Management (IAM) 정책을 설정하는 roles/resourcemanager.projectIamAdmin
  • 서비스 계정 역할을 하는 roles/iam.serviceAccountUser
• Google Cloud CLI를 설치하고 초기화합니다.
• 워크스테이션에 Python 3.12 이상이 설치되어 있는지 확인합니다.

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

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

환경 설정

다음 스크립트를 실행하여 환경 변수를 설정합니다.

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
export HF_TOKEN="YOUR_HF_TOKEN"

다음을 바꿉니다.

• YOUR_PROJECT_ID: 프로젝트 ID Google Cloud
• YOUR_REGION: 사용할 리전
• YOUR_ZONE: 사용할 영역
• YOUR_CLUSTER_NAME: Google Kubernetes Engine 클러스터의 이름
• YOUR_GCS_BUCKET: Cloud Storage 버킷의 고유한 이름
• YOUR_RESERVATION_NAME: 용량 예약
• YOUR_HF_TOKEN: Hugging Face 액세스 토큰

MaxText 컨테이너 이미지 준비

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

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

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

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

   gcloud artifacts repositories create maxtext-images \
       --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.1 --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}"

Google Kubernetes Engine 클러스터 만들기

Qwen3 14b 모델에서 SFT 학습을 실행하려면 TPU 칩이 장착된 Google Kubernetes Engine 클러스터가 필요합니다. 가속 처리 키트 (XPK)를 설치하고 Pathways 지원으로 GKE 클러스터를 만듭니다.

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

학습을 위한 모델 준비

CPU 기반 워크로드를 사용하여 기본 모델을 MaxText 형식으로 변환합니다. 이 태스크를 여러 머신에서 동시에 실행하지 마세요. 다음 명령어에는 변환이 하나의 TPU 노드에서만 실행되도록 하는 검사가 포함되어 있습니다.

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

모델 변환 진행 상황 추적

변환 진행 상황을 추적하려면 다음 단계를 따르세요.

1. GKE 클러스터에 예약된 포드를 나열하려면 kubectl get pod 명령어를 실행합니다.
2. qwen-hf-to-mt-slice-job-0-0-HASH라는 포드를 찾습니다.
3. 포드의 출력을 실시간으로 검사하려면 명령어를 실행합니다. kubectl logs -f POD_NAME

학습 워크로드 시작

변환 프로세스가 완료되면 XPK를 사용하여 SFT 미세 조정 워크로드를 시작할 수 있습니다.

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

학습 워크로드 모니터링

XPK 명령줄 인터페이스 (CLI)를 사용하여 워크로드의 상태를 모니터링합니다.

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

로그 및 TPU 사용률을 보려면 Google Cloud 콘솔을 사용하세요. 다음 명령어를 실행하여 로그를 볼 수도 있습니다.

kubectl logs -f qwen-training-pathways-head-0-0-HASH

HASH를 포드 이름의 숫자 해시로 바꿉니다. 이 해시의 값을 확인하려면 kubectl get pod 명령어를 실행하고 반환된 포드 목록을 검사합니다.

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

학습 워크로드가 완료되면 체크포인트를 Hugging Face 형식으로 다시 변환합니다.

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

변환 진행 상황을 추적하려면 kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH 명령어를 실행하고 HASH를 포드 이름의 숫자 해시로 바꿉니다.

변환이 완료되면 gs://$GCS_BUCKET/qwen-3-14b/hf-trained/에 저장된 조정된 모델을 사용할 수 있습니다.

정리

추가 비용이 청구되지 않도록 하려면 Google Kubernetes Engine 클러스터, Cloud Storage 버킷, Artifact Registry 저장소를 비롯하여 이 튜토리얼 중에 만든 리소스를 삭제합니다.

주의: 이 튜토리얼에서 학습시킨 모델을 계속 사용하려면 모델이 포함된 Cloud Storage 버킷을 삭제하기 전에 액세스할 수 있는 버전을 저장하세요.

이 튜토리얼에서 만든 리소스를 삭제하려면 다음 명령어를 실행합니다.

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

다음 단계

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