Modelle auf TPUs mit mehreren Hosts mit Ray und TorchTPU optimieren

Autopilot Standard

In dieser Anleitung erfahren Sie, wie Sie verteilte PyTorch-Arbeitslasten ausführen und ein Gemma 4-Modell auf Cloud TPU-Slices mit mehreren Hosts in Google Kubernetes Engine (GKE) optimieren. Dazu verwenden Sie Ray Core, Ray Train, PyTorch Fully Sharded Data Parallel (FSDP2) und TorchTPU.

Diese Anleitung richtet sich an Entwickler für maschinelles Lernen (ML), Plattformadministratoren und -operatoren sowie Daten- und KI-Spezialisten, die verteilte PyTorch-Trainingsarbeitslasten auf TPU-Slices mit mehreren Hosts in Kubernetes orchestrieren möchten. Weitere Informationen zu gängigen Rollen und Beispielaufgaben, auf die wir in Google Cloud -Inhalten verweisen, finden Sie unter Häufig verwendete GKE-Nutzerrollen und -Aufgaben.

Bevor Sie mit dieser Anleitung beginnen, sollten Sie mit Folgendem vertraut sein:

• Informationen zu GKE-Betriebsmodi
• TPUs in GKE
• TPU-Slice-Knotenpools mit mehreren Hosts
• Ray in GKE

Hintergrund

In diesem Abschnitt werden die in dieser Anleitung verwendeten Schlüsseltechnologien beschrieben.

TPUs

Tensor Processing Units (TPUs) sind von Google speziell entwickelte Beschleuniger für Arbeitslasten des maschinellen Lernens. Diese Anleitung unterstützt sowohl TPU Trillium (v6e) als auch Ironwood (TPU7x). In einem TPU-Slice mit mehreren Hosts kommunizieren mehrere TPU-Hosts direkt über ein Hochgeschwindigkeits-ICI-Netzwerk (Inter-Chip Interconnect). Dadurch werden der aggregierte High Bandwidth Memory (HBM) und die kollektive Kommunikation mit niedriger Latenz bereitgestellt, die zum Trainieren von Modellen erforderlich sind, die die Speicherkapazität eines einzelnen Chips oder TPU-Hosts überschreiten.

TorchTPU

TorchTPU ist ein Open-Source-PyTorch-Backend für Cloud TPU, das in Zusammenarbeit von Google und Meta entwickelt wurde. Mit TorchTPU können Sie vorhandene PyTorch-Modelle mit minimalen Codeänderungen auf TPU-Hardware ausführen und Trainings- und Inferenzarbeitslasten über den XLA-Compiler (Accelerated Linear Algebra) von einem einzelnen Chip auf TPU-Pods mit mehreren Slices skalieren. TorchTPU verwendet ein Ausführungsmodell mit einem Prozess pro Gerät über das TPU-ICI-Netzwerk und unterstützt Standard-PyTorch-APIs, einschließlich sofortiger Ausführung, torch.compile, torch.distributed, DeviceMesh, DTensor und Fully Sharded Data Parallel (FSDP2).

KubeRay und Ray auf TPUs

KubeRay ist ein Kubernetes-Operator zum Bereitstellen und Verwalten von Ray-Anwendungen. Wenn Sie das Ray on GKE-Add-on aktivieren, werden der verwaltete KubeRay-Operator und der KubeRay-TPU-Webhook bereitgestellt. Wenn KubeRay Worker-Pods erstellt, die google.com/tpu-Ressourcen anfordern, konfiguriert der Webhook automatisch die Pod-Hostnamen, den monitorlosen Dienst und die TPU-Umgebungsvariablen, die für die Koordination mehrerer Hosts erforderlich sind.

Ray erkennt die TPU-Topologie auf jedem Knoten beim Start und plant verteilte Arbeitslasten über TPU-Slices hinweg. Sie können verteilte TorchTPU-Arbeitslasten auf zwei Ebenen orchestrieren:

• Ray Core bietet topologiebewusste Platzierungsgruppen und Akteurplanung für benutzerdefinierte verteilte Ausführungsschleifen, z. B. tensorparallele Berechnungen oder Pipelines für bestärkendes Lernen (Reinforcement Learning, RL).
• Ray Train bietet verwaltetes verteiltes datenparalleles und modellparalleles Training, einschließlich PyTorch Distributed Data Parallel (DDP), PyTorch FSDP2 und JAX Single Program, Multiple Data (SPMD), mit integrierter Prüfpunkterstellung und automatischer Fehlerbehebung.

Ziele

In dieser Anleitung wird Folgendes beschrieben:

1. Erstellen Sie einen GKE-Cluster im Autopilot- oder Standardmodus auf TPU Trillium (v6e) oder Ironwood (TPU7x).
2. Erstellen Sie ein Container-Image, das Ray- und Hugging Face-Trainingsbibliotheken über das TorchTPU-Basis-Image legt.
3. Stellen Sie eine benutzerdefinierte RayCluster-Ressource bereit, um die Head- und TPU-Worker-Knoten bereitzustellen.
4. Verteilte tensorparallele Berechnungen mit PyTorch DeviceMesh und DTensor auf einem TPU-Slice mit Ray Core ausführen.
5. Feinabstimmung eines Gemma 4-Modells mit PyTorch FSDP2 über ein Multi-Host-TPU-Slice mit Ray Train.

Kosten

In dieser Anleitung verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

• GKE
• Cloud TPU
• Artifact Registry
• Cloud Storage

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen.

Nach Abschluss der in dieser Anleitung beschriebenen Aufgaben können Sie weitere Kosten vermeiden, indem Sie die erstellten Ressourcen löschen. Weitere Informationen finden Sie unter Bereinigen.

Hinweis

Führen Sie die folgenden Aufgaben aus, bevor Sie beginnen:

• Aktivieren Sie die Artifact Registry API, die Cloud Storage API und die Google Kubernetes Engine API.
APIs aktivieren
• Wenn Sie die Google Cloud CLI für diese Aufgabe verwenden möchten, müssen Sie die gcloud CLI installieren und dann initialisieren. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem Befehl gcloud components update ab. Ältere gcloud CLI-Versionen unterstützen möglicherweise nicht die Ausführung der Befehle in diesem Dokument.
  Hinweis: Legen Sie für vorhandene Installationen der gcloud CLI das compute/region-Attribut fest. Wenn Sie hauptsächlich zonale Cluster verwenden, legen Sie stattdessen compute/zone fest. Wenn Sie einen standardmäßigen Speicherort festlegen, können Sie in der gcloud CLI Fehler wie diesen vermeiden: One of [--zone, --region] must be supplied: Please specify location. Bei bestimmten Befehlen müssen Sie möglicherweise den Speicherort angeben, wenn er sich von dem von Ihnen festgelegten Standard unterscheidet.

• Prüfen Sie, ob Ihr Projekt über ein ausreichendes Kontingent für TPU Trillium (v6e)- oder Ironwood (TPU7x)-Kapazität in der ausgewählten Zone verfügt. Weitere Informationen finden Sie unter Cloud TPU-Kontingente.
• Prüfen Sie, ob Ihr Konto die folgenden IAM-Rollen hat:
  • Kubernetes Engine-Administrator (roles/container.admin)
  • Artifact Registry-Administrator (roles/artifactregistry.admin)
  • Storage-Administrator (roles/storage.admin)

Umgebung vorbereiten

In dieser Anleitung verwenden Sie Cloud Shell zum Verwalten von Ressourcen, die in Google Cloudgehostet werden. Die Software, die Sie für diese Anleitung benötigen, ist in Cloud Shell vorinstalliert, einschließlich kubectl, docker, python3 und der gcloud-Befehlszeile.

So richten Sie Ihre Umgebung mit Cloud Shell ein:

1. Starten Sie in der Google Cloud -Console eine Cloud Shell-Sitzung. Klicken Sie dazu in der Google Cloud -Console auf Cloud Shell aktivieren. Dadurch wird im unteren Bereich der Google Cloud Console eine Sitzung gestartet.

