Modell mit TPU v6e trainieren

Achtung: Diese Seite bezieht sich auf die Cloud TPU API, die nicht mehr aktiv weiterentwickelt wird. Die Cloud TPU API erhält nur Fehlerkorrekturen und Sicherheitsupdates. Für die neuesten Funktionen und die Unterstützung der neuesten TPU-Versionen empfehlen wir, Compute Engine oder Google Kubernetes Engine zum Verwalten von TPU-Ressourcen zu verwenden. Weitere Informationen finden Sie unter TPU-Ressourcen in Compute Engine und TPUs in GKE. Die Cloud TPU API wird auf TPU7x und höheren TPU-Versionen nicht unterstützt.

In diesem Dokument wird beschrieben, wie Sie Modelle auf Cloud TPU v6e (Trillium) trainieren. Wir behandeln dabei die Einrichtung der Umgebung, die Leistungsoptimierung und praktische Trainingsbeispiele mit JAX und PyTorch.

TPU v6e (Trillium) ist die 6. Generation der TPUs von Google. Auf allen technischen Oberflächen, z. B. in der API und in Logs, sowie in diesem Dokument wird Trillium als v6e bezeichnet. Mit 256 Chips pro Pod hat die Architektur von TPU v6e viele Ähnlichkeiten mit v5e. TPU v6e ist für das Training, die Feinabstimmung und die Bereitstellung von Transformer-, Text-zu-Bild- und CNN-Modellen (Convolutional Neural Network) optimiert. Weitere Informationen zur Systemarchitektur und zu den Konfigurationen von TPU v6e finden Sie unter TPU v6e.

Informationen zum Ausführen von Inferenzen auf Cloud TPU v6e finden Sie in den folgenden Anleitungen:

• MaxDiffusion-Inferenz auf v6e
• vLLM-Inferenz auf v6e
• Inferenz auf mehreren Hosts mit Pathways durchführen

Vorbereitung

Zur Vorbereitung sind folgende Schritte erforderlich:

• Google Cloud -Projekt für TPUs einrichten
• Google Cloud CLI installieren
• Compute Engine API aktivieren
• Compute Engine-Dienstkonto erstellen und Berechtigungen erteilen

Weitere Informationen finden Sie unter Cloud TPU-Umgebung einrichten.

Kontingent und Berechtigungen prüfen

Prüfen Sie, ob Ihr Projekt die folgenden Kontingente hat:

• Compute Engine-TPU-Beschleunigerkontingent
• Compute Engine-VM-Kontingent
• Kontingent für IP-Adressen

• Kontingent für Hyperdisk Balanced und alle anderen Laufwerkstypen, die Sie verwenden möchten

  Hinweis: v6e unterstützt Hyperdisk Balanced und Hyperdisk ML. Persistent Disk (PD) wird von v6e nicht unterstützt. Weitere Informationen finden Sie unter Speicheroptionen für Cloud TPU-Daten.

Wenn Sie Google Kubernetes Engine (GKE) mit XPK (Accelerated Processing Kit) verwenden, benötigen Sie zusätzliche Berechtigungen in der Console von Google Cloud . Weitere Informationen finden Sie unter Erforderliche Berechtigungen für die Console von Google Cloud .

Bereitstellungsoptionen

Sie können TPU v6e mit den folgenden Methoden bereitstellen und verwalten:

• Compute Engine (Google Cloud CLI): Sie können das Google Cloud CLI-Tool verwenden, um TPU-VMs direkt als Compute Engine-Ressourcen zu verwalten. Sie können einzelne TPU-VMs bereitstellen oder verwaltete Instanzgruppen (MIGs) verwenden, um Gruppen von TPU-VMs zu verwalten. Weitere Informationen finden Sie unter TPU-Erstellung – Übersicht.
• GKE: Mit der GKE können Sie TPUs als Pool von Beschleunigern für Ihre containerisierten ML-Arbeitslasten bereitstellen und verwalten. Weitere Informationen finden Sie unter TPUs in GKE.
• GKE und XPK: XPK ist ein Befehlszeilentool, das die Clustererstellung und die Ausführung von Arbeitslasten in der GKE vereinfacht. Es wurde für ML-Nutzer ohne umfassende Kubernetes-Kenntnisse entwickelt, die TPUs bereitstellen und Trainingsjobs ausführen möchten. Weitere Informationen finden Sie im XPK-GitHub-Repository.

v6e-TPUs mit GKE und XPK bereitstellen

