Modell mit TPU v6e trainieren

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/XLA.

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:

Vorbereitung

Zur Vorbereitung sind folgende Schritte erforderlich:

Konto und Projekt in Google Cloud mit aktivierter Abrechnung erstellen
Google Cloud CLI-Alphakomponenten installieren
Cloud TPU API aktivieren
Cloud TPU-Dienst-Agent erstellen
Cloud TPU-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:

Kontingent auf Abruf oder On-Demand-Kontingent für TPU v6e
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 in der XPK-README unter Erforderliche Berechtigungen für die Console von Google Cloud .

Bereitstellungsoptionen

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

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.
In die Warteschlange gestellte Cloud TPU-Ressourcen: Sie können TPU-Kapazität anfordern, die bereitgestellt wird, sobald sie verfügbar ist. Diese Möglichkeit eignet sich ideal für Batchjobs und fehlertolerante Arbeitslasten, die in einer Warteschlange warten können. Sie können ein Zeitfenster für Ihre Anfrage angeben. Weitere Informationen finden Sie unter In die Warteschlange gestellte Ressourcen verwalten.

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, mit denen Sie den Durchsatz über mehrere Netzwerke hinweg skalieren können. In den folgenden Abschnitten finden Sie Befehle zum Erstellen eines GKE-Clusters mit Unterstützung für eine einzelne NIC oder für mehrere NICs mit XPK. Für die meisten Single-Slice-Arbeitslasten bietet Single-NIC eine ausreichende Leistung bei weniger Konfiguration. Verwenden Sie multi-NIC für Multislice-Arbeitslasten und Arbeitslasten, die hohe Datenaufnahmegeschwindigkeiten erfordern.

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

Für die meisten Arbeitslasten mit einem einzelnen Slice bietet eine einzelne NIC mit weniger Konfiguration eine ausreichende Leistung. Verwenden Sie multi-NIC für Multislice-Arbeitslasten und Arbeitslasten, die hohe Datenaufnahmegeschwindigkeiten erfordern.

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

XPK installieren und Umgebungsvariablen einrichten

Installieren Sie XPK. Folgen Sie der Anleitung im XPK-GitHub-Repository.
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. Jeder Slice enthält 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:

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.

So 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

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}

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}"
```

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

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 Datenerfassung, können Sie Multi-NIC verwenden, um die Leistung zu verbessern. Wenn Sie mehrere NICs verwenden, werden jeder TPU-VM zusätzliche Netzwerkschnittstellen zugewiesen, die jeweils mit einem eindeutigen VPC-Netzwerk verbunden sind. Dadurch wird der gesamte Netzwerkdurchsatz erhöht. Für die meisten Arbeitslasten mit einem einzelnen Slice bietet eine einzelne NIC mit weniger Konfiguration eine ausreichende Leistung.

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

XPK installieren und Umgebungsvariablen einrichten

Installieren Sie XPK. Folgen Sie der Anleitung im XPK-GitHub-Repository.
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. Jeder Slice enthält 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

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}
```
Die Verwendung eines benutzerdefinierten Netzwerks mit einer MTU von 8.896 bietet eine bessere Leistung. Weitere Informationen finden Sie unter MTU konfigurieren.

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}

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}

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}

Konfigurieren Sie NAT für das primäre Netzwerk. Der folgende Befehl ermöglicht es, dass Traffic von Ihrem Cluster das Internet erreicht:

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

Ressourcen für das sekundäre Netzwerk erstellen

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}

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}

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}

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}

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

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 für die Verwendung der beiden Netzwerke konfiguriert, die Sie für die Unterstützung mehrerer NICs erstellt haben.
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 für viele Arbeitslasten die empfohlene Standardeinstellung, Sie können aber einen anderen Maschinentyp auswählen, wenn Ihre Arbeitslast bestimmte CPU- oder Arbeitsspeicheranforderungen für Aufgaben wie die Datenvorverarbeitung hat.

Multi-NIC-Einrichtung validieren

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

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
```
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 einem Host mit PyTorch: Berechnung mit PyTorch auf einer Cloud TPU-VM ausführen
Cloud TPU mit mehreren Hosts mit PyTorch: PyTorch-Code auf TPU-Slices ausführen

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

TCP-Einstellungen verbessern

Wenn Sie Ihre v6e-TPUs mit in die Warteschlange gestellten Ressourcen bereitgestellt haben, können Sie die Netzwerkleistung verbessern, indem Sie die Limits für den TCP-Empfangszwischenspeicher erhöhen. Führen Sie dazu den folgenden Befehl aus.

gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \
    --project "${PROJECT_ID}" \
    --zone "${ZONE}" \
    --node=all \
    --worker=all \
    --command='
    sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"'

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:

Basis-Image für die Arbeitslast erstellen.
Arbeitslast mit XPK ausführen.
1. Trainingsbefehl für die Arbeitslast erstellen.
2. Arbeitslast bereitstellen.
Arbeitslast verfolgen und Messwerte ansehen.
XPK-Arbeitslast löschen, wenn sie nicht benötigt wird.
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:

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

MaxText:

git clone https://github.com/google/maxtext.git && cd maxtext

MaxDiffusion:

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

Konfigurieren Sie Docker für die Verwendung der Google Cloud CLI:
```
gcloud auth configure-docker
```

Erstellen Sie das Docker-Image mit dem folgenden Befehl oder mit einem JAX AI-Image. Weitere Informationen zu JAX AI-Images finden Sie unter JAX AI-Images.

MaxText:

bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35

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

Legen Sie Ihre Projekt-ID in der aktiven gcloud CLI-Konfiguration fest:
```
gcloud config set project ${PROJECT_ID}
```
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 docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
```

