Google menggunakan teknologi AI untuk menerjemahkan konten ke dalam bahasa pilihan Anda. Terjemahan AI mungkin mengandung kesalahan.

Menyesuaikan model di TPU multi-host dengan Ray dan TorchTPU

Autopilot Standard

Tutorial ini menunjukkan cara menjalankan workload PyTorch terdistribusi dan menyesuaikan model Gemma 4 pada slice Cloud TPU multi-host di Google Kubernetes Engine (GKE) menggunakan Ray Core, Ray Train, PyTorch Fully Sharded Data Parallel (FSDP2), dan TorchTPU.

Tutorial ini ditujukan untuk engineer machine learning (ML), admin dan operator platform, serta spesialis data dan AI yang ingin mengorkestrasi workload pelatihan PyTorch terdistribusi di slice TPU multi-host di Kubernetes. Untuk mempelajari lebih lanjut peran umum dan contoh tugas yang dirujuk dalam konten Google Cloud , lihat Peran dan tugas pengguna GKE umum.

Sebelum memulai tutorial ini, pastikan Anda memahami hal-hal berikut:

Latar belakang

Bagian ini menjelaskan teknologi utama yang digunakan dalam tutorial ini.

TPU

Tensor Processing Unit (TPU) adalah akselerator yang dirancang khusus oleh Google untuk workload machine learning. Tutorial ini mendukung TPU Trillium (v6e) dan Ironwood (TPU7x). Dalam slice TPU multi-host, beberapa host TPU berkomunikasi langsung melalui jaringan Inter-Chip Interconnect (ICI) berkecepatan tinggi, yang menyediakan High Bandwidth Memory (HBM) gabungan dan komunikasi kolektif latensi rendah yang diperlukan untuk melatih model yang melebihi kapasitas memori satu chip atau host TPU.

TorchTPU

TorchTPU adalah backend PyTorch open source untuk Cloud TPU, yang dikembangkan melalui kolaborasi antara Google dan Meta. Dengan TorchTPU, Anda dapat menjalankan model PyTorch yang ada di hardware TPU dengan perubahan kode minimal serta menskalakan workload pelatihan dan inferensi dari satu chip ke pod TPU multi-slice melalui compiler Accelerated Linear Algebra (XLA). TorchTPU menggunakan model eksekusi satu proses per perangkat melalui jaringan ICI TPU dan mendukung API PyTorch standar, termasuk eksekusi segera, torch.compile, torch.distributed, DeviceMesh, DTensor, dan Fully Sharded Data Parallel (FSDP2).

KubeRay dan Ray di TPU

KubeRay adalah operator Kubernetes untuk men-deploy dan mengelola aplikasi Ray. Mengaktifkan add-on Ray di GKE akan men-deploy operator KubeRay terkelola dan webhook TPU KubeRay. Saat KubeRay membuat Pod worker yang meminta resource google.com/tpu, webhook akan otomatis mengonfigurasi nama host Pod, Service headless, dan variabel lingkungan TPU yang diperlukan untuk koordinasi multi-host.

Ray mendeteksi topologi TPU di setiap node saat memulai dan menjadwalkan workload terdistribusi di seluruh slice TPU. Anda dapat mengorkestrasi workload TorchTPU terdistribusi di dua level:

Ray Core menyediakan grup penempatan berbasis topologi dan penjadwalan aktor untuk loop eksekusi terdistribusi kustom, seperti komputasi paralel tensor atau pipeline reinforcement learning (RL).
Ray Train menyediakan pelatihan paralel data dan paralel model terdistribusi terkelola, termasuk PyTorch Distributed Data Parallel (DDP), PyTorch FSDP2, dan JAX single program, multiple data (SPMD), dengan checkpointing bawaan dan pemulihan kegagalan otomatis.

Tujuan

Tutorial ini menunjukkan kepada Anda cara melakukan hal berikut:

Buat cluster GKE di Autopilot atau Standard di TPU Trillium (v6e) atau Ironwood (TPU7x).
Bangun image container yang melapisi library pelatihan Ray dan Hugging Face di atas image dasar TorchTPU.
Men-deploy resource kustom RayCluster untuk menyediakan node worker TPU dan head.
Jalankan komputasi tensor-paralel terdistribusi dengan PyTorch DeviceMesh dan DTensor di seluruh slice TPU menggunakan Ray Core.
Menyempurnakan model Gemma 4 dengan PyTorch FSDP2 di seluruh slice TPU multi-host menggunakan Ray Train.

Biaya

Dalam tutorial ini, Anda akan menggunakan komponen Google Cloudyang dapat ditagih berikut:

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga.

Setelah menyelesaikan tugas yang dijelaskan dalam tutorial ini, Anda dapat menghindari penagihan berkelanjutan dengan menghapus resource yang Anda buat. Untuk informasi selengkapnya, lihat Pembersihan.

Sebelum memulai

Sebelum memulai, pastikan Anda telah melakukan tugas berikut:

Aktifkan Artifact Registry API dan Cloud Storage API dan Google Kubernetes Engine API.

Mengaktifkan API

Untuk menggunakan Google Cloud CLI untuk tugas ini, instal lalu lakukan inisialisasi gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan perintah gcloud components update. Versi gcloud CLI yang lebih lama mungkin tidak mendukung menjalankan perintah dalam dokumen ini.
Catatan: Untuk penginstalan gcloud CLI yang ada, pastikan untuk menyetel properti compute/region. Jika Anda menggunakan cluster zona, tetapkan compute/zone. Dengan menetapkan lokasi default, Anda dapat menghindari error di gcloud CLI seperti berikut: One of [--zone, --region] must be supplied: Please specify location. Anda mungkin perlu menentukan lokasi dalam perintah tertentu jika lokasi cluster Anda berbeda dengan default yang Anda tetapkan.

Pastikan project Anda memiliki kuota yang cukup untuk kapasitas TPU Trillium (v6e) atau Ironwood (TPU7x) di zona yang Anda pilih. Untuk mengetahui informasi selengkapnya, lihat Kuota Cloud TPU.
Pastikan akun Anda memiliki peran IAM berikut:
- Kubernetes Engine Admin (roles/container.admin)
- Artifact Registry Administrator (roles/artifactregistry.admin)
- Storage Admin (roles/storage.admin)

Menyiapkan lingkungan Anda

Dalam tutorial ini, Anda akan menggunakan Cloud Shell untuk mengelola resource yang dihosting di Google Cloud. Cloud Shell telah diinstal dengan software yang Anda perlukan untuk tutorial ini, termasuk kubectl, docker, python3, dan gcloud CLI.

Untuk menyiapkan lingkungan Anda dengan Cloud Shell, ikuti langkah-langkah berikut:

Di konsol Google Cloud , luncurkan sesi Cloud Shell dengan mengklik Activate Cloud Shell di konsolGoogle Cloud . Tindakan ini akan meluncurkan sesi di panel bawah konsol Google Cloud .

Buat clone repositori kubernetes-engine-samples dan ubah ke direktori contoh untuk tutorial ini:

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

Buat dan aktifkan lingkungan virtual Python lokal dengan Ray CLI yang terinstal:

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

Untuk mendownload bobot model Gemma 4, setujui perjanjian lisensi Google di Hugging Face:
1. Buka halaman model Gemma 4 di Hugging Face.
2. Login dan setujui persyaratan lisensi dengan mengklik Agree and access repository.
3. Di setelan akun Hugging Face Anda, buat token akses pengguna dengan izin Read.
Tetapkan variabel lingkungan untuk project, cluster, bucket Cloud Storage, dan token Hugging Face Anda Google Cloud :
```
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}
```
Ganti kode berikut:
- PROJECT_ID: Project ID Google Cloud Anda.
- CLUSTER_NAME: nama cluster GKE Anda, seperti ray-torchtpu-cluster.
- REGION: region tempat repositori Artifact Registry, bucket Cloud Storage, dan kapasitas TPU Anda berada.
- ZONE: zona tempat kapasitas TPU Anda tersedia. Untuk mengetahui region dan zona yang tersedia berdasarkan versi TPU, lihat Region dan zona TPU.
- ARTIFACT_REGISTRY: nama repositori Docker Artifact Registry Anda, seperti ray-torchtpu.
- BUCKET_NAME: nama bucket Cloud Storage Anda untuk menyimpan checkpoint Ray Train, seperti PROJECT_ID-ray-torchtpu.
- HF_TOKEN: token akses Hugging Face Anda.