Wenn Sie GKE mit v6e-TPUs verwenden, können Sie Kubernetes-Befehle oder XPK verwenden, um TPUs bereitzustellen und Modelle zu trainieren oder bereitzustellen. Weitere Informationen zur Verwendung von GKE mit TPUs finden Sie unter TPUs in GKE.

Cloud TPU v6e unterstützt Netzwerkkartenkonfigurationen (NIC), mit denen Sie den Durchsatz über mehrere Netzwerke hinweg skalieren können. In den folgenden Abschnitten finden Sie Befehle zum Erstellen eines GKE-Cluster mit Unterstützung für eine einzelne NIC oder für multi-NIC mithilfe von XPK. Für die meisten Single-Slice-Arbeitslasten bietet Single-NIC eine ausreichende Leistung bei geringerem Konfigurationsaufwand. Verwenden Sie für Multislice-Arbeitslasten und Arbeitslasten, die eine hohe Datenaufnahmegeschwindigkeit erfordern, Multi-NIC.

Cluster mit Unterstützung für eine einzelne NIC mit XPK erstellen

Für die meisten Single-Slice-Arbeitslasten bietet eine einzelne NIC ausreichende Leistung bei weniger Konfigurationsaufwand. Verwenden Sie für Multislice-Arbeitslasten und Arbeitslasten, die eine hohe Datenaufnahmegeschwindigkeit erfordern, Multi-NIC.

In den folgenden Abschnitten wird beschrieben, wie Sie mit XPK einen GKE-Cluster mit Unterstützung für eine einzelne NIC erstellen.

XPK installieren und Umgebungsvariablen einrichten

1. Installieren Sie XPK. Folgen Sie der Anleitung im XPK-GitHub-Repository.

2. Richten Sie Umgebungsvariablen für Ihren Cluster ein:

   export CLUSTER_NAME=XPK_CLUSTER_NAME
   export ZONE=us-east1-d
   export PROJECT_ID=PROJECT_ID
   export ACCELERATOR_TYPE=ACCELERATOR_TYPE
   export NUM_SLICES=1

   Legen Sie die folgenden Umgebungsvariablen fest:

   • CLUSTER_NAME: Ein Name für Ihren Cluster.
   • ZONE: Die Zone, in der der TPU-Cluster erstellt wird. Weitere Informationen zu unterstützten Zonen finden Sie unter Regionen und Zonen.
   • PROJECT_ID: Ihre Google Cloud Projekt-ID
   • ACCELERATOR_TYPE: Der TPU-Typ, auch Beschleunigertyp genannt, gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Beispiel: v6e-256 Weitere Informationen zu den unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
   • NUM_SLICES: Die Anzahl der TPU-Slices für Ihren Cluster. Jedes Slice hat die in ACCELERATOR_TYPE angegebene Anzahl von Chips. Legen Sie für einen Cluster mit einem einzelnen Slice NUM_SLICES auf 1 fest. Geben Sie für einen Multislice-Cluster die Anzahl der Slices basierend auf den Skalierbarkeitsanforderungen Ihrer Arbeitslast an. Die Gesamtzahl der Chips im Cluster ist die Anzahl der Chips in ACCELERATOR_TYPE multipliziert mit NUM_SLICES.

Cluster erstellen

Wählen Sie eine der folgenden Optionen aus, um den Cluster zu erstellen. Für eine optimale Leistung wird die Verwendung eines benutzerdefinierten Netzwerks mit einer MTU von 8.896 empfohlen. Weitere Informationen finden Sie unter MTU konfigurieren.

Benutzerdefinierte Netzwerke

So erstellen Sie ein benutzerdefiniertes Netzwerk mit einer MTU von 8.896 und verwenden es für Ihren Cluster:

1. Legen Sie Umgebungsvariablen für die Netzwerk- und Firewallnamen fest:

   export NETWORK_NAME=NETWORK_NAME
   export NETWORK_FW_NAME=FIREWALL_NAME

   Ersetzen Sie Folgendes:

   • NETWORK_NAME: Ein Name für das Netzwerk.
   • FIREWALL_NAME: Ein Name für die Netzwerk-Firewallregel.

2. Erstellen Sie ein benutzerdefiniertes Netzwerk mit einer MTU von 8.896:

   gcloud compute networks create ${NETWORK_NAME} \
       --mtu=8896 \
       --project=${PROJECT_ID} \
       --subnet-mode=auto \
       --bgp-routing-mode=regional

