TPU7x(Ironwood)를 사용한 모델 학습

이 문서에서는 TPU7x 리소스를 프로비저닝하는 방법을 설명하고 MaxText 및 XPK를 사용하여 학습 워크로드를 배포하는 예를 보여줍니다.

TPU7x는 Ironwood 제품군으로서 Google Cloud의 7세대 TPU의 첫 번째 버전입니다. Ironwood 세대는 대규모 AI 학습 및 추론을 위해 설계되었습니다. 자세한 내용은 TPU7x를 참조하세요.

TPU7x에 최적화된 더 많은 예시는 GitHub의 Ironwood TPU 학습 레시피를 참조하세요.

TPU 프로비저닝

다음 방법을 사용하여 TPU7x를 프로비저닝하고 관리할 수 있습니다.

GKE: GKE를 사용하여 컨테이너화된 머신러닝 워크로드용 가속기 풀로 TPU를 프로비저닝하고 관리할 수 있습니다. Google Cloud CLI를 사용하여 GKE 클러스터 인스턴스를 수동으로 만들어 기존 프로덕션 GKE 환경을 정확하게 맞춤설정하거나 확장합니다. 자세한 내용은 GKE의 TPU 정보를 참조하세요.
GKE 및 XPK: XPK는 GKE에서 클러스터 생성과 워크로드 실행을 간소화하는 명령줄 도구입니다. ML 실무자가 심도 있는 Kubernetes 전문 지식이 없어도 TPU를 프로비저닝하고 학습 작업을 실행할 수 있도록 설계되었습니다. XPK를 사용하여 GKE 클러스터를 빠르게 만들고 개념 증명 및 테스트를 위한 워크로드를 실행합니다. 자세한 내용은 XPK GitHub 저장소를 참조하세요.
GKE 및 TPU Cluster Director: TPU Cluster Director는 모든 용량 모드 예약을 통해 사용할 수 있으며, 이를 통해 예약된 모든 용량에 제한 없이 전부 액세스하고 TPU 하드웨어 토폴로지, 사용률 상태, 상태를 완전히 파악할 수 있습니다. 자세한 내용은 모든 용량 모드 개요를 참조하세요.

MaxText 및 XPK로 학습 워크로드 배포

가속 처리 키트(XPK)를 사용하여 개념 증명 및 테스트를 위한 GKE 클러스터를 만듭니다. XPK는 머신러닝 워크로드의 프로비저닝, 관리, 실행을 간소화하도록 설계된 명령줄 도구입니다.

다음 섹션에서는 MaxText 및 XPK를 사용하여 학습 워크로드를 배포하는 방법을 보여줍니다.

시작하기 전에

시작하기 전에 다음 단계를 완료하세요.

결제가 사용 설정된 Google Cloud 프로젝트가 있어야 합니다.
TPU7x 액세스 권한을 획득합니다. 자세한 내용은 계정팀에 문의하세요.
XPK와 함께 사용하는 계정에 XPK GitHub 저장소에 나열된 역할이 있는지 확인합니다.

XPK 및 종속 항목 설치

XPK를 설치합니다. XPK GitHub 저장소의 안내를 따르세요.
관리자가 제공한 안내에 따라 Docker를 설치하거나 공식 설치 안내를 따릅니다. 설치가 완료되면 다음 명령어를 실행하여 Docker를 구성하고 설치를 테스트합니다.
```
gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker
```
다음 환경 변수를 설정합니다.
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export RESERVATION_NAME=YOUR_RESERVATION_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
다음을 바꿉니다.
- YOUR_PROJECT_ID: Google Cloud 프로젝트 ID입니다.
- YOUR_ZONE: 클러스터를 만들 영역입니다. 프리뷰 버전에서는 us-central1-c만 지원됩니다.
- YOUR_CLUSTER_NAME: 새 클러스터의 이름입니다.
- YOUR_ACCELERATOR_TYPE: TPU 버전 및 토폴로지입니다. 예를 들면 tpu7x-4x4x8입니다. 지원되는 토폴로지 목록은 지원되는 구성을 참조하세요.
- YOUR_RESERVATION_NAME: 예약의 이름입니다. 공유 예약의 경우 projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME을 사용합니다.
- YOUR_BUCKET_NAME: Cloud Storage 버킷의 이름입니다. 모델 학습의 출력 디렉터리가 됩니다.