Membuat cluster GKE

Anda dapat menjalankan workload pelatihan TorchTPU terdistribusi di cluster GKE Autopilot atau Standard dengan add-on operator Ray terkelola yang diaktifkan.

Buat cluster GKE dan konfigurasikan kredensial untuk alat command line kubectl:
Autopilot
1. Buat cluster GKE Autopilot regional dengan add-on operator Ray yang diaktifkan:
  
  gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator
  
  Saat Anda men-deploy resource kustom RayCluster di langkah berikutnya, GKE Autopilot akan otomatis menyediakan node head CPU dan slice TPU multi-host yang sesuai dengan permintaan resource Pod dan pemilih node Anda.
2. Konfigurasi alat command line kubectl untuk berkomunikasi dengan cluster Anda:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
Standar
1. Buat cluster GKE Standard zona dengan node pool CPU untuk menghosting Pod head Ray, Streaming image, GKE Dataplane V2, dan Workload Identity Federation untuk GKE yang diaktifkan:
  
  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
  Catatan: Jika Anda mengelola sendiri KubeRay open source, bukan menggunakan add-on RayOperator GKE, instal KubeRay dan webhook TPU KubeRay versi 1.4.0 atau yang lebih baru untuk injeksi variabel lingkungan TorchTPU.
2. Konfigurasi alat command line kubectl untuk berkomunikasi dengan cluster Anda:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
3. Buat node pool slice TPU multi-host 16 perangkat dengan topologi 4x4 di empat host pada TPU Trillium (v6e) atau topologi 2x2x2 di dua host pada Ironwood (TPU7x):
  
  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
  
  Tips: Untuk menggunakan pemesanan TPU tertentu, tambahkan flag --reservation-affinity=specific dan --reservation=RESERVATION_NAME ke perintah node-pools create, dengan RESERVATION_NAME adalah nama pemesanan TPU Anda. Untuk menggunakan Spot VM, tambahkan flag --spot.

Buat bucket Cloud Storage untuk menyimpan checkpoint Ray Train bersama di seluruh node TPU multi-host dan berikan akses resource Kubernetes ServiceAccount default ke bucket:

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"

Buat Secret Kubernetes dengan token akses Hugging Face Anda agar Pod Ray dapat mendownload bobot Gemma 4 yang dibatasi:
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

Membangun image container

Bangun image container yang menginstal library pelatihan Ray dan Hugging Face di atas image dasar TorchTPU bawaan, yang mencakup TorchTPU dan library libtpu.

Buat repositori Docker Artifact Registry dan konfigurasi autentikasi Docker:

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

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

Periksa Dockerfile di direktori contoh:

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

Pelatihan single-slice dengan class TorchTrainer Ray dan penjadwalan per chiplet di Ironwood (TPU7x) memerlukan Ray versi 2.58.0 atau yang lebih baru, dan penyesuaian Gemma 4 memerlukan library transformers versi 5.5.0 atau yang lebih baru.

Bangun dan kirim image ke Artifact Registry:

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

Men-deploy resource kustom RayCluster

Men-deploy resource kustom RayCluster yang menyediakan Pod head CPU, volume /dev/shm dalam memori untuk penyimpanan objek memori bersama Ray, dan grup worker TPU multi-host.

Tinjau manifes RayCluster untuk versi TPU Anda di repositori contoh:

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

Di Ironwood (TPU7x), setiap chip fisik berisi dua chiplet, yang menyediakan delapan core TPU logis per host empat chip. Manifes menetapkan nilai kolom TPU ke "8" di bagian resources dan menetapkan variabel lingkungan RAY_TPU_RESOURCE_PER_CHIP ke "2" untuk memberitahukan dan menjadwalkan delapan resource TPU logis per host.