3. Erstellen Sie eine Firewallregel, die TCP-, ICMP- und UDP-Traffic in Ihrem Netzwerk zulässt:

   gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
       --network=${NETWORK_NAME} \
       --allow tcp,icmp,udp \
       --project=${PROJECT_ID}

4. Legen Sie eine Umgebungsvariable für die XPK-Clusterargumente fest, um das von Ihnen erstellte Netzwerk zu verwenden:

   export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

5. Erstellen Sie den XPK-Cluster. Mit dem folgenden Befehl wird On-Demand-Kapazität bereitgestellt:

   xpk cluster create --cluster=${CLUSTER_NAME} \
       --cluster-cpu-machine-type=e2-standard-8 \
       --num-slices=${NUM_SLICES} \
       --tpu-type=${ACCELERATOR_TYPE} \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --on-demand \
       --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

   Wenn Sie reservierte Kapazität verwenden möchten, ersetzen Sie --on-demand durch --reservation=RESERVATION_NAME. Wenn Sie TPU-Spot-VMs verwenden möchten, ersetzen Sie --on-demand durch --spot.

Standardnetzwerk

Wenn Sie kein Netzwerk mit hoher MTU benötigen, können Sie einen Cluster erstellen, der das Standard-VPC-Netzwerk verwendet. Mit dem folgenden Befehl wird On-Demand-Kapazität bereitgestellt:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

Wenn Sie reservierte Kapazität verwenden möchten, ersetzen Sie --on-demand durch --reservation=RESERVATION_NAME. Wenn Sie TPU-Spot-VMs verwenden möchten, ersetzen Sie --on-demand durch --spot.

Hinweis: Mit dem Flag --cluster-cpu-machine-type wird der Maschinentyp für die CPU-Host-VMs im Cluster angegeben. e2-standard-8 ist zwar ein empfohlener Standard für viele Arbeitslasten, Sie können jedoch einen anderen Maschinentyp auswählen, wenn Ihre Arbeitslast bestimmte CPU- oder Arbeitsspeicheranforderungen für Aufgaben wie die Datenvorverarbeitung hat.

Cluster mit Unterstützung für mehrere NICs mit XPK erstellen

Für Multislice-Arbeitslasten oder andere Arbeitslasten, die eine hohe Netzwerkbandbreite erfordern, z. B. für die Datenaufnahme, können Sie die Leistung mit mehreren NICs verbessern. Wenn Sie mit mehreren NICs arbeiten, werden jeder TPU-VM zusätzliche Netzwerkschnittstellen zugewiesen, die jeweils mit einem eindeutigen VPC-Netzwerk verbunden sind. Dadurch wird der Gesamtdurchsatz des Netzwerks erhöht. Für die meisten Single-Slice-Arbeitslasten bietet Single-NIC eine ausreichende Leistung bei geringerem Konfigurationsaufwand.

Hinweis: Die beiden Subnetzwerke müssen sich gegenseitig ausschließende IP-Bereiche haben. Dies umfasst sowohl die primären als auch die sekundären IP-CIDR-Bereiche im Subnetzwerk.

In den folgenden Abschnitten wird gezeigt, wie Sie mit XPK einen GKE-Cluster mit Unterstützung für mehrere NICs erstellen.

XPK installieren und Umgebungsvariablen einrichten

1. Installieren Sie XPK. Folgen Sie der Anleitung im XPK-GitHub-Repository.

2. Richten Sie Umgebungsvariablen für Ihren Cluster und Ihr primäres Netzwerk ein:

   export CLUSTER_NAME=XPK_CLUSTER_NAME
   export REGION=REGION
   export ZONE=us-east1-d
   export PROJECT_ID=PROJECT_ID
   export ACCELERATOR_TYPE=ACCELERATOR_TYPE
   export NUM_SLICES=2
   
   export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
   export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
   export FIREWALL_RULE_NAME_1=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
   export ROUTER_NAME_1=${CLUSTER_NAME}-network-1-${ZONE}
   export NAT_CONFIG_1=${CLUSTER_NAME}-natconfig-1-${ZONE}
   
   export NETWORK_NAME_2=${CLUSTER_NAME}-mtu9k-2-${ZONE}
   export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
   export FIREWALL_RULE_NAME_2=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
   export ROUTER_NAME_2=${CLUSTER_NAME}-network-2-${ZONE}
   export NAT_CONFIG_2=${CLUSTER_NAME}-natconfig-2-${ZONE}

   Legen Sie die folgenden Umgebungsvariablen fest:

   • CLUSTER_NAME: Ein Name für Ihren Cluster.
   • REGION: Die Region, in der Ihr TPU-Cluster erstellt wird.
   • ZONE: Die Zone, in der der TPU-Cluster erstellt wird. Weitere Informationen zu unterstützten Zonen finden Sie unter Regionen und Zonen.
   • PROJECT_ID: Ihre Google Cloud Projekt-ID
   • ACCELERATOR_TYPE: Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Beispiel: v6e-256. Weitere Informationen zu den unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
   • NUM_SLICES: Die Anzahl der TPU-Slices für Ihren Cluster. Jedes Slice hat die in ACCELERATOR_TYPE angegebene Anzahl von Chips. Legen Sie für einen Cluster mit einem einzelnen Slice NUM_SLICES auf 1 fest. Geben Sie für einen Multislice-Cluster die Anzahl der Slices basierend auf den Skalierbarkeitsanforderungen Ihrer Arbeitslast an. Die Gesamtzahl der Chips im Cluster ist die Anzahl der Chips in ACCELERATOR_TYPE multipliziert mit NUM_SLICES.

