Google은 AI 기술을 사용하여 콘텐츠를 사용자의 기본 언어로 번역합니다. AI 번역에는 오류가 있을 수 있습니다.

Ray 및 TorchTPU를 사용하여 멀티 호스트 TPU에서 모델 미세 조정

Autopilot Standard

이 튜토리얼에서는 Ray Core, Ray Train, PyTorch 완전 분할 데이터 동시 로드 (FSDP2), TorchTPU를 사용하여 Google Kubernetes Engine (GKE)의 멀티 호스트 Cloud TPU 슬라이스에서 분산 PyTorch 워크로드를 실행하고 Gemma 4 모델을 미세 조정하는 방법을 보여줍니다.

이 튜토리얼은 Kubernetes의 멀티 호스트 TPU 슬라이스에서 분산 PyTorch 학습 워크로드를 조정하려는 머신러닝 (ML) 엔지니어, 플랫폼 관리자 및 운영자, 데이터 및 AI 전문가를 대상으로 합니다. Google Cloud 콘텐츠에서 언급된 일반적인 역할 및 예시 태스크에 대해 자세히 알아보려면 일반 GKE 사용자 역할 및 태스크를 참조하세요.

이 튜토리얼을 시작하기 전에 다음 사항을 숙지해야 합니다.

배경

이 섹션에서는 이 튜토리얼에서 사용되는 주요 기술을 설명합니다.

TPU

Tensor Processing Unit (TPU)은 머신러닝 워크로드를 위해 Google에서 커스텀 설계한 가속기입니다. 이 튜토리얼은 TPU Trillium (v6e) 및 Ironwood (TPU7x)를 모두 지원합니다. 멀티 호스트 TPU 슬라이스에서는 여러 TPU 호스트가 고속 칩 간 상호 연결 (ICI) 네트워크를 통해 직접 통신하여 단일 칩 또는 TPU 호스트의 메모리 용량을 초과하는 모델을 학습시키는 데 필요한 집계 고대역폭 메모리 (HBM)와 지연 시간이 짧은 집합 통신을 제공합니다.

TorchTPU

TorchTPU는 Google과 Meta가 공동으로 개발한 Cloud TPU용 오픈소스 PyTorch 백엔드입니다. TorchTPU를 사용하면 최소한의 코드 변경으로 TPU 하드웨어에서 기존 PyTorch 모델을 실행하고, XLA (Accelerated Linear Algebra) 컴파일러를 통해 단일 칩에서 멀티 슬라이스 TPU Pod로 학습 및 추론 워크로드를 확장할 수 있습니다. TorchTPU는 TPU ICI 네트워크에서 기기당 하나의 프로세스 실행 모델을 사용하며 즉시 실행, torch.compile, torch.distributed, DeviceMesh, DTensor, 완전 샤딩 데이터 동시 로드 (FSDP2)를 비롯한 표준 PyTorch API를 지원합니다.

TPU의 KubeRay 및 Ray

KubeRay는 Ray 애플리케이션을 배포하고 관리하기 위한 Kubernetes 연산자입니다. GKE 기반 Ray 부가기능을 사용 설정하면 관리형 KubeRay 연산자와 KubeRay TPU 웹훅이 배포됩니다. KubeRay가 google.com/tpu 리소스를 요청하는 워커 포드를 만들면 웹훅은 멀티 호스트 조정에 필요한 포드 호스트 이름, 헤드리스 서비스, TPU 환경 변수를 자동으로 구성합니다.

Ray는 시작 시 각 노드의 TPU 토폴로지를 감지하고 TPU 슬라이스 전반에 분산 워크로드를 예약합니다. 두 가지 수준에서 분산 TorchTPU 워크로드를 조정할 수 있습니다.

Ray Core는 텐서 병렬 계산 또는 강화 학습 (RL) 파이프라인과 같은 커스텀 분산 실행 루프를 위한 토폴로지 인식 배치 그룹 및 액터 예약을 제공합니다.
Ray Train은 PyTorch 분산 데이터 병렬 (DDP), PyTorch FSDP2, JAX 단일 프로그램, 다중 데이터 (SPMD)를 포함한 관리형 분산 데이터 병렬 및 모델 병렬 학습을 제공하며, 체크포인트 및 자동 장애 복구 기능이 기본 제공됩니다.

목표

이 튜토리얼에서는 다음 작업을 처리하는 방법을 보여줍니다.

TPU Trillium (v6e) 또는 Ironwood (TPU7x)에서 Autopilot 또는 Standard로 GKE 클러스터를 만듭니다.
TorchTPU 기본 이미지 위에 Ray 및 Hugging Face 학습 라이브러리를 레이어링하는 컨테이너 이미지를 빌드합니다.
RayCluster 커스텀 리소스를 배포하여 헤드 및 TPU 워커 노드를 프로비저닝합니다.
Ray Core를 사용하여 TPU 슬라이스에서 PyTorch DeviceMesh 및 DTensor로 분산 텐서 병렬 계산을 실행합니다.
Ray Train을 사용하여 멀티 호스트 TPU 슬라이스에서 PyTorch FSDP2로 Gemma 4 모델을 파인 튜닝합니다.

비용

이 튜토리얼에서는 비용이 청구될 수 있는 Google Cloud의 구성요소를 사용합니다.

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

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

시작하기 전에

시작하기 전에 다음 태스크를 수행했는지 확인합니다.

Artifact Registry API, Cloud Storage API, Google Kubernetes Engine API를 사용 설정합니다.

API 사용 설정

이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치했으면 gcloud components update 명령어를 실행하여 최신 버전을 가져옵니다. 이전 gcloud CLI 버전에서는 이 문서의 명령어를 실행하지 못할 수 있습니다.
참고: 기존 gcloud CLI 설치의 경우 compute/region 속성을 설정해야 합니다. 주로 영역 클러스터를 사용하는 경우에는 대신 compute/zone을 설정합니다. 기본 위치를 설정하면 gcloud CLI에서 One of [--zone, --region] must be supplied: Please specify location과 같은 오류를 방지할 수 있습니다. 클러스터의 위치가 설정한 기본값과 다른 경우 특정 명령어에서 위치를 지정해야 할 수 있습니다.

