LLM mit JAX, Ray Train und TPU Trillium in GKE trainieren

Autopilot Standard

In dieser Anleitung erfahren Sie, wie Sie das Llama 3 8B-LLM (Large Language Model) mit MaxText, Ray Train und TPUs in Google Kubernetes Engine (GKE) trainieren.

In dieser Anleitung wird der gesamte Prozess von der Konfiguration der erforderlichen Cloud-Infrastruktur bis zum Einreichen und erfolgreichen Ausführen der Trainingsarbeitslast auf TPUs mit mehreren Hosts beschrieben.

Diese Anleitung richtet sich an Plattformadministratoren und ‑operatoren sowie an Daten- und KI-Spezialisten, die lernen möchten, wie sie große Modelle auf einem verteilten TPU-Slice mit mehreren Hosts trainieren.

Hintergrund

Die Kombination aus GKE, KubeRay, MaxText und TPUs bietet eine leistungsstarke und skalierbare Plattform für das Training von Modellen im großen Maßstab. In diesem Abschnitt werden die in diesem Leitfaden verwendeten Schlüsseltechnologien beschrieben:

JAX

JAX ist eine Python-Bibliothek für die beschleunigerorientierte Array-Berechnung und Programmtransformation, die für numerisches Hochleistungs-Computing und maschinelles Lernen im großen Maßstab entwickelt wurde.

JAX bietet ein erweiterbares System zum Transformieren numerischer Funktionen wie jax.grad, jax.jit und jax.vmap. Dabei wird der XLA-Compiler verwendet, um hochoptimierten Code zu erstellen, der effizient auf Beschleunigern wie GPUs und TPUs skaliert. Die Stärke von JAX liegt in seiner Zusammensetzbarkeit, die es Nutzern ermöglicht, diese Transformationen zu kombinieren, um komplexe, leistungsstarke numerische Programme für die verteilte Ausführung zu erstellen.

MaxText

MaxText ist ein leistungsstarkes Open-Source-LLM (Large Language Model), das für Skalierbarkeit und Anpassbarkeit entwickelt wurde. MaxText basiert auf JAX und ist für die effiziente Ausführung auf Cloud TPU und GPUs optimiert.

TPUs

Tensor Processing Units (TPUs) sind von Google speziell entwickelte Beschleuniger zur Optimierung von Arbeitslasten für maschinelles Lernen. Im Gegensatz zu CPUs für allgemeine Zwecke oder GPUs für die Parallelverarbeitung sind TPUs hochgradig auf die massiven Matrix- und Tensorberechnungen spezialisiert, die die Grundlage von Deep Learning bilden. Dadurch sind sie für diese spezielle Aufgabe effizient. Der Hauptvorteil von TPUs ist die Leistung bei großem Umfang.

In dieser Anleitung wird TPU Trillium verwendet, die sechste Generation von TPUs. Weitere Informationen finden Sie unter Vorteile der Verwendung von TPU Trillium.

KubeRay

KubeRay ist ein Kubernetes-Operator, der eine einheitliche Möglichkeit zum Bereitstellen, Verwalten und Überwachen von Ray-Anwendungen in Kubernetes bietet. Der KubeRay-Operator wird über das Ray on GKE-Add-on installiert und verwaltet. Dies ist die empfohlene Methode zum Bereitstellen und Verwalten von Ray-Clustern in GKE.

Ziele

In dieser Anleitung wird Folgendes beschrieben:

1. Richten Sie einen GKE-Cluster mit einem TPU-Knotenpool mit mehreren Hosts ein.
2. KubeRay so konfigurieren, dass die Umgebung für verteiltes Training verwaltet wird.
3. Erstellen Sie ein benutzerdefiniertes Docker-Image, das MaxText-, Ray- und JAX-Abhängigkeiten enthält.
4. Erstellen Sie ein Python-Trainingsscript, in dem der MaxText-Trainingsloop mit JaxTrainer von Ray Train über das TPU-Slice hinweg orchestriert wird.
5. Definieren Sie eine benutzerdefinierte RayCluster-Ressource, um die Head- und Worker-Knoten mit den erforderlichen TPU-Ressourcen bereitzustellen.
6. Senden Sie den Trainingsjob an RayCluster und überwachen Sie den Fortschritt.
7. Verwenden Sie Cloud Storage zum Speichern von Modellprüfpunkten.

Hinweise

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

Install the Google Cloud CLI.

Hinweis: Wenn Sie die gcloud CLI bereits installiert haben, prüfen Sie, ob Sie die aktuelle Version haben, indem Sie gcloud components update ausführen.