Primäre Netzwerkressourcen erstellen

1. Erstellen Sie das primäre Netzwerk mit einer maximalen Übertragungseinheit (MTU) von 8.896:

   gcloud compute networks create ${NETWORK_NAME_1} \
       --mtu=8896 \
       --bgp-routing-mode=regional \
       --subnet-mode=custom \
       --project=${PROJECT_ID}

   Ein benutzerdefiniertes Netzwerk mit einer MTU von 8.896 bietet eine bessere Leistung. Weitere Informationen finden Sie unter MTU konfigurieren.

2. Erstellen Sie das primäre Subnetzwerk:

   gcloud compute networks subnets create ${SUBNET_NAME_1} \
       --network=${NETWORK_NAME_1} \
       --range=10.11.0.0/18 \
       --region=${REGION} \
       --project=${PROJECT_ID}

3. Erstellen Sie eine Firewallregel für das primäre Netzwerk, die tcp-, icmp- und udp-Traffic im primären Netzwerk zulässt:

   gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_1} \
       --network=${NETWORK_NAME_1} \
       --allow tcp,icmp,udp \
       --project=${PROJECT_ID}

4. Erstellen Sie einen Cloud Router für das primäre Netzwerk:

   gcloud compute routers create ${ROUTER_NAME_1} \
       --project=${PROJECT_ID} \
       --network=${NETWORK_NAME_1} \
       --region=${REGION}

5. Konfigurieren Sie NAT für das primäre Netzwerk. Mit dem folgenden Befehl kann Traffic von Ihrem Cluster das Internet erreichen:

   gcloud compute routers nats create ${NAT_CONFIG_1} \
       --router=${ROUTER_NAME_1} \
       --region=${REGION} \
       --auto-allocate-nat-external-ips \
       --nat-all-subnet-ip-ranges \
       --project=${PROJECT_ID} \
       --enable-logging

Sekundäre Netzwerkressourcen erstellen

1. Erstellen Sie das sekundäre Netzwerk:

   gcloud compute networks create ${NETWORK_NAME_2} --mtu=8896 \
       --bgp-routing-mode=regional \
       --subnet-mode=custom \
       --project=${PROJECT_ID}
   

2. Erstellen Sie ein Subnetzwerk für das sekundäre Netzwerk:

   gcloud compute networks subnets create ${SUBNET_NAME_2} \
       --network=${NETWORK_NAME_2} \
       --range=10.10.0.0/18 \
       --region=${REGION} \
       --project=${PROJECT_ID}
   

3. Erstellen Sie eine Firewallregel, die Traffic innerhalb des neuen Netzwerks zulässt:

   gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_2} \
       --network=${NETWORK_NAME_2} \
       --allow tcp,icmp,udp \
       --source-ranges 10.10.0.0/18 \
       --project=${PROJECT_ID}
   

4. Erstellen Sie einen Cloud Router für das sekundäre Netzwerk:

   gcloud compute routers create ${ROUTER_NAME_2} \
       --project=${PROJECT_ID} \
       --network=${NETWORK_NAME_2} \
       --region=${REGION}
   

5. Erstellen Sie eine NAT-Konfiguration für den Cloud Router:

   gcloud compute routers nats create ${NAT_CONFIG_2} \
       --router=${ROUTER_NAME_2} \
       --region=${REGION} \
       --auto-allocate-nat-external-ips \
       --nat-all-subnet-ip-ranges \
       --project=${PROJECT_ID} \
       --enable-logging
   