Terapkan manifes untuk membuat resource kustom 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 -

Tunggu hingga Pod head RayCluster dan semua Pod worker TPU mencapai status Running:

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

Outputnya mirip dengan hal berikut ini:

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

Periksa variabel lingkungan di salah satu Pod worker untuk mengonfirmasi bahwa webhook TPU KubeRay memasukkan koordinasi TPU multi-host dan variabel topologi 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"

Outputnya mirip dengan hal berikut ini:

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

Buat sesi penerusan port ke layanan head Ray:
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
Anda dapat memeriksa node worker TPU dan label ray.io/tpu-* yang terdeteksi otomatis di tab Cluster di Ray Dashboard.

Menjalankan paralelisme tensor terdistribusi dengan Ray Core

Ray Core menyediakan fungsi slice_placement_group dalam modul ray.util.tpu untuk memesan slice TPU multi-host yang utuh dan class PlacementGroupSchedulingStrategy untuk menjadwalkan satu aktor Ray per perangkat TPU logis di seluruh slice.

Tinjau skrip ray_core_collectives.py di repositori contoh:

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()

Skrip ini menjalankan paralelisme tensor terdistribusi di seluruh slice TPU dalam tiga langkah:

Fungsi slice_placement_group mencadangkan slice TPU 16 perangkat dan menjadwalkan satu aktor TorchTPUSliceWorker di setiap perangkat TPU.
Setiap aktor menghitung peringkat globalnya dari indeks host dan ID perangkat lokal, lalu menginisialisasi grup proses tpu_dist.
Fungsi init_device_mesh dan parallelize_module melakukan sharding pada lapisan feed-forward di 16 perangkat TPU, memverifikasi output terhadap referensi CPU perangkat tunggal, dan menjalankan kolektif all_reduce melalui ICI.

Kirim tugas Ray Core untuk versi TPU Anda:

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

Outputnya mirip dengan hal berikut ini:

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
------------------------------------------

Menyesuaikan model Gemma 4 dengan Ray Train dan PyTorch FSDP2

Class TorchTrainer Ray Train mengotomatiskan pemesanan slice TPU, penempatan worker, dan inisialisasi torch.distributed di seluruh slice TPU multi-host.

Tinjau skrip ray_train_llm_finetune.py di repositori contoh, yang melakukan fine-tuning pada model google/gemma-4-E2B-it di set data instruksi databricks/databricks-dolly-15k:

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()

Skrip ini mengonfigurasi class TorchTrainer dengan komponen berikut:

Class ScalingConfig mengonfigurasi topologi worker terdistribusi untuk tugas pelatihan. Menetapkan parameter use_tpu=True, num_workers=16, topology, dan accelerator_type akan mereservasi slice TPU 16 perangkat dan menjadwalkan satu worker per TPU core.
Class TorchConfig mengonfigurasi grup proses torch.distributed di seluruh worker. Menetapkan parameter backend="tpu_dist" akan menghubungkan ke-16 worker melalui jaringan ICI TPU sebelum fungsi train_func berjalan.
Fungsi create_dataloader mendownload dan melakukan tokenisasi pada set data petunjuk, mengelompokkan urutan token dengan objek DataLoader PyTorch, dan memanggil fungsi ray.train.torch.prepare_data_loader sehingga setiap worker menerima shard yang tidak saling tumpang tindih dari setiap batch global di perangkat TPU yang ditetapkan.
Fungsi train_func memuat model dalam presisi bfloat16 dan memanggil fungsi fully_shard PyTorch FSDP2 pada setiap blok feed-forward dan model root untuk melakukan sharding parameter, gradien, dan status pengoptimal di 16 core TPU.
Class Checkpoint dan fungsi ray.train.report mengumpulkan kamus status model lengkap pada peringkat 0 di akhir setiap epoch dan mengupload checkpoint ke bucket Cloud Storage Anda.