Arbeitslast mit XPK ausführen

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

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

Führen Sie das Modellscript noch nicht aus.

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

python3 -m MaxText.train 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.

python3 -m MaxText.train 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

Llama3-8b

Llama ist eine LLM-Reihe mit offenen Gewichtungen, die von Meta entwickelt wurden.

Ein Beispiel für die Ausführung von Llama3 in PyTorch finden Sie in den torch_xla-Modellen im torchprime-Repository.

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

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 Modellscript, das als Trainingsbefehl ausgeführt werden soll.
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:

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.

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 .

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

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

Llama-Modelle mit PyTorch/XLA auf Cloud TPU v6e trainieren

In diesem Abschnitt wird beschrieben, wie Sie Llama-Modelle mit PyTorch/XLA auf Cloud TPU v6e mit dem WikiText-Dataset trainieren.

Zugriff auf Hugging Face und das Llama 3-Modell erhalten

Für dieses Beispiel benötigen Sie ein Hugging Face-Nutzerzugriffstoken. Informationen zum Erstellen von Nutzerzugriffstokens finden Sie in der Hugging Face-Dokumentation zu Nutzerzugriffstokens.

Außerdem benötigen Sie die Berechtigung für den Zugriff auf das Modell „Llama-3-8B“ auf Hugging Face. Rufen Sie dazu das Meta-Llama-3-8B-Modell auf Hugging Face auf und beantragen Sie den Zugriff.

Cloud TPU-VM erstellen

Erstellen Sie für dieses Beispiel eine Cloud TPU v6e mit 8 Chips.

Richten Sie Umgebungsvariablen ein:
```
export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east1-d
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e
```
Beschreibungen von Umgebungsvariablen
- PROJECT_ID: Ihre Google Cloud Projekt-ID. Verwenden Sie ein vorhandenes Projekt oder erstellen Sie ein neues Projekt.
- TPU_NAME: Der Name der TPU.
- ZONE: Die Zone, in der die TPU-VM erstellt werden soll. Weitere Informationen zu unterstützten Zonen finden Sie unter TPU-Regionen und ‑Zonen.
- 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.
- RUNTIME_VERSION: Die Cloud TPU-Softwareversion.

Erstellen Sie eine Cloud TPU-VM:

gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${RUNTIME_VERSION} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID}

Installation

Installieren Sie den pytorch-tpu/transformers-Fork von Hugging Face Transformers und die Abhängigkeiten. Dieses Beispiel wurde mit den folgenden Abhängigkeitsversionen getestet:

torch: kompatibel mit 2.5.0
torch_xla[tpu]: kompatibel mit 2.5.0
jax: 0.4.33
jaxlib: 0.4.33

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
   cd transformers
   sudo pip3 install -e .
   pip3 install datasets
   pip3 install evaluate
   pip3 install scikit-learn
   pip3 install accelerate
   pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
   pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/'

Modellkonfigurationsdateien einrichten

Im Trainingsbefehl im nächsten Abschnitt, Modell ausführen, werden zwei JSON-Konfigurationsdateien verwendet, um Modellparameter und die FSDP-Konfiguration zu definieren. Mit der FSDP-Fragmentierung (Fully Sharded Data Parallelism, vollständig fragmentierte Datenparallelität) können Sie beim Training eine größere Batchgröße verwenden, indem die Modellgewichtungen auf mehrere TPUs verteilt werden. Beim Training mit kleineren Modellen kann es ausreichen, Datenparallelität zu verwenden und die Gewichtungen auf jedem Gerät zu replizieren. Weitere Informationen zur Fragmentierung von Tensoren über mehrere Geräte in PyTorch/XLA finden Sie im PyTorch/XLA SPMD User Guide.

Erstellen Sie die Modellparameter-Konfigurationsdatei. Im Folgenden finden Sie die Modellparameterkonfiguration für Llama-3-8B. Die Konfigurationsdatei für andere Modelle finden Sie auf Hugging Face. Ein Beispiel finden Sie in der Llama-2-7B-Konfiguration.