2. Klonen Sie das Repository kubernetes-engine-samples und wechseln Sie in das Beispielverzeichnis für diese Anleitung:

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

3. Erstellen und aktivieren Sie eine lokale virtuelle Python-Umgebung mit installierter Ray-Befehlszeile:

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

4. Wenn Sie die Gemma 4-Modellgewichtungen herunterladen möchten, müssen Sie die Google-Lizenzvereinbarung auf Hugging Face akzeptieren:

   1. Rufen Sie die Gemma 4-Modellseite auf Hugging Face auf.
   2. Melden Sie sich an und akzeptieren Sie die Lizenzbedingungen, indem Sie auf Agree and access repository klicken.
   3. Generieren Sie in den Einstellungen Ihres Hugging Face-Kontos ein Nutzerzugriffstoken mit Read-Berechtigungen.

5. Legen Sie die Umgebungsvariablen für Ihr Google Cloud -Projekt, Ihren Cluster, Ihren Cloud Storage-Bucket und Ihr Hugging Face-Token fest:

   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}
   

   Ersetzen Sie Folgendes:

   • PROJECT_ID: Projekt-ID in Google Cloud .
   • CLUSTER_NAME: der Name Ihres GKE-Cluster, z. B. ray-torchtpu-cluster.
   • REGION: Die Region, in der sich Ihr Artifact Registry-Repository, Ihr Cloud Storage-Bucket und Ihre TPU-Kapazität befinden.
   • ZONE: die Zone, in der Ihre TPU-Kapazität verfügbar ist. Informationen zu den verfügbaren Regionen und Zonen nach TPU-Version finden Sie unter TPU-Regionen und -Zonen.
   • ARTIFACT_REGISTRY: der Name Ihres Artifact Registry-Docker-Repositorys, z. B. ray-torchtpu.
   • BUCKET_NAME: der Name Ihres Cloud Storage-Bucket zum Speichern von Ray Train-Prüfpunkten, z. B. PROJECT_ID-ray-torchtpu.
   • HF_TOKEN: Ihr Hugging Face-Zugriffstoken.

GKE-Cluster erstellen

Sie können verteilte TorchTPU-Trainingsarbeitslasten in einem GKE Autopilot- oder Standardcluster mit aktiviertem verwalteten Ray-Operator-Add-on ausführen.

1. Erstellen Sie Ihren GKE-Cluster und konfigurieren Sie die Anmeldedaten für das kubectl-Befehlszeilentool:

   Autopilot

   1. Erstellen Sie einen regionalen GKE Autopilot-Cluster mit aktiviertem Ray-Operator-Add-on:

      gcloud container clusters create-auto ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --region=${REGION} \
          --cluster-version=1.36 \
          --enable-ray-operator
      

      Wenn Sie die benutzerdefinierte RayCluster-Ressource in einem späteren Schritt bereitstellen, stellt GKE Autopilot automatisch den CPU-Head-Knoten und das Multi-Host-TPU-Slice bereit, die Ihren Pod-Ressourcenanfragen und Knotenselektoren entsprechen.

   2. Konfigurieren Sie das kubectl-Befehlszeilentool für die Kommunikation mit Ihrem Cluster:

      gcloud container clusters get-credentials ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --region=${REGION}
      

   Standard

   1. Erstellen Sie einen zonalen GKE Standard-Cluster mit einem CPU-Knotenpool zum Hosten des Ray-Head-Pods, Image-Streaming, GKE Dataplane V2 und Workload Identity-Föderation für GKE aktiviert:

      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
      

      Hinweis: Wenn Sie Open-Source-KubeRay selbst verwalten, anstatt das GKE-RayOperator-Add-on zu verwenden, installieren Sie KubeRay und den KubeRay-TPU-Webhook in Version 1.4.0 oder höher, um die TorchTPU-Umgebungsvariable einzufügen.

   2. Konfigurieren Sie das kubectl-Befehlszeilentool für die Kommunikation mit Ihrem Cluster:

      gcloud container clusters get-credentials ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --zone=${ZONE}
      

   3. Erstellen Sie einen TPU-Slice-Knotenpool mit mehreren Hosts und 16 Geräten mit einer 4x4-Topologie auf vier Hosts auf TPU Trillium (v6e) oder einer 2x2x2-Topologie auf zwei Hosts auf 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
      

      Tipp: Wenn Sie eine bestimmte TPU-Reservierung nutzen möchten, hängen Sie die Flags --reservation-affinity=specific und --reservation=RESERVATION_NAME an den Befehl node-pools create an, wobei RESERVATION_NAME der Name Ihrer TPU-Reservierung ist. Wenn Sie Spot-VMs verwenden möchten, hängen Sie das Flag --spot an.