Kirim tugas fine-tuning FSDP2 TorchTrainer untuk versi TPU Anda:

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

Outputnya mirip dengan hal berikut ini:

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
------------------------------------------

Output pelatihan menampilkan hal berikut:

Epoch 1 mencakup kompilasi grafik XLA awal. Epoch 2 dan 3 menggunakan kembali file yang dapat dieksekusi yang di-cache dari /dev/shm/torch_tpu_cache, sehingga mencapai lebih dari 20.000 token per detik di TPU Trillium (v6e) dan 31.000 token per detik di Ironwood (TPU7x).
Kerugian entropi silang global yang disinkronkan di semua 16 peringkat worker FSDP2 menurun dari 3,14 di Epoch 1 menjadi 2,20 di Epoch 3 di tiga epoch fine-tuning pada set data instruksi databricks/databricks-dolly-15k.
Di akhir setiap epoch, peringkat 0 mengumpulkan kamus status model lengkap dan mengupload checkpoint ke bucket Cloud Storage Anda, dengan mempertahankan checkpoint dengan loss terendah seperti yang dikonfigurasi dalam objek CheckpointConfig(num_to_keep=1, checkpoint_score_attribute="loss", checkpoint_score_order="min").

Mengamati performa workload

Untuk memantau pemakaian TPU TensorCore, penggunaan HBM, dan metrik cluster Ray di Cloud Monitoring, lihat Mengumpulkan dan melihat log serta metrik untuk cluster Ray di GKE dan Melihat metrik TPU di Dasbor Ray.

Mengakses Dasbor Ray

Untuk memeriksa status aktor Ray, grup penempatan, dan tugas pelatihan di Dasbor Ray bawaan, ikuti langkah-langkah berikut:

Jika penerusan port ke layanan head Ray belum berjalan, mulai di latar belakang:

kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &

Di Cloud Shell, klik Web Preview dan pilih Preview on port 8265, atau buka http://localhost:8265 di browser lokal Anda.
Klik tab Jobs, Cluster, dan Metrics untuk memeriksa log worker, status grup penempatan, label node TPU, dan grafik TPU Utilization by Node.

Pembersihan

Agar akun Google Cloud Anda tidak dikenai biaya untuk resource yang digunakan dalam tutorial ini, hapus resource kustom RayCluster, cluster GKE, bucket Cloud Storage, dan repositori Artifact Registry:

Hentikan proses penerusan port latar belakang dan hapus resource kustom RayCluster dan Secret Kubernetes:

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

Hapus cluster GKE:

Autopilot

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

Standar

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

Hapus bucket Cloud Storage dan konten checkpoint-nya:

gcloud storage rm --recursive gs://${BUCKET_NAME}

Hapus repositori Artifact Registry:

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

Langkah berikutnya

Pelajari lebih lanjut PyTorch di Cloud TPU.
Pelajari dokumentasi TorchTPU.
Pelajari cara menyajikan LLM di TPU multi-host dengan Ray dan TorchTPU.
Pelajari cara mengoptimalkan pelatihan AI di TPU dengan DWS dan Kueue.
Baca Tentang TPU di GKE.
Pelajari dokumentasi Ray Train.

Menyesuaikan model di TPU multi-host dengan Ray dan TorchTPU Tetap teratur dengan koleksi Simpan dan kategorikan konten berdasarkan preferensi Anda.

Latar belakang

TPU

TorchTPU

KubeRay dan Ray di TPU

Tujuan

Biaya

Sebelum memulai

Menyiapkan lingkungan Anda

Membuat cluster GKE

Autopilot

Standar

Trillium (v6e)

Ironwood (TPU7x)

Membangun image container

Men-deploy resource kustom RayCluster

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Menjalankan paralelisme tensor terdistribusi dengan Ray Core

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Menyesuaikan model Gemma 4 dengan Ray Train dan PyTorch FSDP2

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Mengamati performa workload

Mengakses Dasbor Ray

Pembersihan

Autopilot

Standar

Langkah berikutnya

Menyesuaikan model di TPU multi-host dengan Ray dan TorchTPU