Cluster erstellen

1. Legen Sie eine Umgebungsvariable für die Cluster- und Knotenpoolargumente fest, um die von Ihnen erstellten Netzwerke und Subnetzwerke zu verwenden:

   export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
   export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"
   

   Mit diesen Argumenten wird der Cluster so konfiguriert, dass die beiden Netzwerke verwendet werden, die Sie für die Unterstützung mehrerer NICs erstellt haben.

2. Erstellen Sie den Cluster. Mit dem folgenden Befehl wird On-Demand-Kapazität bereitgestellt:

   xpk cluster create \
       --cluster=${CLUSTER_NAME} \
       --cluster-cpu-machine-type=e2-standard-8 \
       --num-slices=${NUM_SLICES} \
       --tpu-type=${ACCELERATOR_TYPE} \
       --zone=${ZONE}  \
       --project=${PROJECT_ID} \
       --on-demand \
       --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
       --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
       --create-vertex-tensorboard

   Wenn Sie reservierte Kapazität verwenden möchten, ersetzen Sie --on-demand durch --reservation=RESERVATION_NAME. Wenn Sie TPU-Spot-VMs verwenden möchten, ersetzen Sie --on-demand durch --spot.

   Hinweis: Mit dem Flag --cluster-cpu-machine-type wird der Maschinentyp für die CPU-Host-VMs im Cluster angegeben. e2-standard-8 ist zwar ein empfohlener Standard für viele Arbeitslasten, Sie können jedoch einen anderen Maschinentyp auswählen, wenn Ihre Arbeitslast bestimmte CPU- oder Arbeitsspeicheranforderungen für Aufgaben wie die Datenvorverarbeitung hat.

Einrichtung mit mehreren NICs validieren

Nachdem Sie einen Cluster mit Unterstützung für multi-NIC erstellt haben, können Sie prüfen, ob beide NICs verwendet werden. Erstellen Sie dazu eine XPK-Arbeitslast und fügen Sie das Flag --command ifconfig hinzu.

1. Verwenden Sie den folgenden Befehl, um die Ausgabe des Befehls ifconfig in den Logs der Console vonGoogle Cloud anzuzeigen. Sie müssen entweder das Flag --base-docker-image maxtext_base_image angeben, um das MaxText-Basis-Image zu verwenden, wie im folgenden Beispiel, oder das Flag --docker-image und das gewünschte Image.

   xpk workload create \
       --cluster ${CLUSTER_NAME} \
       --base-docker-image maxtext_base_image \
       --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
       --tpu-type=${ACCELERATOR_TYPE} \
       --num-slices=${NUM_SLICES}  \
       --on-demand \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --command "ifconfig"

   Wenn Sie Debugging-Logs aktivieren oder Vertex AI TensorBoard verwenden möchten, fügen Sie dem Befehl die folgenden optionalen Argumente hinzu:

   --enable-debug-logs \
   --use-vertex-tensorboard

2. Prüfen Sie, ob die MTU für „eth0“ und „eth1“ auf 8.896 festgelegt ist. Sehen Sie dazu in der Ausgabe der XPK-Arbeitslast in den Logs der Console von Google Cloud nach.

JAX oder PyTorch einrichten

In den folgenden Ressourcen wird beschrieben, wie Sie JAX oder PyTorch auf Ihrer TPU einrichten, je nachdem, welche Bereitstellungs- und Verwaltungsmethode Sie verwenden:

• GKE Autopilot: TPU-Anwendung vorbereiten
• GKE Standard: Arbeitslastvorbereitung
• GKE und XPK: XPK-README
• Cloud TPU mit einem Host mit JAX: Berechnung mit JAX auf einer Cloud TPU-VM ausführen
• Cloud TPU mit mehreren Hosts mit JAX: JAX-Code auf TPU-Slices ausführen
• Cloud TPU mit PyTorch: PyTorch auf Cloud TPU

Informationen zum Einrichten und Ausführen von XPK mit MaxText finden Sie unter Running MaxText at Scale with XPK.

SkyPilot verwenden

Sie können Cloud TPU v6e mit SkyPilot verwenden. SkyPilot ist ein Open-Source-Framework, das das Ausführen, Verwalten und Skalieren von KI-Arbeitslasten vereinfacht. Sie können SkyPilot v6e-bezogene Standort- und Preisinformationen hinzufügen. Weitere Informationen finden Sie im SkyPilot-Beispiel für TPU v6e.

Trainingsbeispiele