Wenn Sie einen externen Identitätsanbieter (IdP) verwenden, müssen Sie sich zuerst mit Ihrer föderierten Identität in der gcloud CLI anmelden.

Führen Sie den folgenden Befehl aus, um die gcloud CLI zu initialisieren:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

Verify that billing is enabled for your Google Cloud project.

Enable the required API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable container.googleapis.com

Install the Google Cloud CLI.

Hinweis: Wenn Sie die gcloud CLI bereits installiert haben, prüfen Sie, ob Sie die aktuelle Version haben, indem Sie gcloud components update ausführen.

Wenn Sie einen externen Identitätsanbieter (IdP) verwenden, müssen Sie sich zuerst mit Ihrer föderierten Identität in der gcloud CLI anmelden.

Führen Sie den folgenden Befehl aus, um die gcloud CLI zu initialisieren:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

Verify that billing is enabled for your Google Cloud project.

Enable the required API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable container.googleapis.com

• Grant roles to your user account. Run the following command once for each of the following IAM roles: roles/container.admin, roles/iam.serviceAccountAdmin

  gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

  Replace the following:

  • PROJECT_ID: Your project ID.
  • USER_IDENTIFIER: The identifier for your user account. For example, [e-mail].
  • ROLE: The IAM role that you grant to your user account.

• Da in dieser Anleitung TPU Trillium (v6e) verwendet wird, müssen Sie eine Region oder Zone mit Verfügbarkeit auswählen. Weitere Informationen finden Sie unter Cloud TPU-Kontingente.

Umgebung vorbereiten

In dieser Anleitung verwenden Sie Cloud Shell. Die in dieser Anleitung verwendeten Befehlszeilentools gcloud, helm und kubectl sind in Cloud Shell vorinstalliert.

1. Rufen Sie die Google Cloud Console auf.

2. Klicken Sie oben im Google Cloud Console-Fenster auf die Schaltfläche Cloud Shell aktivieren .

   In einem neuen Frame in derGoogle Cloud Console wird eine Cloud Shell-Sitzung geöffnet und darin eine Eingabeaufforderung angezeigt.

3. Erstellen und aktivieren Sie eine virtuelle Python-Umgebung:

   python3 -m venv ray-env
   source ray-env/bin/activate
   

4. Installieren Sie die Ray-Befehlszeile und andere Abhängigkeiten:

   pip install "ray[default]==2.49.1"
   

5. Legen Sie die folgenden Umgebungsvariablen fest:

   export PROJECT_ID=$(gcloud config get project)
   export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
   export GS_BUCKET=GS_BUCKET
   export KSA_NAME=KSA_NAME
   export NAMESPACE=default
   export CLUSTER_NAME=CLUSTER_NAME
   export REGION=REGION
   export ZONE=ZONE
   export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
   

   Ersetzen Sie Folgendes:

   • GS_BUCKET: der Name des Cloud Storage-Buckets.
   • KSA_NAME: Der Name des Kubernetes-ServiceAccount.
   • CLUSTER_NAME ist der Name des neuen Clusters.
   • REGION: Die Region, in der Ihre TPU-Trillium-Kapazität verfügbar ist.
   • ZONE: Die Zone, in der Ihre TPU-Trillium-Kapazität verfügbar ist. Weitere Informationen finden Sie unter TPU-Verfügbarkeit in GKE.
   • ARTIFACT_REGISTRY: der Name des Artifact Registry-Repositorys.

GKE-Cluster erstellen

Sie können KubeRay auf TPUs in einem GKE Autopilot- oder Standardcluster konfigurieren. Für eine vollständig verwaltete Kubernetes-Umgebung empfehlen wir die Verwendung eines Autopilot-Clusters. Informationen zum Auswählen des GKE-Betriebsmodus, der für Ihre Arbeitslasten am besten geeignet ist, finden Sie unter GKE-Betriebsmodi.

Autopilot

1. Führen Sie in Cloud Shell den folgenden Befehl aus:

   gcloud container clusters create-auto $CLUSTER_NAME \
       --enable-ray-operator \
       --machine-type=n1-standard-16 \
       --location=$REGION
   

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

   gcloud container clusters get-credentials CLUSTER_NAME \
       --location=$ZONE
   

Standard