2. Erstellen Sie einen Cloud Storage-Bucket, um freigegebene Ray Train-Prüfpunkte auf den TPU-Knoten mit mehreren Hosts zu speichern, und gewähren Sie der Kubernetes-ServiceAccount-Ressource default Zugriff auf den 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"
   

3. Erstellen Sie ein Kubernetes-Secret mit Ihrem Hugging Face-Zugriffstoken, damit Ray-Pods die geschützten Gemma 4-Gewichtungen herunterladen können:

   kubectl create secret generic hf-secret \
       --from-literal=hf_api_token=${HF_TOKEN}
   

Container-Image erstellen

Erstellen Sie ein Container-Image, das Ray und die Hugging Face-Trainingsbibliotheken zusätzlich zum vordefinierten TorchTPU-Basis-Image installiert, das TorchTPU und die libtpu-Bibliothek enthält.

1. Erstellen Sie ein Artifact Registry-Docker-Repository und konfigurieren Sie die Docker-Authentifizierung:

   gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
       --repository-format=docker \
       --location=${REGION}
   
   gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev
   

2. Prüfen Sie die Dockerfile im Beispielverzeichnis:

   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

   Für das Single-Slice-Training mit der Ray-Klasse TorchTrainer und die Planung pro Chiplet auf Ironwood (TPU7x) ist Ray Version 2.58.0 oder höher erforderlich. Für die Feinabstimmung von Gemma 4 ist die transformers-Bibliothek Version 5.5.0 oder höher erforderlich.

3. Erstellen Sie das Image und übertragen Sie es per Push in die Artifact Registry:

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

Benutzerdefinierte RayCluster-Ressource bereitstellen

Stellen Sie eine benutzerdefinierte RayCluster-Ressource bereit, die einen CPU-Head-Pod, ein In-Memory-/dev/shm-Volume für den Ray-Objektspeicher mit gemeinsamem Arbeitsspeicher und eine Multi-Host-TPU-Worker-Gruppe bereitstellt.

1. Sehen Sie sich das RayCluster-Manifest für Ihre TPU-Version im Beispiel-Repository an:

   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

   Auf Ironwood (TPU7x) enthält jeder physische Chip zwei Chiplets, was acht logische TPU-Kerne pro Host mit vier Chips ergibt. Im Manifest wird der Wert des Felds TPU im Abschnitt resources auf "8" festgelegt. Außerdem wird die Umgebungsvariable RAY_TPU_RESOURCE_PER_CHIP auf "2" gesetzt, um acht logische TPU-Ressourcen pro Host anzukündigen und zu planen.

2. Wenden Sie das Manifest an, um die benutzerdefinierte RayCluster-Ressource zu erstellen:

   Trillium (v6e)

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

   Ironwood (TPU7x)

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

3. Warten Sie, bis der RayCluster-Head-Pod und alle TPU-Worker-Pods den Status Running erreicht haben:

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

   Die Ausgabe sieht etwa so aus:

   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
   

4. Prüfen Sie die Umgebungsvariablen auf einem der Worker-Pods, um zu bestätigen, dass der KubeRay-TPU-Webhook die TPU-Koordinations- und TorchTPU-Topologievariablen für mehrere Hosts eingefügt hat:

   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"
   

   Die Ausgabe sieht etwa so aus:

   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
   