선택한 영역에서 TPU Trillium (v6e) 또는 Ironwood (TPU7x) 용량에 대한 프로젝트 할당량이 충분한지 확인합니다. 자세한 내용은 Cloud TPU 할당량을 참조하세요.
계정에 다음 IAM 역할이 있는지 확인합니다.
- Kubernetes Engine 관리자(roles/container.admin)
- Artifact Registry 관리자 (roles/artifactregistry.admin)
- 스토리지 관리자(roles/storage.admin)

개발 환경 준비

이 튜토리얼에서는 Cloud Shell을 사용하여 Google Cloud에서 호스팅되는 리소스를 관리합니다. Cloud Shell에는 kubectl, docker, python3, gcloud CLI를 비롯해 이 튜토리얼에 필요한 소프트웨어가 사전 설치되어 있습니다.

Cloud Shell로 환경을 설정하려면 다음 단계를 따르세요.

Google Cloud 콘솔의 경우 Google Cloud 콘솔에서 Cloud Shell 활성화를 클릭하여 Cloud Shell 세션을 시작합니다. 이 작업을 수행하면 Google Cloud 콘솔의 하단 창에서 세션이 실행됩니다.

kubernetes-engine-samples 저장소를 클론하고 이 튜토리얼의 샘플 디렉터리로 변경합니다.

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples/ai-ml/gke-ray/raytrain/torchtpu

Ray CLI가 설치된 로컬 Python 가상 환경을 만들고 활성화합니다.

python3 -m venv ~/ray-env
source ~/ray-env/bin/activate
pip install "ray[default]==2.58.0"