1. Erstellen Sie in Cloud Shell einen Standardcluster, in dem das Add-on Ray-Operator aktiviert ist, indem Sie den folgenden Befehl ausführen:

   gcloud container clusters create $CLUSTER_NAME \
       --addons=RayOperator \
       --addons GcsFuseCsiDriver \
       --machine-type=n1-standard-16 \
       --workload-pool=$PROJECT_ID.svc.id.goog \
       --location=$ZONE
   

   Mit diesem Befehl wird auch GcsFuseCsiDriver aktiviert, sodass Pods Cloud Storage-Buckets als lokale Dateisysteme bereitstellen können. Die Erstellung eines Clusters kann einige Minuten dauern.

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

   gcloud container clusters get-credentials CLUSTER_NAME \
       --location=LOCATION
   

3. So erstellen Sie einen TPU-Slice-Knotenpool mit mehreren Hosts:

   gcloud container node-pools create v6e-16 \
       --location=$ZONE \
       --cluster=$CLUSTER_NAME \
       --machine-type=ct6e-standard-4t \
       --threads-per-core=1 \
       --tpu-topology=4x4 \
       --num-nodes=4
   

GKE stellt einen Knotenpool mit vier TPU Trillium-VMs (v6e) bereit, die als TPU-Slice mit mehreren Hosts mit einer 4x4-Topologie konfiguriert sind und für verteilte Trainingsarbeitslasten bereit sind.

In einem GKE-Cluster mit aktiviertem Ray-Operator werden KubeRay und der KubeRay-TPU-Webhook automatisch in Ihrem Cluster installiert.

Cloud Storage-Bucket und Dienstkonto konfigurieren

1. Erstellen Sie einen Cloud Storage-Bucket für freigegebene Prüfpunkte zwischen den TPU-Knoten mit mehreren Hosts.

   gsutil mb -p ${PROJECT_ID} -c STANDARD -l ${REGION} gs://${GS_BUCKET}
   

2. So aktivieren Sie den Zugriff auf den Cloud Storage-Bucket:

   kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}
   

3. Fügen Sie dem Dienstkonto die erforderlichen IAM-Richtlinienbindungen hinzu, um den Zugriff auf den Cloud Storage-Bucket zu ermöglichen:

   gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
       --member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
       --role "roles/storage.objectUser"
   

Trainingsskript erstellen

Im folgenden Skript wird JaxTrainer von Ray Train verwendet, um einen verteilten MaxText-Trainingsjob auszuführen. Das Skript konfiguriert die Trainingsumgebung für einen TPU-Slice-Knotenpool mit mehreren Hosts und führt den MaxText-Trainingsjob auf jedem Worker-Knoten aus. Die Funktion train_loop_per_worker umschließt den Haupteinstiegspunkt von MaxText und verwendet den verteilten Scheduler von Ray, um den MaxText-Trainer auf einem TPU-Slice mit mehreren Hosts auszuführen.

1. Speichern Sie das folgende Python-Skript als maxtext_ray_trainer.py:

   import os
   from absl import app
   import logging
   from typing import Sequence
   import ray
   from ray.train.v2.api.config import ScalingConfig, RunConfig
   from ray.train.v2.jax import JaxTrainer
   
   def train_loop_per_worker(config):
       from MaxText.train import main as maxtext_main
   
       argv = config["argv"]
       maxtext_main(argv)
   
   def main(argv: Sequence[str]):
       trainer = JaxTrainer(
           train_loop_per_worker=train_loop_per_worker,
           train_loop_config={"argv": argv},
           scaling_config=ScalingConfig(
               use_tpu=True,
               num_workers=4,
               topology="4x4",
               accelerator_type="TPU-V6E",
               resources_per_worker={"TPU": 4},
               placement_strategy="SPREAD",
           ),
           run_config=RunConfig(
               name="maxtext_jaxtrainer",
               worker_runtime_env={
                   "env_vars": {
                       "JAX_PLATFORMS": "tpu",
                       "ENABLE_PJRT_COMPATIBILITY": "true",
                       "TPU_SLICE_BUILDER_DUMP_CHIP_FORCE": "true",
                       "TPU_SLICE_BUILDER_DUMP_ICI": "true",
                       "XLA_FLAGS": "--xla_dump_to=/tmp/xla_dump_file --xla_dump_hlo_as_proto",
                   }
               },
           ),
       )
       result = trainer.fit()
       logging.info("Training complete!")
       ray.shutdown()
   
   if __name__ == "__main__":
       app.run(main)

2. Erstellen Sie ein Artifact Registry-Repository zum Hosten des benutzerdefinierten Images:

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