5. Richten Sie eine Portweiterleitungssitzung zum Ray-Head-Dienst ein:

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

   Sie können die TPU-Worker-Knoten und ihre automatisch erkannten ray.io/tpu-*-Labels auf dem Tab Cluster des Ray-Dashboards prüfen.

Verteilte Tensor-Parallelität mit Ray Core ausführen

Ray Core bietet die Funktion slice_placement_group im Modul ray.util.tpu, um ein intaktes TPU-Slice mit mehreren Hosts zu reservieren, und die Klasse PlacementGroupSchedulingStrategy, um einen Ray-Akteur pro logischem TPU-Gerät im Slice zu planen.

1. Sehen Sie sich das Script ray_core_collectives.py im Beispiel-Repository an:

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

   Das Script führt die verteilte Tensor-Parallelität auf dem TPU-Slice in drei Schritten aus:

   1. Die Funktion slice_placement_group reserviert ein TPU-Slice mit 16 Geräten und plant einen TorchTPUSliceWorker-Actor auf jedem TPU-Gerät.
   2. Jeder Akteur berechnet seinen globalen Rang aus seinem Hostindex und seiner lokalen Geräte-ID und initialisiert dann die Prozessgruppe tpu_dist.
   3. Die Funktionen init_device_mesh und parallelize_module teilen die Feedforward-Ebenen auf alle 16 TPU-Geräte auf, vergleichen die Ausgabe mit einer CPU-Referenz auf einem einzelnen Gerät und führen ein all_reduce-Kollektiv über ICI aus.

2. Senden Sie den Ray Core-Job für Ihre TPU-Version:

   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
   

   Die Ausgabe sieht etwa so aus:

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

Gemma 4-Modell mit Ray Train und PyTorch FSDP2 optimieren

Die Ray Train-Klasse TorchTrainer automatisiert die TPU-Slice-Reservierung, die Worker-Platzierung und die torch.distributed-Initialisierung über ein TPU-Slice mit mehreren Hosts.

1. Sehen Sie sich das Script ray_train_llm_finetune.py im Beispiel-Repository an, mit dem das Modell google/gemma-4-E2B-it anhand des Anweisungs-Datasets databricks/databricks-dolly-15k optimiert wird:

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

   Das Script konfiguriert die Klasse TorchTrainer mit den folgenden Komponenten:

   • Die Klasse ScalingConfig konfiguriert die verteilte Worker-Topologie für den Trainingsjob. Wenn Sie die Parameter use_tpu=True, num_workers=16, topology und accelerator_type festlegen, wird ein TPU-Slice mit 16 Geräten reserviert und ein Worker pro TPU-Kern geplant.
   • Die Klasse TorchConfig konfiguriert die Prozessgruppe torch.distributed auf den Workern. Wenn Sie den Parameter backend="tpu_dist" festlegen, werden alle 16 Worker über das TPU-ICI-Netzwerk verbunden, bevor die Funktion train_func ausgeführt wird.
   • Die Funktion create_dataloader lädt das Anweisungs-Dataset herunter und tokenisiert es, fasst die Token-Sequenzen mit einem PyTorch-DataLoader-Objekt zusammen und ruft die Funktion ray.train.torch.prepare_data_loader auf, damit jeder Worker ein disjunktes Shard jedes globalen Batches auf seinem zugewiesenen TPU-Gerät erhält.
   • Die Funktion train_func lädt das Modell mit der Genauigkeit bfloat16 und ruft die PyTorch FSDP2-Funktion fully_shard für jeden Feedforward-Block und das Stammmodell auf, um Parameter, Gradienten und Optimierungszustände auf alle 16 TPU-Kerne zu verteilen.
   • Die Klasse Checkpoint und die Funktion ray.train.report erfassen am Ende jeder Epoche das vollständige Modellstatus-Dictionary auf Rang 0 und laden den Prüfpunkt in Ihren Cloud Storage-Bucket hoch.