In den folgenden Abschnitten finden Sie Beispiele für das Training von MaxText-, MaxDiffusion- und PyTorch-Modellen auf Cloud TPU v6e.

Diese Beispiele wurden mit den folgenden Softwareversionen getestet:

• Python 3.10 oder höher
• Nightly-Softwareversionen:
  • Nightly JAX 0.4.32.dev20240912
  • Nightly LibTPU 0.1.dev20240912+nightly
• Stabile Softwareversionen:
  • JAX + jaxlib v0.4.37

MaxText und MaxDiffusion auf Cloud TPU v6e trainieren

In den folgenden Abschnitten wird der Trainingslebenszyklus der MaxText- und MaxDiffusion-Modelle beschrieben.

Es müssen allgemein folgende Schritte ausgeführt werden:

1. Basis-Image für die Arbeitslast erstellen.
2. Arbeitslast mit XPK ausführen.
   1. Trainingsbefehl für die Arbeitslast erstellen.
   2. Arbeitslast bereitstellen.
3. Arbeitslast verfolgen und Messwerte ansehen.
4. XPK-Arbeitslast löschen, wenn sie nicht benötigt wird.
5. Löschen Sie den Cluster, wenn er nicht mehr benötigt wird.

Basis-Image erstellen

Installieren Sie MaxText oder MaxDiffusion und erstellen Sie das Docker-Image:

1. Klonen Sie das gewünschte Repository und wechseln Sie zum Verzeichnis des Repositorys:

   MaxText:

   git clone https://github.com/google/maxtext.git
   cd maxtext
   git checkout maxtext-v0.2.0
   

   MaxDiffusion:

   git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
   

2. Konfigurieren Sie Docker für die Verwendung der Google Cloud CLI:

   gcloud auth configure-docker
   

3. Erstellen Sie das Docker-Image mit dem folgenden Befehl oder mit einem JAX KI-Bild. Weitere Informationen zu JAX KI-Bildern finden Sie unter JAX KI-Bildern.

   MaxText:

   bash dependencies/scripts/docker_build_dependency_image.sh MODE=stable
   

   MaxDiffusion:

   bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
   

4. Legen Sie Ihre Projekt-ID in der aktiven gcloud CLI-Konfiguration fest:

   gcloud config set project ${PROJECT_ID}
   

5. Wenn Sie die Arbeitslast von einem Computer aus starten, auf dem das Image nicht lokal erstellt wurde, laden Sie das Image hoch.

   Hinweis: Sie benötigen die Bash-Version 4.2 oder höher.

   1. Legen Sie die Umgebungsvariable CLOUD_IMAGE_NAME fest:

      export CLOUD_IMAGE_NAME=${USER}_runner
      

   2. Laden Sie das Image hoch:

      bash dependencies/scripts/docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
      

Arbeitslast mit XPK ausführen

1. Legen Sie die folgenden Umgebungsvariablen fest, wenn Sie nicht die Standardwerte von MaxText oder MaxDiffusion verwenden:

   export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
   export PER_DEVICE_BATCH_SIZE=2
   export NUM_STEPS=30
   export MAX_TARGET_LENGTH=8192