기존 Cloud Storage 버킷이 없으면 다음 명령어를 사용하여 버킷을 만듭니다.

gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

단일 NIC, 단일 슬라이스 클러스터 만들기

MTU 구성 섹션의 안내에 따라 네트워크 구성을 최적화합니다.
xpk cluster create 명령어에서 사용할 ${CLUSTER_ARGUMENTS} 변수를 채웁니다.
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"
```
xpk cluster create 명령어를 사용하여 TPU7x 노드 풀이 있는 GKE 클러스터를 만듭니다.
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
--cluster-cpu-machine-type 플래그를 n1-standard-8 이상으로 설정하면 기본 노드 풀에 시스템 포드(예: JobSet 웹훅)를 위한 충분한 CPU가 확보되어 오류를 방지할 수 있습니다. 기본적으로 XPK는 e2-standard-16을 사용합니다. 일부 영역에서는 특정 CPU 유형만 지원하므로 n1, n2, e2 유형 간에 변경해야 할 수 있습니다. 그러지 않으면 할당량 오류가 발생할 수 있습니다.
클러스터의 업그레이드를 방지하기 위해 유지보수 제외를 추가합니다.
```
gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"
```
다음을 바꿉니다.
- EXCLUSION_START_TIME: 자신이 선택한 YYYY-MM-DDTHH:MM:SSZ 형식의 유지보수 제외 시작 시간입니다.
- EXCLUSION_END_TIME: 자신이 선택한 YYYY-MM-DDTHH:MM:SSZ 형식의 유지보수 제외 종료 시간입니다.

MaxText Docker 이미지 빌드 또는 업로드

MaxText에서 제공하는 스크립트를 사용하여 로컬로 Docker 이미지를 빌드하거나 사전 빌드된 이미지를 사용할 수 있습니다.

로컬에서 빌드

다음 명령어는 로컬 디렉터리를 컨테이너에 복사합니다.

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

명령어를 성공적으로 실행하면 maxtext_base_image라는 이미지가 로컬에 생성됩니다. xpk 워크로드 명령어에서 로컬 이미지를 직접 사용할 수 있습니다.

이미지 업로드(선택사항)

이전 섹션의 안내에 따라 Docker 이미지를 로컬로 빌드한 후 다음 명령어를 사용하여 MaxText Docker 이미지를 레지스트리에 업로드할 수 있습니다.

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

이 명령어를 실행하면 gcr.io에 gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME이라는 이름의 MaxText 이미지가 표시됩니다.

MaxText 학습 명령어 정의

Docker 컨테이너 내에서 학습 스크립트를 실행할 명령어를 준비합니다.

MaxText 1B 모델은 약 10억 개의 파라미터로 언어 모델을 학습하도록 설계된 MaxText 프레임워크 내의 구성입니다. 소규모 칩 실험에 이 모델을 사용합니다. 성능이 최적화되지 않았습니다.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_xpk \
        steps=30"

학습 워크로드 배포

xpk workload create 명령어를 실행하여 학습 작업을 배포합니다. MaxText 기본 이미지를 사용하려면 --base-docker-image 플래그를 지정하거나 --docker-image 플래그와 사용할 이미지를 지정해야 합니다. --enable-debug-logs 플래그를 포함하여 디버그 로깅을 사용 설정할 수 있습니다.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

워크로드 이름은 클러스터 내에서 고유해야 합니다. 이 예에서는 고유성을 보장하기 위해 워크로드 이름에 $(date +%H%M)이 추가됩니다.