2. Senden Sie den TorchTrainer-FSDP2-Feinabstimmungsjob für Ihre TPU-Version:

   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
   

   Tipp: Standardmäßig wird das Script mit einer Teilmenge von 1.024 Beispielen für drei Epochen trainiert, sodass der Job in wenigen Minuten abgeschlossen ist. Wenn Sie mit allen 15.011 Beispielen trainieren oder die Anzahl der Epochen ändern möchten, übergeben Sie das Flag --dataset-split=train oder --num-epochs an das Script ray_train_llm_finetune.py.

   Die Ausgabe sieht etwa so aus:

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

   Die Trainingsausgabe sieht so aus:

   • Epoche 1 umfasst die anfängliche XLA-Grafikkompilierung. In den Epochen 2 und 3 werden die im Cache gespeicherten ausführbaren Dateien aus /dev/shm/torch_tpu_cache wiederverwendet. Dabei werden über 20.000 Tokens pro Sekunde auf TPU Trillium (v6e) und 31.000 Tokens pro Sekunde auf Ironwood (TPU7x) erreicht.
   • Der globale Kreuzentropieverlust, der über alle 16 FSDP2-Worker-Ränge synchronisiert wird, sinkt von 3,14 in Epoche 1 auf 2,20 in Epoche 3 über die drei Feinabstimmungsepochen im databricks/databricks-dolly-15k-Anweisungs-Dataset.
   • Am Ende jeder Epoche erfasst Rang 0 das vollständige Modellstatus-Dictionary und lädt den Prüfpunkt in Ihren Cloud Storage-Bucket hoch. Dabei wird der Prüfpunkt mit dem geringsten Verlust beibehalten, wie im CheckpointConfig(num_to_keep=1, checkpoint_score_attribute="loss", checkpoint_score_order="min")-Objekt konfiguriert.

Erfolg: Sie haben eine verteilte tensorparallele Berechnung mit Ray Core ausgeführt und ein Gemma 4-Modell mit PyTorch FSDP2 auf TPUs in der GKE mithilfe von Ray Train und TorchTPU optimiert.

Arbeitslastleistung beobachten

Informationen zum Überwachen der TPU-TensorCore-Auslastung, der HBM-Nutzung und der Ray-Cluster-Messwerte in Cloud Monitoring finden Sie unter Logs und Messwerte für Ray-Cluster in GKE erfassen und ansehen und TPU-Messwerte im Ray-Dashboard ansehen.

Auf das Ray-Dashboard zugreifen

So prüfen Sie den Status Ihrer Ray-Akteure, Platzierungsgruppen und Trainingsjobs im integrierten Ray-Dashboard:

1. Wenn die Portweiterleitung zum Ray-Head-Dienst noch nicht ausgeführt wird, starten Sie sie im Hintergrund:

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

2. Klicken Sie in Cloud Shell auf Webvorschau und wählen Sie Vorschau auf Port 8265 aus oder öffnen Sie http://localhost:8265 in Ihrem lokalen Browser.

3. Klicken Sie auf die Tabs Jobs, Cluster und Messwerte, um Worker-Logs, den Status der Platzierungsgruppe, TPU-Knotenlabels und das Diagramm TPU-Auslastung nach Knoten zu prüfen.

Bereinigen

Damit Ihrem Google Cloud Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie die benutzerdefinierte RayCluster-Ressource, den GKE-Cluster, den Cloud Storage-Bucket und das Artifact Registry-Repository:

1. Beenden Sie den Portweiterleitungsprozess im Hintergrund und löschen Sie die benutzerdefinierte RayCluster-Ressource und das Kubernetes-Secret:

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

2. Löschen Sie den GKE-Cluster:

   Autopilot

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

   Standard

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

3. Löschen Sie den Cloud Storage-Bucket und seinen Prüfpunktinhalt:

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

4. Löschen Sie das Artifact Registry-Repository:

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

Nächste Schritte

• Weitere Informationen zu PyTorch auf Cloud TPU
• TorchTPU-Dokumentation lesen
• LLMs auf TPUs mit mehreren Hosts mit Ray und TorchTPU bereitstellen
• KI-Training auf TPUs mit DWS und Kueue optimieren
• TPUs in GKE
• Ray Train-Dokumentation lesen