Gemma 4 모델 가중치를 다운로드하려면 Hugging Face에서 Google 라이선스 계약을 확인하세요.
1. Hugging Face의 Gemma 4 모델 페이지로 이동합니다.
2. 로그인하고 동의 및 저장소 액세스를 클릭하여 라이선스 약관에 동의합니다.
3. Hugging Face 계정 설정에서 Read 권한이 있는 사용자 액세스 토큰을 생성합니다.
Google Cloud 프로젝트, 클러스터, Cloud Storage 버킷, Hugging Face 토큰의 환경 변수를 설정합니다.
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=CLUSTER_NAME
export REGION=REGION
export ZONE=ZONE
export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
export BUCKET_NAME=BUCKET_NAME
export HF_TOKEN=HF_TOKEN
export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu:latest
gcloud config set project ${PROJECT_ID}
```
다음을 바꿉니다.
- PROJECT_ID: Google Cloud 프로젝트 ID입니다.
- CLUSTER_NAME: GKE 클러스터의 이름입니다(예: ray-torchtpu-cluster).
- REGION: Artifact Registry 저장소, Cloud Storage 버킷, TPU 용량이 있는 리전입니다.
- ZONE: TPU 용량을 사용할 수 있는 영역입니다. TPU 버전별로 사용 가능한 리전 및 영역은 TPU 리전 및 영역을 참조하세요.
- ARTIFACT_REGISTRY: Artifact Registry Docker 저장소의 이름입니다(예: ray-torchtpu).
- BUCKET_NAME: Ray Train 체크포인트를 저장할 Cloud Storage 버킷의 이름입니다(예: PROJECT_ID-ray-torchtpu).
- HF_TOKEN: Hugging Face 액세스 토큰입니다.

GKE 클러스터 만들기

관리형 Ray 연산자 부가기능이 사용 설정된 GKE Autopilot 또는 Standard 클러스터에서 분산 TorchTPU 학습 워크로드를 실행할 수 있습니다.

GKE 클러스터를 만들고 kubectl 명령줄 도구의 사용자 인증 정보를 구성합니다.
Autopilot
1. Ray Operator 부가기능을 사용 설정하여 리전 GKE Autopilot 클러스터를 만듭니다.
  
  gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator
  
  이후 단계에서 RayCluster 커스텀 리소스를 배포하면 GKE Autopilot은 포드 리소스 요청 및 노드 선택기와 일치하는 CPU 헤드 노드와 멀티 호스트 TPU 슬라이스를 자동으로 프로비저닝합니다.
2. 클러스터와 통신하도록 kubectl 명령줄 도구를 구성합니다.
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
표준
1. Ray 헤드 포드, 이미지 스트리밍, GKE Dataplane V2, GKE용 워크로드 아이덴티티 제휴가 사용 설정된 CPU 노드 풀이 있는 영역 GKE Standard 클러스터를 만듭니다.
  
  gcloud container clusters create ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --cluster-version=1.36 \ --machine-type=n2-standard-8 \ --num-nodes=1 \ --addons=RayOperator \ --enable-image-streaming \ --enable-dataplane-v2 \ --workload-pool=${PROJECT_ID}.svc.id.goog
  참고: GKE RayOperator 부가기능을 사용하는 대신 오픈소스 KubeRay를 자체 관리하는 경우 TorchTPU 환경 변수 삽입을 위해 KubeRay 및 KubeRay TPU 웹훅 버전 1.4.0 이상을 설치합니다.
2. 클러스터와 통신하도록 kubectl 명령줄 도구를 구성합니다.
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
3. TPU Trillium (v6e)의 호스트 4개에 걸친 4x4 토폴로지 또는 Ironwood (TPU7x)의 호스트 2개에 걸친 2x2x2 토폴로지를 사용하여 16개 기기 멀티 호스트 TPU 슬라이스 노드 풀을 만듭니다.
  
  Trillium(v6e)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=ct6e-standard-4t \ --tpu-topology=4x4 \ --num-nodes=4
  
  Ironwood (TPU7x)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=tpu7x-standard-4t \ --tpu-topology=2x2x2 \ --num-nodes=2
  
  도움말: 특정 TPU 예약을 사용하려면 --reservation-affinity=specific 및 --reservation=RESERVATION_NAME 플래그를 node-pools create 명령어에 추가합니다. 여기서 RESERVATION_NAME는 TPU 예약의 이름입니다. 스팟 VM을 사용하려면 --spot 플래그를 추가합니다.

멀티 호스트 TPU 노드에서 공유 Ray Train 체크포인트를 저장할 Cloud Storage 버킷을 만들고 default Kubernetes ServiceAccount 리소스에 버킷에 대한 액세스 권한을 부여합니다.

gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}

PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
gcloud storage buckets add-iam-policy-binding gs://${BUCKET_NAME} \
    --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/default/sa/default" \
    --role="roles/storage.objectAdmin"

Ray 포드가 제어된 Gemma 4 가중치를 다운로드할 수 있도록 Hugging Face 액세스 토큰으로 Kubernetes 보안 비밀을 만듭니다.
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

컨테이너 이미지 빌드

TorchTPU 및 libtpu 라이브러리가 포함된 사전 빌드된 TorchTPU 기본 이미지 위에 Ray 및 Hugging Face 학습 라이브러리를 설치하는 컨테이너 이미지를 빌드합니다.

Artifact Registry Docker 저장소를 만들고 Docker 인증을 구성합니다.

gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
    --repository-format=docker \
    --location=${REGION}

gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev

샘플 디렉터리에서 Dockerfile를 검사합니다.

FROM us-docker.pkg.dev/ml-oss-artifacts-transient/torch-tpu-docker-container/torch-tpu:nightly

# Install wget for KubeRay health probes, Ray (Core, Dashboard, and Train), and Hugging Face libraries for fine-tuning
RUN apt-get update && apt-get install -y --no-install-recommends wget \
    && rm -rf /var/lib/apt/lists/* \
    && pip install --no-cache-dir \
    "ray[default,train]==2.58.0" \
    "transformers>=5.5.0" \
    datasets

Ray TorchTrainer 클래스를 사용한 단일 슬라이스 학습과 Ironwood (TPU7x)의 칩렛별 예약에는 Ray 버전 2.58.0 이상이 필요하며, Gemma 4를 파인 튜닝하려면 transformers 라이브러리 버전 5.5.0 이상이 필요합니다.

이미지를 빌드하고 Artifact Registry에 푸시합니다.

docker build -t ${IMAGE_URI} .
docker push ${IMAGE_URI}

RayCluster 커스텀 리소스 배포

CPU 헤드 포드, Ray 공유 메모리 객체 스토어용 메모리 내 /dev/shm 볼륨, 멀티 호스트 TPU 워커 그룹을 프로비저닝하는 RayCluster 커스텀 리소스를 배포합니다.

샘플 저장소에서 TPU 버전에 대한 RayCluster 매니페스트를 검토합니다.

Trillium(v6e)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-torchtpu-cluster
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: v6e-16
    replicas: 1
    numOfHosts: 4
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          resources:
            limits:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 100Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4

Ironwood (TPU7x)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-torchtpu-cluster
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: tpu7x-16
    replicas: 1
    numOfHosts: 2
    resources:
      # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
      TPU: "8"
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          resources:
            limits:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 128Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu7x
          cloud.google.com/gke-tpu-topology: 2x2x2

Ironwood (TPU7x)에서는 각 물리적 칩에 2개의 칩렛이 포함되어 있어 4칩 호스트당 8개의 논리적 TPU 코어를 제공합니다. 매니페스트는 resources 섹션에서 TPU 필드의 값을 "8"로 설정하고 RAY_TPU_RESOURCE_PER_CHIP 환경 변수를 "2"로 설정하여 호스트당 8개의 논리적 TPU 리소스를 공지하고 예약합니다.

매니페스트를 적용하여 RayCluster 커스텀 리소스를 만듭니다.

Trillium(v6e)

envsubst < ray-cluster.tpu-v6e-multihost.yaml | kubectl apply -f -

Ironwood (TPU7x)

envsubst < ray-cluster.tpu-v7x-multihost.yaml | kubectl apply -f -

RayCluster 헤드 포드와 모든 TPU 워커 포드가 Running 상태가 될 때까지 기다립니다.

kubectl get pods -l ray.io/cluster=ray-torchtpu-cluster

출력은 다음과 비슷합니다.

Trillium(v6e)

NAME                                         READY   STATUS    RESTARTS   AGE
ray-torchtpu-cluster-head-pdpk6              1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-brwpc     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-dx4w7     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-h4d4w     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-t2g24     1/1     Running   0          2m30s

Ironwood (TPU7x)

NAME                                         READY   STATUS    RESTARTS   AGE
ray-torchtpu-cluster-head-k9x2m              1/1     Running   0          2m30s
ray-torchtpu-cluster-tpu7x-16-worker-7cwpq   1/1     Running   0          2m30s
ray-torchtpu-cluster-tpu7x-16-worker-m98r9   1/1     Running   0          2m30s

작업자 포드 중 하나의 환경 변수를 검사하여 KubeRay TPU 웹훅이 멀티 호스트 TPU 조정 및 TorchTPU 토폴로지 변수를 삽입했는지 확인합니다.

kubectl get pods -l ray.io/node-type=worker \
    -o=jsonpath='{range .items[0].spec.containers[0].env[*]}{.name}={.value}{"\n"}{end}' \
    | grep -E "TPU_WORKER_ID|TPU_WORKER_HOSTNAMES|TPU_NAME|TORCH_TPU_TOPOLOGY|TORCH_TPU_SLICEBUILDER_ADDRESSES"

출력은 다음과 비슷합니다.

Trillium(v6e)

TPU_WORKER_HOSTNAMES=v6e-16-0-0.ray-torchtpu-cluster-headless,v6e-16-0-1.ray-torchtpu-cluster-headless,v6e-16-0-2.ray-torchtpu-cluster-headless,v6e-16-0-3.ray-torchtpu-cluster-headless
TORCH_TPU_TOPOLOGY=4,4,1
TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-torchtpu-cluster-headless:8471,v6e-16-0-0.ray-torchtpu-cluster-headless:8472,v6e-16-0-0.ray-torchtpu-cluster-headless:8473,v6e-16-0-0.ray-torchtpu-cluster-headless:8474,v6e-16-0-1.ray-torchtpu-cluster-headless:8471,v6e-16-0-1.ray-torchtpu-cluster-headless:8472,v6e-16-0-1.ray-torchtpu-cluster-headless:8473,v6e-16-0-1.ray-torchtpu-cluster-headless:8474,v6e-16-0-2.ray-torchtpu-cluster-headless:8471,v6e-16-0-2.ray-torchtpu-cluster-headless:8472,v6e-16-0-2.ray-torchtpu-cluster-headless:8473,v6e-16-0-2.ray-torchtpu-cluster-headless:8474,v6e-16-0-3.ray-torchtpu-cluster-headless:8471,v6e-16-0-3.ray-torchtpu-cluster-headless:8472,v6e-16-0-3.ray-torchtpu-cluster-headless:8473,v6e-16-0-3.ray-torchtpu-cluster-headless:8474
TPU_WORKER_ID=3
TPU_NAME=v6e-16-0

Ironwood (TPU7x)

TPU_WORKER_HOSTNAMES=tpu7x-16-0-0.ray-torchtpu-cluster-headless,tpu7x-16-0-1.ray-torchtpu-cluster-headless
TORCH_TPU_TOPOLOGY=2,2,2,2
TORCH_TPU_SLICEBUILDER_ADDRESSES=tpu7x-16-0-0.ray-torchtpu-cluster-headless:8471,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8472,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8473,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8474,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8475,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8476,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8477,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8478,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8471,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8472,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8473,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8474,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8475,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8476,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8477,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8478
TPU_WORKER_ID=1
TPU_NAME=tpu7x-16-0

Ray 헤드 서비스에 대한 포트 전달 세션을 설정합니다.
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
Ray 대시보드의 클러스터 탭에서 TPU 워커 노드와 자동 감지된 ray.io/tpu-* 라벨을 검사할 수 있습니다.

Ray Core로 분산 텐서 동시 로드 실행

Ray Core는 ray.util.tpu 모듈의 slice_placement_group 함수를 제공하여 손상되지 않은 멀티 호스트 TPU 슬라이스를 예약하고 PlacementGroupSchedulingStrategy 클래스를 제공하여 슬라이스 전반에서 논리적 TPU 기기당 하나의 Ray 행위자를 예약합니다.

샘플 저장소에서 ray_core_collectives.py 스크립트를 검토합니다.

import argparse
import os

import ray
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from ray.util.tpu import (
    get_tpu_version_from_type,
    slice_placement_group,
)
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    ColwiseParallel,
    RowwiseParallel,
    parallelize_module,
)
import torch.nn as nn
import torch_tpu  # noqa: F401

RANDOM_SEED = 42


class FeedForward(nn.Module):
    """Two-layer feed-forward block for tensor-parallel sharding across a TPU slice."""

    def __init__(self, hidden_dim: int):
        super().__init__()
        self.w1 = nn.Linear(hidden_dim, 4 * hidden_dim, bias=False)
        self.act = nn.SiLU()
        self.w2 = nn.Linear(4 * hidden_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.w2(self.act(self.w1(x)))


@ray.remote
class TorchTPUSliceWorker:
    """Ray actor bound to a single logical TPU device within a slice placement group."""

    def __init__(self, devices_per_host: int):
        self.node_ip = ray.util.get_node_ip_address()
        hosts = os.environ["TPU_WORKER_HOSTNAMES"].split(",")
        self.worker_id = int(os.environ.get("TPU_WORKER_ID", "0"))
        self.local_rank = int(
            ray.get_runtime_context().get_accelerator_ids()["TPU"][0]
        )
        self.rank = self.worker_id * devices_per_host + self.local_rank
        # Set per-process PyTorch env:// variables from the webhook's host-level
        # TPU metadata and Ray's assigned local device ID (TorchTrainer sets these
        # automatically in the Ray Train section below).
        os.environ.update(
            {
                "MASTER_ADDR": hosts[0],
                "MASTER_PORT": "29500",
                "WORLD_SIZE": str(len(hosts) * devices_per_host),
                "RANK": str(self.rank),
                "LOCAL_RANK": str(self.local_rank),
                "LOCAL_WORLD_SIZE": str(devices_per_host),
            }
        )

    def run_tensor_parallel_step(
        self,
        batch_size: int = 64,
        hidden_dim: int = 128,
    ) -> dict:
        dist.init_process_group(backend="tpu_dist")
        world_size = dist.get_world_size()
        device = torch.device("tpu")

        # 1. Compute single-device reference output on CPU with identical seed.
        torch.manual_seed(RANDOM_SEED)
        ref_model = FeedForward(hidden_dim).eval()
        inputs_cpu = torch.randn(batch_size, hidden_dim)
        with torch.no_grad():
            ref_output = ref_model(inputs_cpu)

        # 2. Shard the model across all TPU devices in the slice with DTensor TP.
        torch.manual_seed(RANDOM_SEED)
        tp_model = FeedForward(hidden_dim).to(device).eval()
        mesh = init_device_mesh("tpu", (world_size,))
        tp_model = parallelize_module(
            tp_model,
            mesh,
            {"w1": ColwiseParallel(), "w2": RowwiseParallel()},
        )

        with torch.no_grad():
            tp_output = tp_model(inputs_cpu.to(device)).cpu()

        torch.testing.assert_close(tp_output, ref_output, rtol=1e-3, atol=1e-2)
        max_abs_diff = float((tp_output - ref_output).abs().max().item())

        # 3. Verify direct ICI all_reduce across all ranks in the slice.
        rank_tensor = torch.full(
            (4,), float(self.rank + 1), dtype=torch.float32, device=device
        )
        dist.all_reduce(rank_tensor, op=dist.ReduceOp.SUM)
        all_reduce_sum = rank_tensor.cpu().tolist()

        dist.destroy_process_group()
        return {
            "rank": self.rank,
            "local_rank": self.local_rank,
            "worker_id": self.worker_id,
            "node_ip": self.node_ip,
            "all_reduce_sum": all_reduce_sum,
            "max_abs_diff": max_abs_diff,
        }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run distributed TorchTPU tensor parallelism across a Ray TPU slice."
    )
    parser.add_argument(
        "--topology",
        type=str,
        default="4x4",
        help="TPU physical topology.",
    )
    parser.add_argument(
        "--accelerator-type",
        type=str,
        default="TPU-V6E",
        help="Ray TPU accelerator type (TPU-V6E, TPU-V7X, etc.).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ray.init()

    slice_handle = slice_placement_group(
        topology=args.topology,
        accelerator_version=get_tpu_version_from_type(args.accelerator_type),
        num_slices=1,
        resources_per_bundle={"TPU": 1},
    )
    ray.get(slice_handle.placement_group.ready())

    devices_per_host = slice_handle.num_bundles // slice_handle.num_hosts
    workers = [
        TorchTPUSliceWorker.options(
            resources={"TPU": 1},
            scheduling_strategy=PlacementGroupSchedulingStrategy(
                placement_group=slice_handle.placement_group,
                placement_group_bundle_index=bundle_idx,
            ),
        ).remote(devices_per_host)
        for bundle_idx in range(slice_handle.num_bundles)
    ]
    results = sorted(
        ray.get([w.run_tensor_parallel_step.remote() for w in workers]),
        key=lambda r: r["rank"],
    )

    expected_sum = sum(range(1, slice_handle.num_bundles + 1))
    for r in results:
        print(
            f"Rank {r['rank']:2d} (worker_id={r['worker_id']}, "
            f"local_rank={r['local_rank']}, ip={r['node_ip']}): "
            f"all_reduce_sum={r['all_reduce_sum']}"
        )
        assert r["all_reduce_sum"] == [float(expected_sum)] * 4

    print(
        f"All {len(results)} TorchTPU ranks verified DTensor TP parity "
        f"(max_abs_diff={results[0]['max_abs_diff']:.6f}) and "
        f"all_reduce sum = {expected_sum:.1f}"
    )
    for w in workers:
        ray.kill(w)
    slice_handle.shutdown()


if __name__ == "__main__":
    main()

스크립트는 TPU 슬라이스에서 분산 텐서 동시 로드를 3단계로 실행합니다.

slice_placement_group 함수는 16개 기기 TPU 슬라이스를 예약하고 각 TPU 기기에 TorchTPUSliceWorker 액터 하나를 예약합니다.
각 액터는 호스트 색인과 로컬 기기 ID에서 전역 순위를 계산한 후 tpu_dist 프로세스 그룹을 초기화합니다.
init_device_mesh 및 parallelize_module 함수는 16개의 모든 TPU 기기에서 피드포워드 레이어를 샤딩하고, 단일 기기 CPU 참조와 비교하여 출력을 확인하고, ICI를 통해 all_reduce 집합을 실행합니다.

TPU 버전에 대한 Ray Core 작업을 제출합니다.

Trillium(v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_core_collectives.py --topology=4x4 --accelerator-type=TPU-V6E

Ironwood (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_core_collectives.py --topology=2x2x2 --accelerator-type=TPU-V7X

출력은 다음과 비슷합니다.

Trillium(v6e)

Rank  0 (worker_id=0, local_rank=0, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  1 (worker_id=0, local_rank=1, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  2 (worker_id=0, local_rank=2, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  3 (worker_id=0, local_rank=3, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  4 (worker_id=1, local_rank=0, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  5 (worker_id=1, local_rank=1, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  6 (worker_id=1, local_rank=2, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  7 (worker_id=1, local_rank=3, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  8 (worker_id=2, local_rank=0, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  9 (worker_id=2, local_rank=1, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 10 (worker_id=2, local_rank=2, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 11 (worker_id=2, local_rank=3, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 12 (worker_id=3, local_rank=0, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 13 (worker_id=3, local_rank=1, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 14 (worker_id=3, local_rank=2, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 15 (worker_id=3, local_rank=3, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
All 16 TorchTPU ranks verified DTensor TP parity (max_abs_diff=0.002470) and all_reduce sum = 136.0
------------------------------------------
Job 'raysubmit_HFTkKrjSnfza3ZFD' succeeded
------------------------------------------

Ironwood (TPU7x)

Rank  0 (worker_id=0, local_rank=0, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  1 (worker_id=0, local_rank=1, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  2 (worker_id=0, local_rank=2, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  3 (worker_id=0, local_rank=3, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  4 (worker_id=0, local_rank=4, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  5 (worker_id=0, local_rank=5, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  6 (worker_id=0, local_rank=6, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  7 (worker_id=0, local_rank=7, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  8 (worker_id=1, local_rank=0, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  9 (worker_id=1, local_rank=1, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 10 (worker_id=1, local_rank=2, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 11 (worker_id=1, local_rank=3, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 12 (worker_id=1, local_rank=4, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 13 (worker_id=1, local_rank=5, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 14 (worker_id=1, local_rank=6, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 15 (worker_id=1, local_rank=7, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
All 16 TorchTPU ranks verified DTensor TP parity (max_abs_diff=0.002470) and all_reduce sum = 136.0
------------------------------------------
Job 'raysubmit_MrJbAe3hvARBcBN2' succeeded
------------------------------------------

Ray Train 및 PyTorch FSDP2를 사용하여 Gemma 4 모델 파인 튜닝

Ray Train TorchTrainer 클래스는 멀티 호스트 TPU 슬라이스에서 TPU 슬라이스 예약, 작업자 배치, torch.distributed 초기화를 자동화합니다.

샘플 저장소에서 databricks/databricks-dolly-15k 명령어 데이터 세트에서 google/gemma-4-E2B-it 모델을 미세 조정하는 ray_train_llm_finetune.py 스크립트를 검토합니다.

import argparse
import os
import tempfile
import time

from datasets import load_dataset
import ray
import ray.train
from ray.train import (
    Checkpoint,
    CheckpointConfig,
    RunConfig,
    ScalingConfig,
)
from ray.train.torch import TorchConfig, TorchTrainer
import torch
from torch.distributed import ReduceOp, all_reduce
from torch.distributed.checkpoint.state_dict import (
    StateDictOptions,
    get_model_state_dict,
)
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
from torch.utils.data import DataLoader
import torch_tpu  # noqa: F401
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
)


def create_dataloader(
    tokenizer: AutoTokenizer,
    batch_size: int,
    dataset_split: str = "train[:1024]",
    max_seq_len: int = 512,
) -> DataLoader:
    """Loads and tokenizes the Dolly 15k dataset for causal LM fine-tuning."""
    dataset = load_dataset(
        "databricks/databricks-dolly-15k", split=dataset_split
    )
    tokenized = dataset.map(
        lambda x: tokenizer(
            f"Instruction: {x['instruction']}\nResponse: {x['response']}",
            truncation=True,
            max_length=max_seq_len,
            padding="max_length",
        ),
        remove_columns=dataset.column_names,
    )
    dataloader = DataLoader(
        tokenized,
        batch_size=batch_size,
        shuffle=True,
        drop_last=True,
        collate_fn=DataCollatorForLanguageModeling(
            tokenizer=tokenizer, mlm=False
        ),
    )
    return ray.train.torch.prepare_data_loader(dataloader)


def train_func(config: dict) -> None:
    """Per-worker distributed FSDP2 training loop executed on each logical TPU core."""
    ctx = ray.train.get_context()
    world_rank = ctx.get_world_rank()
    world_size = ctx.get_world_size()
    device = ray.train.torch.get_device()

    model_id = config["model_id"]
    dataset_split = config["dataset_split"]
    batch_size = config["batch_size_per_device"]
    num_epochs = config["num_epochs"]
    max_seq_len = 512

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    dataloader = create_dataloader(
        tokenizer=tokenizer,
        batch_size=batch_size,
        dataset_split=dataset_split,
        max_seq_len=max_seq_len,
    )

    torch.manual_seed(42)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        dtype=torch.bfloat16,
        attn_implementation="eager",
        vision_config=None,
        audio_config=None,
    ).to(device)
    model.config.get_text_config().use_cache = False

    mesh = init_device_mesh("tpu", (world_size,))
    mp_policy = MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    )
    text_model = model.model.language_model
    text_model.embed_tokens_per_layer.requires_grad_(False)
    for layer in text_model.layers:
        fully_shard(layer.mlp, mesh=mesh, mp_policy=mp_policy)
    fully_shard(
        model,
        mesh=mesh,
        mp_policy=mp_policy,
        ignored_params=set(text_model.embed_tokens_per_layer.parameters()),
    )

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=2e-5, weight_decay=0.01
    )
    global_batch_size = batch_size * world_size

    for epoch in range(num_epochs):
        dataloader.sampler.set_epoch(epoch)
        model.train()
        epoch_start = time.perf_counter()
        running_loss = torch.zeros(1, device=device, dtype=torch.float32)

        for batch in dataloader:
            optimizer.zero_grad()
            outputs = model(
                input_ids=batch["input_ids"],
                labels=batch["labels"],
            )
            outputs.loss.backward()
            optimizer.step()
            running_loss += outputs.loss.detach().float()

        all_reduce(running_loss, op=ReduceOp.AVG)
        avg_loss = running_loss.item() / len(dataloader)
        epoch_time = time.perf_counter() - epoch_start
        tokens_per_sec = (
            len(dataloader) * global_batch_size * max_seq_len
        ) / epoch_time

        metrics = {
            "epoch": epoch + 1,
            "loss": round(avg_loss, 4),
            "epoch_time_s": round(epoch_time, 2),
            "tokens_per_s": round(tokens_per_sec, 1),
        }
        if world_rank == 0:
            print(
                f"Epoch {epoch + 1}/{num_epochs} | loss={metrics['loss']:.4f} | "
                f"epoch_time={metrics['epoch_time_s']:.2f}s | "
                f"tokens/s={metrics['tokens_per_s']:.1f}"
            )

        full_state_dict = get_model_state_dict(
            model,
            options=StateDictOptions(
                full_state_dict=True,
                cpu_offload=True,
                ignore_frozen_params=True,
            ),
        )
        if world_rank == 0:
            with tempfile.TemporaryDirectory() as tmpdir:
                torch.save(full_state_dict, os.path.join(tmpdir, "model.pt"))
                ray.train.report(
                    metrics, checkpoint=Checkpoint.from_directory(tmpdir)
                )
        else:
            ray.train.report(metrics)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fine-tune a Hugging Face CausalLM on TPUs with Ray Train and TorchTPU FSDP2."
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="google/gemma-4-E2B-it",
        help="Hugging Face model ID to fine-tune.",
    )
    parser.add_argument(
        "--topology",
        type=str,
        default="4x4",
        help="TPU slice topology.",
    )
    parser.add_argument(
        "--accelerator-type",
        type=str,
        default="TPU-V6E",
        help="Ray TPU accelerator type (TPU-V6E, TPU-V7X, etc.).",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=16,
        help="Total TPU worker processes across the slice.",
    )
    parser.add_argument(
        "--dataset-split",
        type=str,
        default="train[:1024]",
        help="Hugging Face dataset split to load.",
    )
    parser.add_argument(
        "--num-epochs",
        type=int,
        default=3,
        help="Number of fine-tuning epochs.",
    )
    parser.add_argument(
        "--batch-size-per-device",
        type=int,
        default=4,
        help="Per-device training batch size.",
    )
    parser.add_argument(
        "--storage-path",
        type=str,
        default=os.environ.get("STORAGE_PATH"),
        required="STORAGE_PATH" not in os.environ,
        help="Cloud Storage or shared file system path for checkpoints.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ray.init()
    run_name = f"torchtpu_gemma4_fsdp2_sft_{int(time.time())}"

    print(
        f"Starting Ray Train FSDP2 SFT job: model={args.model_id}, "
        f"topology={args.topology}, accelerator_type={args.accelerator_type}, "
        f"num_workers={args.num_workers}"
    )

    trainer = TorchTrainer(
        train_loop_per_worker=train_func,
        train_loop_config={
            "model_id": args.model_id,
            "dataset_split": args.dataset_split,
            "num_epochs": args.num_epochs,
            "batch_size_per_device": args.batch_size_per_device,
        },
        torch_config=TorchConfig(backend="tpu_dist"),
        scaling_config=ScalingConfig(
            use_tpu=True,
            num_workers=args.num_workers,
            topology=args.topology,
            accelerator_type=args.accelerator_type,
            resources_per_worker={"TPU": 1},
        ),
        run_config=RunConfig(
            name=run_name,
            storage_path=args.storage_path,
            checkpoint_config=CheckpointConfig(
                num_to_keep=1,
                checkpoint_score_attribute="loss",
                checkpoint_score_order="min",
            ),
        ),
    )

    result = trainer.fit()
    print(f"Training completed! Final metrics: {result.metrics}")
    print(f"Saved checkpoint: {result.checkpoint}")


if __name__ == "__main__":
    main()

스크립트는 다음 구성요소로 TorchTrainer 클래스를 구성합니다.

ScalingConfig 클래스는 학습 작업의 분산 작업자 토폴로지를 구성합니다. use_tpu=True, num_workers=16, topology, accelerator_type 매개변수를 설정하면 16개 기기 TPU 슬라이스가 예약되고 TPU 코어당 작업자 1명이 예약됩니다.
TorchConfig 클래스는 작업자 전반에서 torch.distributed 프로세스 그룹을 구성합니다. backend="tpu_dist" 매개변수를 설정하면 train_func 함수가 실행되기 전에 TPU ICI 네트워크를 통해 16개의 워커가 모두 연결됩니다.
create_dataloader 함수는 안내 데이터 세트를 다운로드하고 토큰화하며, PyTorch DataLoader 객체로 토큰 시퀀스를 일괄 처리하고, ray.train.torch.prepare_data_loader 함수를 호출하여 각 작업자가 할당된 TPU 기기에서 각 전역 배치의 분리된 샤드를 수신하도록 합니다.
train_func 함수는 bfloat16 정밀도로 모델을 로드하고 각 피드포워드 블록과 루트 모델에서 PyTorch FSDP2 fully_shard 함수를 호출하여 16개의 모든 TPU 코어에서 파라미터, 경사, 옵티마이저 상태를 샤딩합니다.
Checkpoint 클래스와 ray.train.report 함수는 각 에포크가 끝날 때 순위 0에서 전체 모델 상태 사전을 수집하고 체크포인트를 Cloud Storage 버킷에 업로드합니다.

TPU 버전에 대한 TorchTrainer FSDP2 파인 튜닝 작업을 제출합니다.

Trillium(v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_train_llm_finetune.py \
        --topology=4x4 \
        --accelerator-type=TPU-V6E \
        --storage-path=gs://${BUCKET_NAME}/ray_train_torchtpu

Ironwood (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_train_llm_finetune.py \
        --topology=2x2x2 \
        --accelerator-type=TPU-V7X \
        --storage-path=gs://${BUCKET_NAME}/ray_train_torchtpu

출력은 다음과 비슷합니다.

Trillium(v6e)

Starting Ray Train FSDP2 SFT job: model=google/gemma-4-E2B-it, topology=4x4, accelerator_type=TPU-V6E, num_workers=16
(TrainController pid=5407) Attempting to start training worker group of size 16 with the following resources: [{'TPU': 1, 'accelerator_type:TPU-V6E': 0.001}] * 16
(RayTrainWorker pid=19733, ip=10.72.3.20) Setting up process group for: env:// [rank=0, world_size=16]
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 1/3 | loss=3.1402 | epoch_time=28.17s | tokens/s=18613.5
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-08-10.762175)
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 2/3 | loss=2.3318 | epoch_time=26.07s | tokens/s=20112.8
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-10-02.548088)
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 3/3 | loss=2.2050 | epoch_time=25.52s | tokens/s=20545.2
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-11-48.701853)
Training completed! Final metrics: {'epoch': 3, 'loss': 2.205, 'epoch_time_s': 25.52, 'tokens_per_s': 20545.2}
Saved checkpoint: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-11-48.701853)
------------------------------------------
Job 'raysubmit_zbm4pnFZitA1YECA' succeeded
------------------------------------------

Ironwood (TPU7x)

Starting Ray Train FSDP2 SFT job: model=google/gemma-4-E2B-it, topology=2x2x2, accelerator_type=TPU-V7X, num_workers=16
(TrainController pid=1162) Attempting to start training worker group of size 16 with the following resources: [{'TPU': 1, 'accelerator_type:TPU-V7X': 0.001}] * 16
(RayTrainWorker pid=257, ip=172.16.26.6) Setting up process group for: env:// [rank=0, world_size=16]
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 1/3 | loss=3.1407 | epoch_time=26.28s | tokens/s=19953.3
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-52-00.026930)
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 2/3 | loss=2.3235 | epoch_time=16.68s | tokens/s=31429.9
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-53-33.533341)
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 3/3 | loss=2.1964 | epoch_time=16.98s | tokens/s=30884.6
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-55-01.508219)
Training completed! Final metrics: {'epoch': 3, 'loss': 2.1964, 'epoch_time_s': 16.98, 'tokens_per_s': 30884.6}
Saved checkpoint: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-55-01.508219)
------------------------------------------
Job 'raysubmit_zSNnCXhHkwKiSmJh' succeeded
------------------------------------------

학습 출력에는 다음이 표시됩니다.

Epoch 1에는 초기 XLA 그래프 컴파일이 포함됩니다. 에포크 2와 3은 /dev/shm/torch_tpu_cache의 캐시된 실행 파일을 재사용하여 TPU Trillium (v6e)에서 초당 20, 000개 이상의 토큰,Ironwood (TPU7x)에서 초당 31,000개의 토큰을 처리합니다.
16개의 FSDP2 워커 순위 전체에서 동기화된 전역 교차 엔트로피 손실은 databricks/databricks-dolly-15k 명령 데이터 세트의 세 가지 미세 조정 에포크에서 에포크 1의 3.14에서 에포크 3의 2.20으로 감소합니다.
각 세대가 끝날 때 순위 0은 전체 모델 상태 사전을 수집하고 체크포인트를 Cloud Storage 버킷에 업로드하여 CheckpointConfig(num_to_keep=1, checkpoint_score_attribute="loss", checkpoint_score_order="min") 객체에 구성된 대로 손실이 가장 낮은 체크포인트를 유지합니다.

워크로드 성능 관찰

Cloud Monitoring에서 TPU TensorCore 사용률, HBM 사용량, Ray 클러스터 측정항목을 모니터링하려면 GKE의 Ray 클러스터에 대한 로그 및 측정항목 수집 및 보기 및 Ray 대시보드에서 TPU 측정항목 보기를 참조하세요.

Ray 대시보드 액세스

기본 제공 Ray 대시보드에서 Ray 작업자, 배치 그룹, 학습 작업의 상태를 검사하려면 다음 단계를 따르세요.

Ray 헤드 서비스에 대한 포트 전달이 아직 실행되고 있지 않으면 백그라운드에서 시작합니다.
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
Cloud Shell에서 웹 미리보기를 클릭하고 포트 8265에서 미리보기를 선택하거나 로컬 브라우저에서 http://localhost:8265를 엽니다.
작업, 클러스터, 측정항목 탭을 클릭하여 워커 로그, 배치 그룹 상태, TPU 노드 라벨, 노드별 TPU 사용률 그래프를 검사합니다.

삭제

이 튜토리얼에서 사용한 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 RayCluster 커스텀 리소스, GKE 클러스터, Cloud Storage 버킷, Artifact Registry 저장소를 삭제합니다.

백그라운드 포트 전달 프로세스를 중지하고 RayCluster 커스텀 리소스와 Kubernetes 보안 비밀을 삭제합니다.

pkill -f "kubectl port-forward" || true
kubectl delete raycluster ray-torchtpu-cluster
kubectl delete secret hf-secret --ignore-not-found

GKE 클러스터를 삭제합니다.

Autopilot

gcloud container clusters delete ${CLUSTER_NAME} \
    --region=${REGION}

표준

gcloud container clusters delete ${CLUSTER_NAME} \
    --zone=${ZONE}

Cloud Storage 버킷 및 체크포인트 콘텐츠를 삭제합니다.
```
gcloud storage rm --recursive gs://${BUCKET_NAME}
```

Artifact Registry 저장소를 삭제합니다.

gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} \
    --location=${REGION}

다음 단계

Cloud TPU의 PyTorch에 대해 자세히 알아보세요.
TorchTPU 문서 살펴보기
Ray 및 TorchTPU를 사용하여 멀티 호스트 TPU에서 LLM을 제공하는 방법 알아보기
DWS 및 Kueue를 사용하여 TPU에서 AI 학습을 최적화하는 방법 알아보기
GKE의 TPU에 대하여를 읽어보세요.
Ray Train 문서 살펴보기

Ray 및 TorchTPU를 사용하여 멀티 호스트 TPU에서 모델 미세 조정 컬렉션을 사용해 정리하기 내 환경설정을 기준으로 콘텐츠를 저장하고 분류하세요.

배경

TPU

TorchTPU

TPU의 KubeRay 및 Ray

목표

비용

시작하기 전에

개발 환경 준비

GKE 클러스터 만들기

Autopilot

표준

Trillium(v6e)

Ironwood (TPU7x)

컨테이너 이미지 빌드

RayCluster 커스텀 리소스 배포

Trillium(v6e)

Ironwood (TPU7x)

Trillium(v6e)

Ironwood (TPU7x)

Trillium(v6e)

Ironwood (TPU7x)

Trillium(v6e)

Ironwood (TPU7x)

Ray Core로 분산 텐서 동시 로드 실행

Trillium(v6e)

Ironwood (TPU7x)

Trillium(v6e)

Ironwood (TPU7x)

Ray Train 및 PyTorch FSDP2를 사용하여 Gemma 4 모델 파인 튜닝

Trillium(v6e)

Ironwood (TPU7x)

Trillium(v6e)

Ironwood (TPU7x)

워크로드 성능 관찰

Ray 대시보드 액세스

삭제

Autopilot

표준

다음 단계

Ray 및 TorchTPU를 사용하여 멀티 호스트 TPU에서 모델 미세 조정