2. Erstellen Sie Ihr Modellscript. Dieses Script wird in einem späteren Schritt als Trainingsbefehl kopiert.

   Führen Sie das Modellscript noch nicht aus.

   Hinweis: Die automatische Erstellung von Prüfpunkten ist nicht aktiviert. Wenn Fehler wie Could not find a known environment for initializing distributed JAX angezeigt werden, verwenden Sie JAX_FORCE_TPU_INIT=true.

   MaxText

   MaxText ist ein leistungsstarkes, hochgradig skalierbares Open-Source-LLM, das in reinem Python und JAX geschrieben und auf TPUs von Google Cloud und GPUs für Training und Inferenz ausgerichtet ist.

   JAX_PLATFORMS=tpu,cpu \
   ENABLE_PJRT_COMPATIBILITY=true \
   TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
   TPU_SLICE_BUILDER_DUMP_ICI=true && \
   python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
        enable_checkpointing=false \
        gcs_metrics=true \
        profiler=xplane \
        skip_first_n_steps_for_profiler=5 \
        steps=${NUM_STEPS}  # attention='dot_product'"
   

   Gemma2

   Gemma ist eine LLM-Reihe mit offenen Gewichtungen, die von Google DeepMind entwickelt wurde und auf der Forschung und Technologie von Gemini basiert.

   Hinweis: Um dieses Modell auszuführen, müssen Sie eine v6e-256- oder größere TPU bereitstellen.

   python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
       model_name=gemma2-27b \
       run_name=gemma2-27b-run \
       base_output_directory=${BASE_OUTPUT_DIR} \
       max_target_length=${MAX_TARGET_LENGTH} \
       per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
       steps=${NUM_STEPS} \
       enable_checkpointing=false \
       use_iota_embed=true \
       gcs_metrics=true \
       dataset_type=synthetic \
       profiler=xplane \
       attention=flash
   

   Mixtral 8x7B

   Mixtral ist ein hochmodernes KI-Modell, das von Mistral AI entwickelt wurde und eine dünnbesetzte MoE-Architektur (Mixture of Experts) nutzt.

   Hinweis: Um dieses Modell auszuführen, müssen Sie eine v6e-256- oder größere TPU bereitstellen.

   python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
       base_output_directory=${BASE_OUTPUT_DIR} \
       per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
       model_name=mixtral-8x7b \
       steps=${NUM_STEPS} \
       max_target_length=${MAX_TARGET_LENGTH} \
       tokenizer_path=assets/tokenizer.mistral-v1 \
       attention=flash \
       dtype=bfloat16 \
       dataset_type=synthetic \
       profiler=xplane
   

   MaxDiffusion

   MaxDiffusion ist eine Sammlung von Referenzimplementierungen verschiedener latenter Diffusionsmodelle, die in reinem Python und JAX geschrieben sind und auf XLA-Geräten ausgeführt werden, einschließlich von Cloud TPUs und GPUs. Stable Diffusion ist ein latentes Text-zu-Bild-Modell, das fotorealistische Bilder aus beliebigen Texteingaben generiert.

   Sie müssen einen bestimmten Git-Branch installieren, um MaxDiffusion auszuführen, wie im folgenden Trainingsscript gezeigt.

   git clone https://github.com/google/maxdiffusion.git
   && cd maxdiffusion
   && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
   && pip install -r requirements.txt && pip install .
   && pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR}
   && python src/maxdiffusion/train_sdxl.py \
       src/maxdiffusion/configs/base_xl.yml \
       revision=refs/pr/95 \
       activations_dtype=bfloat16 \
       weights_dtype=bfloat16 \
       resolution=1024 \
       per_device_batch_size=1 \
       output_dir=${OUT_DIR} \
       jax_cache_dir=${OUT_DIR}/cache_dir/ \
       max_train_steps=200 \
       attention=flash \
       run_name=sdxl-ddp-v6e
   

   Hinweis: Optional können Sie dcn_data_parallelism für Multislice-Ausführungen festlegen.

3. Exportieren Sie die folgenden Variablen:

   export CLUSTER_NAME=CLUSTER_NAME
   export ACCELERATOR_TYPE=ACCELERATOR_TYPE
   export NUM_SLICES=NUM_SLICES
   export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT

   Beschreibungen von Umgebungsvariablen

   • CLUSTER_NAME ist der Name Ihres Clusters.
   • ACCELERATOR_TYPE: Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
   • NUM_SLICES: Die Anzahl der TPU-Slices.
   • YOUR_MODEL_SCRIPT: Das Modellskript, das als Trainingsbefehl ausgeführt werden soll.

