MaxText를 사용하여 TPU VM에서 지도 미세 조정 실행

이 튜토리얼에서는 대규모 언어 모델 (LLM)을 위한 고성능 JAX 기반 학습 스택인 MaxText를 사용하여 Google Cloud 에서 단일 v6e-8Tensor Processing Unit (TPU) 가상 머신 (VM) 인스턴스에서 지도 미세 조정 (SFT)을 실행하는 방법을 단계별로 안내합니다.

목표

• Cloud TPU VM 인스턴스를 설정합니다.
• MaxText 및 종속 항목을 설치합니다.
• Hugging Face 모델을 MaxText 형식으로 변환합니다.
• TPU에서 SFT 학습 워크로드를 실행합니다.
• 제공을 위해 미세 조정된 모델을 Hugging Face 형식으로 다시 변환합니다.

비용

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

• Cloud TPU

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

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

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

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

시작하기 전에

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

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

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

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

• TPU 관리자 (roles/tpu.admin)
• 서비스 계정 사용자(roles/iam.serviceAccountUser)
• 컴퓨팅 편집자 (roles/compute.editor)

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

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

환경 설정

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

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export NAME="YOUR_TPU_NAME"
export NETWORK="default"

다음을 바꿉니다.

• YOUR_PROJECT_ID: Google Cloud 프로젝트 ID
• YOUR_ZONE: 사용할 영역
• YOUR_RESERVATION_NAME: 용량 예약
• YOUR_TPU_NAME: Cloud TPU VM 인스턴스의 이름

다음 명령어를 실행하여 Google Cloud 로 인증합니다.

gcloud auth login

Cloud TPU VM 만들기

용량 예약에 바인딩된 v6e TPU 칩 8개가 있는 Cloud TPU VM 인스턴스를 만듭니다.

gcloud compute instances create "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --network="${NETWORK}" \
    --tags="${NAME}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

VM 인스턴스를 만든 후 SSH를 사용하여 연결합니다.

gcloud compute ssh "${NAME}" --zone "${ZONE}" --project "${PROJECT}"

TPU VM 인스턴스 내에서 다음 단계를 완료합니다.

MaxText 설치

TPU VM 인스턴스 내에서 시스템 패키지를 업데이트합니다.

sudo apt update && sudo apt upgrade -y --fix-missing

MaxText에 필요한 Python 3.12와 가상 환경 패키지를 설치합니다.

sudo apt install -y build-essential cmake ninja-build

uv을 사용하여 Python 패키지 설치 속도를 높입니다.

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

maxtext_venv라는 가상 환경을 만들고 활성화합니다.

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

MaxText와 학습 후 작업에 필요한 종속 항목을 설치합니다.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.2" --resolution=lowest

다음 명령어를 실행하여 나머지 필수 종속 항목을 설치합니다.

install_tpu_post_train_extra_deps

모델을 MaxText 형식으로 변환

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

Hugging Face 액세스 토큰, 사용하려는 모델의 이름, MaxText 형식으로 모델을 저장할 디렉터리와 같은 환경 변수를 지정합니다.

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='gemma3-4b'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

YOUR_HF_TOKEN을 이전에 만든 Hugging Face 액세스 토큰으로 바꿉니다.

Hugging Face 형식에서 MaxText 형식으로 모델을 변환하려면 다음 스크립트를 실행합니다. 이 변환을 완료하는 데 약 5분이 걸립니다.

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

학습 워크로드 시작

변환 프로세스가 완료되면 SFT 워크로드를 시작할 수 있습니다.

1. SFT 워크로드 학습 매개변수를 구성합니다.

   # -- MaxText configuration --
   export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
   RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   export RUN_NAME
   export STEPS=1000
   export PER_DEVICE_BATCH_SIZE=1
   
   # -- Dataset configuration --
   export DATASET_NAME="HuggingFaceH4/ultrachat_200k"
   export TRAIN_SPLIT="train_sft"
   export TRAIN_DATA_COLUMNS="['messages']"
   
   export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
   export TPU_ACCELERATOR_TYPE=v6e-8
   export TPU_WORKER_ID=0
   TPU_NAME=$(hostname)
   export TPU_NAME
   export TPU_SKIP_MDS_QUERY=1
   export TPU_WORKER_HOSTNAMES=localhost
   export TPU_TOPOLOGY=2x4
   export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
   export TPU_HOST_BOUNDS=1,1,1

2. 학습 작업을 시작합니다. v6e-8 VM 인스턴스에서 약 10분 정도 걸립니다.

   python3 -m maxtext.trainers.post_train.sft.train_sft \
       run_name="${RUN_NAME?}" \
       base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
       model_name="${MODEL_NAME?}" \
       load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
       per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
       steps="${STEPS?}" \
       hf_path="${DATASET_NAME?}" \
       train_split="${TRAIN_SPLIT?}" \
       train_data_columns="${TRAIN_DATA_COLUMNS?}" \
       profiler=xplane

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

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

1. 내보내기 및 학습된 매개변수의 경로를 설정합니다.

   export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
   export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

2. Hugging Face 형식으로 다시 변환을 실행합니다.

   python3 -m maxtext.checkpoint_conversion.to_huggingface \
       model_name="${MODEL_NAME}" \
       load_parameters_path="${POST_TRAIN_PATH}" \
       base_output_directory="${HF_EXPORT}" \
       scan_layers=True \
       use_multimodal=False \
       weight_dtype=bfloat16

변환이 완료되면 /dev/shm/gemma3-4b/hf-trained에 저장된 조정된 모델을 사용할 수 있습니다. VM이 재부팅되면 /dev/shm 폴더의 콘텐츠에 액세스할 수 없으므로 조정된 모델을 영구 스토리지로 이동하거나 Hugging Face Hub에 업로드해야 합니다.

삭제

추가 요금이 발생하지 않도록 이 튜토리얼에서 만든 리소스를 삭제합니다.

TPU VM 인스턴스 삭제

Cloud TPU VM 인스턴스를 종료한 후 삭제합니다.

# 1. Delete TPU instance
echo "Deleting TPU instance: ${NAME}..."
gcloud compute instances delete "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --quiet || true

# 2. Delete IAP firewall rule
echo "Deleting firewall rule: ${FIREWALL_RULE_NAME}..."
gcloud compute firewall-rules delete "${FIREWALL_RULE_NAME}" \
    --project="${PROJECT}" \
    --quiet || true

다음 단계

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