cat > llama-config.json << EOF
{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.40.0.dev0",
  "use_cache": false,
  "vocab_size": 128256
}
EOF

Erstellen Sie die FSDP-Konfigurationsdatei:

cat > fsdp-config.json << EOF
{
  "fsdp_transformer_layer_cls_to_wrap": [
    "LlamaDecoderLayer"
  ],
  "xla": true,
  "xla_fsdp_v2": true,
  "xla_fsdp_grad_ckpt": true
}
EOF

Weitere Informationen zu FSDP finden Sie in der PyTorch-Dokumentation unter Fully Sharded Data Parallel using SPMD.

Laden Sie die Konfigurationsdateien mit dem folgenden Befehl auf Ihre Cloud TPU-VMs hoch:

gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \
   --worker=all \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

Modell ausführen

Führen Sie mit den Konfigurationsdateien, die Sie im vorherigen Abschnitt erstellt haben, das Script run_clm.py aus, um das Llama-3-8B-Modell mit dem WikiText-Dataset zu trainieren. Die Ausführung des Trainingsscripts dauert auf einer Cloud TPU v6e-8 etwa 10 Minuten.

Melden Sie sich mit dem folgenden Befehl auf Ihrer Cloud TPU bei Hugging Face an:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   pip3 install "huggingface_hub[cli]"
   huggingface-cli login --token HUGGING_FACE_TOKEN'

Führen Sie das Modelltraining aus:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_SPMD=1
   export ENABLE_PJRT_COMPATIBILITY=true
   # Optional variables for debugging:
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=100000
   # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
   export PROFILE_LOGDIR=PROFILE_PATH
   python3 transformers/examples/pytorch/language-modeling/run_clm.py \
     --dataset_name wikitext \
     --dataset_config_name wikitext-2-raw-v1 \
     --per_device_train_batch_size 16 \
     --do_train \
     --output_dir /home/$USER/tmp/test-clm \
     --overwrite_output_dir \
     --config_name /home/$USER/llama-config.json \
     --cache_dir /home/$USER/cache \
     --tokenizer_name meta-llama/Meta-Llama-3-8B \
     --block_size 8192 \
     --optim adafactor \
     --save_strategy no \
     --logging_strategy no \
     --fsdp "full_shard" \
     --fsdp_config /home/$USER/fsdp-config.json \
     --torch_dtype bfloat16 \
     --dataloader_drop_last yes \
     --flash_attention \
     --max_steps 20'

Fehlerbehebung bei PyTorch/XLA

Wenn Sie die optionalen Variablen für das Debugging im vorherigen Abschnitt festgelegt haben, wird das Profil für das Modell am Speicherort gespeichert, der in der Variable PROFILE_LOGDIR angegeben ist. Sie können die Datei xplane.pb, die sich an diesem Speicherort befindet, extrahieren und tensorboard verwenden, um die Profile im Browser anzusehen. Beachten Sie dazu die TensorBoard-Anleitung.

Wenn PyTorch/XLA nicht wie erwartet funktioniert, finden Sie im Leitfaden zur Fehlerbehebung Vorschläge zum Debuggen, Profiling und Optimieren Ihres Modells.

Modell mit TPU v6e trainieren

Vorbereitung

Kontingent und Berechtigungen prüfen

Bereitstellungsoptionen

v6e-TPUs mit GKE und XPK bereitstellen

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

XPK installieren und Umgebungsvariablen einrichten

Cluster erstellen

Benutzerdefinierte Netzwerke

Standardnetzwerk

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

XPK installieren und Umgebungsvariablen einrichten

Primäre Netzwerkressourcen erstellen

Ressourcen für das sekundäre Netzwerk erstellen

Cluster erstellen

Multi-NIC-Einrichtung validieren

JAX oder PyTorch einrichten

TCP-Einstellungen verbessern

SkyPilot verwenden

Trainingsbeispiele

MaxText und MaxDiffusion auf Cloud TPU v6e trainieren

Basis-Image erstellen

Arbeitslast mit XPK ausführen

MaxText

Gemma2

Mixtral 8x7B

Llama3-8b

MaxDiffusion

Beschreibungen von Umgebungsvariablen

JAX in MaxText debuggen

JAX mit Vertex AI auf MaxText überwachen

XPK-Arbeitslast löschen

Cluster löschen

Benchmark-Ergebnisse für MaxDiffusion

Llama-Modelle mit PyTorch/XLA auf Cloud TPU v6e trainieren

Zugriff auf Hugging Face und das Llama 3-Modell erhalten

Cloud TPU-VM erstellen

Beschreibungen von Umgebungsvariablen

Installation

Modellkonfigurationsdateien einrichten

Modell ausführen

Fehlerbehebung bei PyTorch/XLA