4. Führen Sie das Modell mit dem Script aus, das Sie im vorherigen Schritt erstellt haben. Sie müssen entweder das Flag --base-docker-image angeben, um das MaxText-Basis-Image zu verwenden, oder das Flag --docker-image und das gewünschte Image.

   Sie können die folgenden optionalen Flags hinzufügen:

   • Sie können das Debugging-Logging aktivieren, indem Sie das Flag --enable-debug-logs einfügen. Weitere Informationen finden Sie unter JAX in MaxText debuggen.
   • Sie können ein Vertex AI-Experiment erstellen, um Daten in Vertex AI TensorBoard hochzuladen. Dazu müssen Sie das Flag --use-vertex-tensorboard einfügen. Weitere Informationen finden Sie unter JAX mit Vertex AI auf MaxText überwachen.

   xpk workload create \
     --cluster ${CLUSTER_NAME} \
     {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \
     --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
     --tpu-type=${ACCELERATOR_TYPE} \
     --num-slices=${NUM_SLICES}  \
     --on-demand \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --command="${YOUR_MODEL_SCRIPT}"

   Die Ausgabe enthält einen Link, über den Sie Ihre Arbeitslast verfolgen können. Öffnen Sie den Link und klicken Sie auf den Tab Logs, um Ihre Arbeitslast in Echtzeit zu verfolgen.

JAX in MaxText debuggen

Verwenden Sie zusätzliche XPK-Befehle, um herauszufinden, warum der Cluster oder die Arbeitslast nicht ausgeführt wird:

• XPK workload list
• XPK inspector
• Aktivieren Sie das ausführliche Logging in Ihren Arbeitslastlogs mit dem Flag --enable-debug-logs, wenn Sie die XPK-Arbeitslast erstellen.

JAX mit Vertex AI auf MaxText überwachen

Damit Sie TensorBoard verwenden können, muss Ihrem Nutzerkonto in Google Cloud die Rolle aiplatform.user zugewiesen sein. Führen Sie den folgenden Befehl aus, um diese Rolle zuzuweisen:

gcloud projects add-iam-policy-binding your-project-id \
   --member='user:your-email' \
   --role='roles/aiplatform.user'

So können Sie skalare und Profildaten über das von Vertex AI verwaltete TensorBoard ansehen:

1. Erhöhen Sie die Resource-Management-Anfragen (CRUD) für die Zone, die Sie verwenden, von 600 auf 5.000. Bei kleinen Arbeitslasten mit weniger als 16 VMs ist das möglicherweise kein Problem.

2. Installieren Sie Abhängigkeiten wie cloud-accelerator-diagnostics für Vertex AI:

   # xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
   cd ~/xpk
   pip install .

3. Erstellen Sie Ihren Cluster mit dem Flag --create-vertex-tensorboard, wie in der XPK-README unter Create Vertex AI TensorBoard beschrieben. Sie können diesen Befehl auch auf vorhandenen Clustern ausführen.

4. Erstellen Sie Ihr Vertex AI-Experiment mit dem Flag --use-vertex-tensorboard und dem optionalen Flag --experiment-name, wenn Sie Ihre XPK-Arbeitslast ausführen. Eine vollständige Liste der Schritte finden Sie in der XPK-README unter Create Vertex AI Experiment to upload data to Vertex AI Tensorboard.

Hinweis: Wenn die von Ihnen ausgeführte Arbeitslast keine MaxText-Arbeitslast ist, verwenden Sie dieses MaxText-Beispiel, um die entsprechenden Änderungen am Datenupload vorzunehmen.

Die Logs enthalten einen Link zu einem Vertex AI TensorBoard, der in etwa so aussieht:

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

Sie finden den Link zum Vertex AI TensorBoard auch in der Console von Google Cloud . Rufen Sie Vertex AI Experiments in der Console von Google Cloud auf. Wählen Sie im Drop-down-Menü die gewünschte Region aus.

Das TensorBoard-Verzeichnis wird auch in den Cloud Storage-Bucket geschrieben, den Sie mit ${BASE_OUTPUT_DIR} angegeben haben.

XPK-Arbeitslast löschen

Verwenden Sie den Befehl xpk workload delete, um eine oder mehrere Arbeitslasten basierend auf dem Jobpräfix oder dem Jobstatus zu löschen. Dieser Befehl kann nützlich sein, wenn Sie XPK-Arbeitslasten gesendet haben, die nicht mehr ausgeführt werden müssen, oder wenn Jobs in der Warteschlange hängen.

Cluster löschen

Verwenden Sie den Befehl xpk cluster delete, um den Cluster zu löschen:

xpk cluster delete --cluster ${CLUSTER_NAME} \
    --zone=${ZONE} --project=${PROJECT_ID}

Benchmark-Ergebnisse für MaxDiffusion

Wir haben das Trainingsscript für MaxDiffusion auf einer v6e-4, einer v6e-16 und zwei v6e-16 ausgeführt. In der folgenden Tabelle sehen Sie die gemessenen Durchsätze.

	v6e-4	v6e-16	Zwei v6e-16
Trainingsschritte	0,069	0,073	0,13
Globale Batchgröße	8	32	64
Durchsatz (Beispiele/s)	115,9	438,4	492,3

Modell mit PyTorch auf Cloud TPU trainieren

Sie können TorchTPU direkt zum Trainieren Ihres Modells verwenden oder die Referenzbibliothek TorchTitan wiederverwenden. Informationen zu den ersten Schritten mit v6e-Arbeitslasten finden Sie in den TorchTitan Trillium-Anleitungen. Weitere Informationen zum PyTorch-Ökosystem in Cloud TPU finden Sie unter PyTorch in Cloud TPU.