3. Wenn Sie ein Image erstellen möchten, das Ray- und MaxText-Abhängigkeiten für das Training enthält, erstellen Sie ein Dockerfile:

   # Start from a Ray base image which includes JaxTrainer API.
   # Maxtext with TPU requires Python 3.12.
   FROM rayproject/ray:2.49.1-py312
   
   USER root
   RUN groupadd -r ray 2>/dev/null || true && usermod -g ray ray
   
   RUN sudo apt-get update -y \
     && sudo apt-get install --no-install-recommends -y git \
     && sudo rm -rf /var/lib/apt/lists/*
   
   WORKDIR /app
   
   # Clone the Maxtext repo and build from source, installing TPU dependencies.
   RUN git clone https://github.com/AI-Hypercomputer/maxtext.git
   
   RUN pip install --no-cache-dir uv
   
   RUN cd maxtext && \
       uv pip install --no-cache --system -e .[tpu] --resolution=lowest && \
       install_maxtext_github_deps
   
   # Copy the Ray Maxtext trainer to run on the remote container.
   COPY maxtext_ray_trainer.py .
   
   RUN chown -R ray:ray .
   ENV PYTHONPATH=/app/maxtext/src:/app/maxtext:/app
   USER ray

   Hinweis :In dieser Anleitung wird ein benutzerdefiniertes Image verwendet, da die Abhängigkeiten zu groß sind, um sie mit dem Ray-Befehl runtime_env zu installieren.

4. Erstellen Sie das Docker-Image, taggen Sie es und übertragen Sie es per Push in Artifact Registry:

   export DOCKER_IMAGE=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-maxtext:latest
   gcloud builds submit --tag ${DOCKER_IMAGE}
   

Modell trainieren

1. Speichern Sie das folgende Beispielmanifest als maxtext-tpu-cluster.yaml:

   apiVersion: ray.io/v1
   kind: RayCluster
   metadata:
     name: maxtext-tpu-cluster
   spec:
     headGroupSpec:
       rayStartParams: {}
       template:
         metadata:
           annotations:
             gke-gcsfuse/volumes: "true"
             gke-gcsfuse/cpu-limit: "0"
             gke-gcsfuse/memory-limit: "0"
             gke-gcsfuse/ephemeral-storage-limit: "0"
         spec:
           serviceAccountName: ${KSA_NAME}
           containers:
             - name: ray-head
               image: ${DOCKER_IMAGE}
               imagePullPolicy: IfNotPresent
               ports:
               - containerPort: 6379
                 name: gcs-server
               - containerPort: 8265
                 name: dashboard
               - containerPort: 10001
                 name: client
               resources:
                 limits:
                   memory: "16Gi"
                 requests:
                   cpu: "8"
                   memory: "16Gi"
               volumeMounts:
               - name: gcs-fuse-csi-ephemeral
                 mountPath: /data
               - name: dshm
                 mountPath: /dev/shm
           volumes:
           - name: gcs-fuse-cache
             emptyDir:
               medium: Memory
           - name: dshm
             emptyDir:
               medium: Memory
           - name: gcs-fuse-csi-ephemeral
             csi:
               driver: gcsfuse.csi.storage.gke.io
               volumeAttributes:
                 bucketName: ${GS_BUCKET}
                 mountOptions: "implicit-dirs"
     workerGroupSpecs:
       - replicas: 1
         numOfHosts: 4
         groupName: tpu-group
         rayStartParams: {}
         template:
           metadata:
             annotations:
               gke-gcsfuse/volumes: "true"
               gke-gcsfuse/cpu-limit: "0"
               gke-gcsfuse/memory-limit: "0"
               gke-gcsfuse/ephemeral-storage-limit: "0"
           spec:
             serviceAccountName: ${KSA_NAME}
             containers:
               - name: ray-worker
                 image: ${DOCKER_IMAGE}
                 imagePullPolicy: IfNotPresent
                 resources:
                   limits:
                     memory: 200G
                     google.com/tpu: "4"
                   requests:
                     cpu: "8"
                     memory: 200G
                     google.com/tpu: "4"
                 env:
                   - name: JAX_PLATFORMS
                     value: tpu
                   - name: ENABLE_PJRT_COMPATIBILITY
                     value: "true"
                 volumeMounts:
                 - name: gcs-fuse-csi-ephemeral
                   mountPath: /data
                 - name: dshm
                   mountPath: /dev/shm
             volumes:
             - name: gcs-fuse-cache
               emptyDir:
                 medium: Memory
             - name: dshm
               emptyDir:
                 medium: Memory
             - name: gcs-fuse-csi-ephemeral
               csi:
                 driver: gcsfuse.csi.storage.gke.io
                 volumeAttributes:
                   bucketName: ${GS_BUCKET}
                   mountOptions: "implicit-dirs"
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4

   Mit der vorherigen RayCluster-Spezifikation wird eine TPU-Worker-Gruppe mit vier Workern (numOfHosts: 4) pro Replikat erstellt. Jeder Worker fordert vier TPU-Chips (google.com/tpu: "4") an. Die Worker werden auf einem Knoten geplant, auf dem TPU Trillium (tpu-v6e-slice) ausgeführt wird und der Teil desselben gemeinsam untergebrachten Multi-Host-Slice ist. KubeRay skaliert alle vier Worker atomar und die erforderlichen JAX-Umgebungsvariablen sowie die Pod-Affinitäten für die Planung werden von GKE über einen mutierenden Webhook gebootstrapped.

2. Um erforderliche Werte in der YAML-Datei zu konfigurieren, erstellen Sie den RayCluster mit envsubst:

   envsubst < maxtext-tpu-cluster.yaml | kubectl apply -f -
   

3. Prüfen Sie, ob der Cluster bereit ist und ausgeführt wird:

   kubectl get rayclusters maxtext-tpu-cluster
   

   Die Ausgabe sollte in etwa so aussehen:

   NAME                  DESIRED WORKERS   AVAILABLE WORKERS   CPUS   MEMORY        GPUS   STATUS   AGE
   maxtext-tpu-cluster   4                 4                   40     798027216Ki   0      ready    11m
   

4. Wenn Sie über den Ray-Head-Dienst auf das Ray-Dashboard zugreifen möchten, richten Sie eine Portweiterleitungssitzung ein:

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

5. Prüfen Sie, ob der RayCluster von Ihrer lokalen Umgebung aus erreichbar ist:

   ray list nodes --address http://localhost:8265
   

   Die Ausgabe sollte in etwa so aussehen:

   ======== List: 2025-09-13 03:53:16.988269 ========
   Stats:
   ------------------------------
   Total: 5
   Table:
   ------------------------------
       NODE_ID                                                   NODE_IP    IS_HEAD_NODE    STATE    STATE_MESSAGE    NODE_NAME    RESOURCES_TOTAL                  LABELS
   0  92c79d04c34b659c1e3044f7642ad3fd47eb16f290785237149fab56  10.84.0.9
   (...)
   

6. Senden Sie das JaxTrainer-Script an den RayCluster und prüfen Sie, ob der RayJob erfolgreich abgeschlossen wurde:

   ray job submit \
     --address http://localhost:8265 \
     -- python /app/maxtext_ray_trainer.py \
         /app/maxtext/src/MaxText/configs/base.yml \
          base_output_directory=/data/ \
         dataset_type=synthetic \
         per_device_batch_size=1 \
         max_target_length=4096 \
         model_name=llama3-8b \
         steps=100 \
         ici_fsdp_parallelism=4 \
         ici_tensor_parallelism=4 \
         run_name=rayjob-8b-4096-tp4-4x4
   

   Mit dem vorherigen Befehl wird das Python-Script, das den JaxTrainer-Ray-Code aufruft, an den Ray-Cluster gesendet. Der Befehl ray job submit enthält einige MaxText-spezifische Argumente, die an die Modellkonfiguration übergeben werden.

   Hinweis :Es kann einige Minuten dauern, bis das Skript abgeschlossen ist.

   Im Terminal sollte eine Ausgabe ähnlich der folgenden angezeigt werden:

   (RayTrainWorker pid=21663, ip=10.12.3.6) completed step: 99, seconds: 1.100, TFLOP/s/device: 179.739, Tokens/s/device: 3725.218, total_weights: 65536, loss: 0.000 [repeated 3x across cluster]
   
   ------------------------------------------
   Job 'raysubmit_zCrJcWnuymMQv4C3' succeeded
   ------------------------------------------
   

Erfolg :Sie haben das Training mit JaxTrainer von Ray Train und MaxText in verteilten TPU-Knotenpools mit mehreren Hosts abgeschlossen.

Bereinigen

Damit Ihrem Google Cloud -Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, können Sie entweder das Projekt löschen, das die Ressourcen enthält, oder das Projekt beibehalten und die einzelnen Ressourcen löschen.

1. Löschen Sie den RayCluster:

   kubectl delete raycluster maxtext-tpu-cluster
   

2. Löschen Sie den GKE-Cluster:

   gcloud container clusters delete $CLUSTER_NAME --zone=$ZONE
   

3. Löschen Sie den Cloud Storage-Bucket:

   gsutil rm -r gs://${GS_BUCKET}
   

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

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

Nächste Schritte

• Weitere Informationen zu Ray in Kubernetes
• Informationen zum Bereitstellen von vLLM in GKE mit TPUs
• Informationen zum Bereitstellen von SDXL in GKE mit TPUs
• Weitere Informationen zu TPUs in GKE