Cloud TPU v5p-Training

Cloud TPU v5p ist die Cloud TPU der fünften Generation von Google Cloud und der Nachfolger der TPU v4. v5p ist für die groß angelegte Modellierung optimiert und soll eine führende Plattform für die Entwicklung grundlegender LLMs, Diffusionsmodelle und generativer KI werden. Allgemein gesprochen, v5p bietet bis zu doppelt so hohe Leistung im Vergleich zu Version 4 und bietet gleichzeitig doppelt so viel Leistung. mehr TPUs in einem Pod (6.000 größtes Slice gegenüber 3.000 in v4) bis zu 4-fache Leistung auf Pod-Ebene. Außerdem läuft sie mit einer höheren Taktfrequenz (1,75 GHz gegenüber 1,05 GHz), bietet SparseCore für groß angelegte Einbettungen und verdreifacht die HBM-Kapazität (High Bandwidth Memory).

Cloud TPU v5p-Konzepte

Wenn Sie noch nicht mit Cloud TPUs vertraut sind, finden Sie hier weitere Informationen: TPU-Dokumentation – Startseite

Cloud TPU-Konzepte (z. B. Slices, Hosts und TensorCores) und die Cloud TPU-Systemarchitektur für alle Cloud TPU-Versionen werden auf der Seite Cloud TPU-Systemarchitektur beschrieben.

Jede Cloud TPU-Version erfordert bestimmte Beschleunigertypen für Training oder Inferenz. Diese Beschleunigertypen werden in v5p-Konfigurationen beschrieben.

TPU-Ressourcen verwalten

Alle Befehle, die Sie zum Verwalten Ihrer TPU-VMs verwenden können, werden unter TPUs verwalten oder Nutzerhandbuch für in der Warteschlange befindliche Ressourcen zur Verwaltung Ressourcen in der Warteschlange.

Framework-Einrichtung

In diesem Abschnitt wird der allgemeine Einrichtungsprozess für das Modell beschrieben. mit JAX oder PyTorch mit TPU v5p trainieren.

Einrichtung für JAX

Wenn Sie Slice-Formen mit mehr als 4 Chips haben, befinden sich mehrere VMs in einem Slice. In diesem Fall müssen Sie das Flag --worker=all verwenden, um die Installation mit einem einzigen Befehl auf allen TPU-VMs auszuführen:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Mit dem folgenden Befehl können Sie die Anzahl der Geräte prüfen. Die hier gezeigten Ausgabedaten wurden mit einem v5p-32-Speicherbereich erstellt. Mit diesem Code wird geprüft, ob alles korrekt installiert ist. Dazu wird überprüft, ob JAX die Cloud TPU-TensorCores erkennt und grundlegende Vorgänge ausführen kann:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Die Ausgabe sollte in etwa so aussehen:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() zeigt die Gesamtzahl der Chips in das angegebene Slice. jax.local_device_count() gibt die Anzahl der Chips, die für eine einzelne VM in diesem Slice zugänglich sind.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

Die Ausgabe sollte in etwa so aussehen:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Verwenden Sie --node=all, um den Befehl für alle Multislice-Worker auszuführen.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

In den JAX-Anleitungen in diesem Dokument mit v5p-Training mit JAX begonnen.

PyTorch einrichten

PJRT-Laufzeit ist die einzige unterstützte Laufzeit für v5p und PyTorch 2.1+ verwendet PJRT als Standardlaufzeit für alle TPU-Versionen. In diesem Abschnitt wird Folgendes beschrieben: Hier erfahren Sie, wie Sie PJRT auf v5p-Pods mit PyTorch/XLA 2.2.0 für alle Worker verwenden.

Abhängigkeiten installieren

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Validierung der Installation mithilfe eines Python-Skripts mit PJRT um verfügbare TPU-Geräte anzuzeigen (die hier gezeigten Ausgaben wurden mit einem v5p-32-Slice).

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

Verwenden Sie --node=all, um den Befehl für alle Multislice-Worker auszuführen.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

Sehen Sie sich die PyTorch-Anleitungen in diesem Dokument an, um mit dem Training mit v5p mit PyTorch zu beginnen.

Überwachen und profilieren

Cloud TPU v5p unterstützt die Überwachung und das Profiling mit denselben Methoden wie frühere Cloud TPU-Generationen. Weitere Informationen zur Profilerstellung finden Sie unter Profilerstellung Ihres Modells mit Cloud TPU-Tools und unter Cloud TPU-VMs überwachen.

Trainingsanleitungen

Dieser Abschnitt konzentriert sich auf Trainingsanleitungen für ein einzelnes Slice. Sie können diese Anleitungen für das Multislice-Training anpassen, indem Sie den SSH-Befehlen das Flag --node=all hinzufügen. Weitere Informationen und Best Practices finden Sie in der Einführung mit mehreren Segmenten

Hinweis: Bevor Sie eine dieser Anleitungen ausführen, müssen Sie Ihre Cloud TPU einrichten. . Nach der Einrichtung Ihre TPU-Umgebung haben, können Sie die folgenden Anleitungen ausführen.

• Diffusion 2.1
• MaxText
• ResNet auf einem einzelnen Host (v5p)
• ResNet auf einem v5p mit mehreren Hosts

JAX-Anleitungen

Train Diffusion 2.1

In dieser Anleitung erfahren Sie, wie Sie das Modell HuggingFace nutzt Pokémon Dataset auf Cloud TPU v5p.

Das Stable Diffusion-Modell ist ein latentes Text-zu-Bild-Modell, das fotorealistische Bilder aus jeder Texteingabe. Weitere Informationen finden Sie in den folgenden Ressourcen:

• Stable Diffusion
• Codequelle für stabile Diffusion

Hinweis: Bevor Sie diese Anleitung ausführen können, müssen Sie Ihre Cloud TPU-Umgebung einrichten.

Einrichten

1. Richten Sie einen Storage-Bucket für die Modellausgabe ein.

   gcloud storage buckets create gs://your_bucket \
    --project=your_project \
    --location=us-east5-a

2. Erstellen Sie Umgebungsvariablen:

   export GCS_BUCKET_NAME=your-bucket
   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5p-32
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export SERVICE_ACCOUNT=your-service-account
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-qr-name
   export QUOTA_TYPE=spot
   export VALID_UNTIL_DURATION=1d

   Beschreibung der Befehls-Flags

   Variable	Beschreibung
   PROJECT_ID	Name des Google Cloud-Projekts
   ACCELERATOR_TYPE	Siehe TPU-Versionen Seite für Ihre TPU-Version.
   ZONE	Informationen zu den unterstützten Zonen finden Sie im Dokument TPU-Regionen und -Zonen.
   RUNTIME_VERSION	Verwenden Sie für v5p „v2-alpha-tpuv5“ für die RUNTIME_VERSION.
   SERVICE_ACCOUNT	Das ist die Adresse Ihres Dienstkontos. Sie finden sie in der Google Cloud Console unter „IAM“ -> „Dienstkonten“. Beispiel: [e-mail]
   TPU_NAME	Die vom Nutzer zugewiesene Text-ID der TPU, die erstellt wird, wenn wird die Ressourcenanfrage in der Warteschlange zugewiesen.
   QUEUED_RESOURCE_ID	Die vom Nutzer zugewiesene Text-ID der anstehenden Ressourcenanfrage. Weitere Informationen finden Sie unter Das Dokument Ressourcen in der Warteschlange finden Sie Informationen zu Ressourcen in der Warteschlange.
   QUOTA_TYPE	Kann reserved oder spot sein. Wenn keine dieser Optionen angegeben ist, wird der Wert für QUOTA_TYPE ist standardmäßig on-demand. Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.
   VALID_UNTIL_DURATION	Die Dauer, für die die Anfrage gültig ist. Weitere Informationen finden Sie unter Ressourcen in der Warteschlange finden Sie Informationen zur gültigen Dauer.

3. Erstellen Sie eine TPU-Ressource:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --valid-until-duration ${VALID_UNTIL_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

   Hinweis: Das Flag QUOTA_TYPE kann so lauten: reserved oder spot. Wenn keines davon angegeben ist, lautet der Standardwert on-demand. Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden, finden Sie unter Kontingente.

   Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die in die Warteschlange gestellte Ressource im Status ACTIVE befindet. Prüfen Sie den Status der Ressource in der Warteschlange mit dem folgenden Befehl:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project ${PROJECT_ID} --zone ${ZONE}

   Wenn sich die in die Warteschlange gestellte Ressource im Status ACTIVE befindet, gibt die Ausgabe sieht in etwa so aus:

   state: ACTIVE
   

4. Modell trainieren

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --project $PROJECT_ID --worker=all --command="
   git clone https://github.com/google/maxdiffusion
   cd maxdiffusion
   git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
   pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html # Install the latest version of JAX
   pip3 install -r requirements.txt
   pip3 install .
   export LIBTPU_INIT_ARGS=""
   python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run base_output_directory=gs://$GCS_BUCKET_NAME enable_profiler=False"

Bereinigen

Sie können Ihre TPU- und Ressourcenanfrage in der Warteschlange am Ende der Sitzung löschen oder Ressourcenanfragen in der Warteschlange entfernen, die den Status „FAILED“ (Fehlgeschlagen) haben. Löschen Sie die Segmente und anschließend die in die Warteschlange eingereihte Ressource, um eine Ressource in der Warteschlange zu löschen. Anfrage in 2 Schritten:

   gcloud compute tpus tpu-vm delete ${TPU_NAME} --project=${PROJECT_ID} \
      --zone=${ZONE} --quiet

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} --zone ${ZONE} --quiet

Mit --force können Sie die Slices und die anstehende Ressourcenanforderung auch in einem Schritt löschen:

# With --force
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE} --quiet --force

Benchmarkergebnisse

Das Stable Diffusion-Trainingsskript wurde auf v5p-8, v5p-32 und v5p-128 ausgeführt. Die folgende Tabelle zeigt den Durchsatz.

	v5p-8	v5p-32	v5p-128
Trainingsschritt	150	150	150
Globale Batchgröße	32	64	64
Durchsatz (Beispiele/s)	12.10	18.08	19.10

MaxText

In dieser Anleitung erfahren Sie, wie Sie die MaxText-Modell mit einem synthetischen Dataset in Cloud TPU.

MaxText ist ein leistungsstarker, beliebig skalierbarer, Open-Source-LLM, der in reiner Python/JAX geschrieben wurde und auf Cloud TPUs ausgerichtet ist. MaxText unterstützt Forscher und Entwickler mit einem barrierefreien und anpassbares Tool, um die Grenzen der Natural Language API voranzutreiben Forschung und Entwicklung in der Verarbeitung (NLP).

Bevor Sie diese Anleitung ausführen, müssen Sie Richten Sie Ihre Cloud TPU-Umgebung ein.

1. Umgebungsvariablen einrichten

   export PROJECT_ID=your_project_ID
   export TPU_NAME=your_tpu_name # user defined TPU name
   export ACCELERATOR_TYPE=v5p-256
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export RUN_NAME=your_experiment_run_name # user defined name for this run
   export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
   export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
   export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

   Beschreibung der Befehls-Flags

   Variable	Beschreibung
   PROJECT_ID	Name des Google Cloud-Projekts
   TPU_NAME	Ein benutzerdefinierter Name für Ihre TPU.
   ACCELERATOR_TYPE	Informationen zu Ihrer TPU-Version finden Sie auf der Seite TPU-Versionen.
   ZONE	Informationen zu den unterstützten Zonen finden Sie im Dokument TPU-Regionen und -Zonen.
   RUNTIME_VERSION	Verwenden Sie für v5p als Laufzeitversion v2-alpha-tpuv5.
   RUN_NAME	Vom Nutzer angegebener Name der Testausführung.

   Optionale Einrichtung für Multislice empfohlen:

   export NETWORK_NAME=your_network_name
   export FIREWALL_RULE_NAME=your_firewall_rule_name

   Wenn Sie Multislice-Arbeitslasten ausführen und optimale Netzwerkleistung erzielen, sollten Sie ein dediziertes Netzwerk mit einer maximalen Übertragungseinheit (MTU) von 8.896 Byte und Konfigurieren der entsprechenden Firewallregeln. Dieser Schritt ist optional, kann aber die Leistung erheblich verbessern, insbesondere wenn die Anzahl der Slices über das Rechenzentrumsnetzwerk (DCN) skaliert wird. Zum Erstellen eines Netzwerks ist die Berechtigung compute.networks.create im Projekt erforderlich. In den folgenden Beispielen wird gezeigt, wie Sie ein dediziertes Netzwerk und eine Firewallregel erstellen.

   So erstellen Sie ein dediziertes Netzwerk:

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

   Firewallregel erstellen

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

2. MaxText-Repository klonen

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

3. Modell trainieren

   In den folgenden Abschnitten werden zwei Optionen zum Trainieren von MaxText beschrieben.

   Option 1

   Wenn Sie möchten, dass ein Skript den gesamten Workflow verwaltet, von der Bereitstellung Cloud TPUs und Installieren von Abhängigkeiten zur Ausführung Ihres Modells und Ressourcen löschen, können Sie multihost_job.py verwenden.

   cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
   --NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
   --VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
   --BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
   --COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

   Hinweis: Fügen Sie --CQR_EXTRA_ARGS="--network=${NETWORK_NAME}" hinzu. wenn Sie ein benutzerdefiniertes Netzwerk.

   Nach der Initiierung des Skripts sollten Sie eine Meldung wie wie folgt im Protokoll: Auf den Logspeicherort wird verwiesen in der Ausgabenachricht. Klicken Sie auf den ersten Link, Logs aller Worker nach Abschluss der TPU-Bereitstellung

   ------------------------------------
   
   multihost_job finished running, TPUs are starting up to run your job remotely.
   
   Logs for your job are displayed here:
   https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID
   
   To see the output of a single host, you may edit the slice and worker
   number in the `log_file_path` property here:
   
   https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID
   
   When your job is finished, the main command log is in your Cloud Storage
   bucket:
   https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID
   
   View the status of the created TPUs using:
   gcloud compute tpus queued-resources list --filter=RUN_NAME
   --zone=ZONE --project=PROJECT_ID
   

Option 2

Um das Trainingsskript mehrmals auf einem bereitgestellten Cloud TPU, verwenden Sie Das Skript multihost_runner.py, um die Ressource zu verwenden.

1. Richten Sie Variablen zum Erstellen einer TPU ein.

   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export VALID_DURATION=1d
   export QUOTA_TYPE=quota_type

   Hinweis: Die folgenden Flags sind nur erforderlich, wenn Sie Multislice verwenden:

   --node-count ${NODE_COUNT} \
   --node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID

   Hinweis: Wenn Sie --node-count und --node-prefix angeben, geben Sie nicht gleichzeitig an. --node-id, wenn Sie im folgenden Schritt die TPU-Ressource erstellen: eine wird generiert, wenn alle drei Flags angegeben sind.

2. Eine TPU-Ressource erstellen

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --valid-until-duration ${VALID_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

   Hinweis: Das Flag QUOTA_TYPE kann so lauten: reserved oder spot. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zu Ihren TPU-VMs herstellen, sobald Ihre QueuedResource den Status ACTIVE hat:

   Verwenden Sie den Befehl describe, um den Status Ihrer in die Warteschlange gestellten Ressource abzufragen.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
   --project ${PROJECT_ID} --zone ${ZONE}

   Wenn die in die Warteschlange gestellte Ressource den Status AKTIV hat, lautet die Ausgabe: etwa so:

    state: ACTIVE
   

3. Über SSH eine Verbindung zur TPU herstellen

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}

4. Abhängigkeiten installieren

   Hinweis: Nachdem Sie ssh zum Herstellen einer Verbindung zu Ihrer TPU-VM verwendet haben, benötigen Sie um die Variablen ${TPU_NAME} und ${MAXTEXT_OUTPUT_PATH} so einzustellen, dass sie verwendet werden in der TPU-VM:

   export TPU_NAME=your_tpu_name
   export MAXTEXT_OUTPUT_PATH=output-path

   cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
   --COMMAND='bash setup.sh'

5. Führen Sie das Modell mit verschiedenen Konfigurationsscripts aus, z. B. 32b.sh oder 64b.sh. Wenn Sie das Script von einer TPU-VM aus ausführen, müssen Sie das Flag --INTERNAL_IP=true hinzufügen.

   python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
   --COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
   OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

   Hinweis: Legen Sie für die Mehrfachschicht-Scans --TPU_PREFIX=${NODE_PREFIX} anstelle von --TPU_PREFIX=${TPU_NAME} fest.

Bereinigen

Löschen Sie Ihre TPU und die in die Warteschlange gestellten Ressourcen.

Benchmarkergebnisse

Das MaxText-Trainingsskript wurde mit bf16-Genauigkeit von 32 B bis 1.160 B ausgeführt. Die Ergebnisse dieser Ausführungen sind in der folgenden Tabelle aufgeführt.

Anzahl der Parameter	Beschleunigertyp	TFLOP/Chip/Sek.	Modell-Flops-Auslastung (MFU)
32B	v5p-128	3,28E+02	71,47%
64 B	v5p-128	3,23E+02	70,31 %
128 B	v5p-256	3.15E+02	68,68 %
128 Mrd.	v5p-512	3,15E+02	68,53%
256 Mrd.	v5p-1024	3,16E+02	68,82%
512 B	v5p-1024	2,94E+02	63,99 %
1024 Mrd.	v5p-2048	2,49E+02	64,05 %
1024 Mrd.	v5p-4096	2,97E+02	64,80%
1160 Mrd.	v5p-7680	2,95E+02	64,27 %
1160 Mrd.	v5p-12288	3,04E+02	66,23 %

Das 256B-Parametermodell wurde mit v5p-512 und v5p-1024 getestet mit bf16- und int8-Gewichtungen. Die Ergebnisse dieser Tests sind in der folgenden Tabelle aufgeführt.

	v5p-512	v5p-512	v5p-1024	v5p-1024
Globale Batchgröße (Tokens)	5,24E+05	5,24E+05	1,05E+06	1,05E+06
Precision	bf16	int8	bf16	int8
TFLOP/Chip/Sek.	307	408	308	414
Auslastung der Modell-Flops (MFU)	66,98%	88,85 %	67,09%	90,23%

TensorFlow-Anleitungen

ResNet auf einem einzelnen Host trainieren (v5p)

In dieser Anleitung wird beschrieben, wie Sie ImageNet mit einem gefälschten Dataset auf einer v5p-8 TPU trainieren. Wenn Sie ein anderes Dataset verwenden möchten, lesen Sie Dataset vorbereiten

Hinweis: Bevor Sie diese Anleitung ausführen, müssen Sie Ihre Cloud TPU-Umgebung einrichten.

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5p-32
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.17.0-pjrt
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   Verwenden Sie für diese Anleitung v5p-8 als ACCELERATOR_TYPE.

2. TPU-Ressource erstellen:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
     --node-id ${TPU_NAME} \
     --project ${PROJECT_ID} \
     --zone ${ZONE} \
     --accelerator-type ${ACCELERATOR_TYPE} \
     --runtime-version ${RUNTIME_VERSION} \
     --${QUOTA_TYPE}

   Hinweis: Das Flag QUOTA_TYPE kann so lauten: reserved oder spot. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden, finden Sie unter Kontingente.

   Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die in der Warteschlange befindliche Ressource im Status ACTIVE befindet. So prüfen Sie den Status der Ressource in der Warteschlange: verwenden Sie den folgenden Befehl:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
     --project ${PROJECT_ID} \
     --zone ${ZONE}

3. Über SSH eine Verbindung zur TPU herstellen

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}

4. Umgebungsvariablen festlegen

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
   export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

5. Wechseln Sie in das Repository-Verzeichnis für die Modelle und installieren Sie die erforderlichen Pakete.

   cd ${MODELS_REPO} && git checkout r2.15.0
   pip install -r official/requirements.txt

Modell trainieren

1. Führen Sie das Trainingsskript aus.

   python3 official/vision/train.py \
     --tpu=local \
     --experiment=resnet_imagenet \
     --mode=train_and_eval \
     --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
     --model_dir=${MODEL_DIR} \
     --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Bereinigen

Löschen Sie Ihre TPU und die in die Warteschlange gestellten Ressourcen.

ResNet auf einem v5p mit mehreren Hosts trainieren

In dieser Anleitung wird beschrieben, wie Sie ImageNet auf v5p-16 oder größer mit ein fiktives Dataset. Wenn Sie ein anderes Dataset verwenden möchten, lesen Sie den Hilfeartikel Dataset vorbereiten.

Hinweis: Bevor Sie diese Anleitung ausführen können, müssen Sie Ihre Cloud TPU-Umgebung einrichten.

1. Erstellen Sie Umgebungsvariablen:

   export PROJECT_ID=your_project_ID
   export TPU_NAME=your_tpu_name
   export ZONE=us-east1-c
   export ACCELERATOR_TYPE=v5p-16
   export RUNTIME_VERSION=tpu-vm-tf-2.17.0-pod-pjrt
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   ACCELERATOR_TYPE kann entweder v5p-16 oder größer sein.

2. Erstellen Sie eine TPU-Ressource:

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
     --node-id ${TPU_NAME} \
     --project ${PROJECT_ID} \
     --zone ${ZONE} \
     --accelerator-type ${ACCELERATOR_TYPE} \
     --runtime-version ${RUNTIME_VERSION} \
     --${QUOTA_TYPE}

   Hinweis: Das Flag QUOTA_TYPE kann so lauten: reserved oder spot. Wenn keines davon angegeben ist, lautet der Standardwert on-demand. Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden, finden Sie unter Kontingente.

   Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die in der Warteschlange befindliche Ressource im Status ACTIVE befindet.

   Verwenden Sie den Befehl describe, um den Status der in der Warteschlange befindlichen Ressource abzufragen:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
     --project ${PROJECT_ID} \
     --zone ${ZONE}

3. Über SSH eine Verbindung zu Ihrer TPU (Worker 0) herstellen

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}

4. Umgebungsvariablen festlegen

   export TPU_NAME=your_tpu_name
   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export TPU_LOAD_LIBRARY=0

5. Wechseln Sie in das Repository-Verzeichnis für die Modelle und installieren Sie die erforderlichen Pakete.

   cd $MODELS_REPO && git checkout r2.15.0
   pip install -r official/requirements.txt

Modell trainieren

1. Führen Sie das Trainingsskript aus.

   python3 official/vision/train.py \
     --tpu=${TPU_NAME} \
     --experiment=resnet_imagenet \
     --mode=train_and_eval \
     --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
     --model_dir=${MODEL_DIR} \
     --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Bereinigen

Löschen Sie Ihre TPU und die in die Warteschlange gestellten Ressourcen.

PyTorch/XLA

Llama 2

In dieser Anleitung wird beschrieben, wie Sie das Llama 2 7B-Modell mit v5p mit einer Fork des HuggingFace-Repositorys auf PyTorch/XLA mit General and Scalable Parallelization for ML Computation Graphs (GSPMD) trainieren.

Einrichtung

1. Erstellen Sie Variablen für Projekt-ID, Beschleunigertyp, Zone, Laufzeitversion und TPU-Namen.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5p-8
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_DURATION=1d

2. Eine TPU-Ressource erstellen

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --valid-until-duration ${VALID_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

   Hinweis: Das Flag QUOTA_TYPE kann so lauten: reserved oder spot. Wenn keines davon angegeben ist, lautet der Standardwert on-demand. Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden, finden Sie unter Kontingente.

   Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald Ihre QueuedResource den Status ACTIVE hat:

   Verwenden Sie den Befehl describe, um den Status Ihrer in die Warteschlange gestellten Ressource abzufragen.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project ${PROJECT_ID} \
   --zone ${ZONE}

   Wenn sich die erwartete Ressource im Status „AKTIV“ befindet, sieht die Ausgabe in etwa so aus:

    state: ACTIVE
   

   Hinweis: Wenn Sie in einer Umgebung mit mehreren Segmenten arbeiten, gcloud compute tpus queued-resources ssh --node=all ... verwenden statt gcloud compute tpus tpu-vm ssh ... für Ihre Befehle. So wird sichergestellt, dass alle TPU-Chips ordnungsgemäß genutzt werden.

3. Pytorch/XLA und die erforderlichen Abhängigkeiten installieren

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project ${PROJECT_ID} \
   --zone  ${ZONE} \
   --worker=all \
   --command='
   sudo apt-get update
   sudo apt-get install libopenblas-dev -y
   pip3 install numpy
   pip3 install typing-extensions
   pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
   '

4. Laden Sie das HuggingFace-Repository herunter und installieren Sie die erforderlichen Komponenten.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
   git clone -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
   cd transformers
   pip3 install git+file://$PWD
   pip3 install datasets accelerate evaluate scikit-learn'

5. Laden Sie die 7B-Modellkonfiguration herunter.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

6. Modell trainieren

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_BF16=1
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   
   export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
   --xla_tpu_enable_async_collective_fusion_multiple_steps=true
   --xla_tpu_enable_async_collective_fusion=true
   --xla_tpu_overlap_compute_collective_tc=true
   --xla_enable_async_all_gather=true
   --xla_jf_spmd_threshold_for_windowed_einsum_mib=0"
   
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=20000
   export PROFILE_LOGDIR=/tmp/home/
   
   cd transformers
   python examples/pytorch/language-modeling/run_clm.py \
    --tokenizer_name hf-internal-testing/llama-tokenizer \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --per_device_train_batch_size 96 \
    --per_device_eval_batch_size 8 \
    --num_train_epochs 1 \
    --do_train \
    --output_dir /tmp/output \
    --overwrite_output_dir \
    --config_name ~/config.json \
    --save_strategy no \
    --logging_strategy no \
    --remove_unused_columns no \
    --optim adafactor \
    --torch_dtype bfloat16 \
    --dataloader_drop_last yes \
    --block_size 2048 \
    --spmd_2d_sharding 1 \
    --spmd_grad_chkpt
   '

Wenn Sie in einer Umgebung mit mehreren Ebenen arbeiten, müssen Sie das Flag --spmd_dcn_parallelism auf die Anzahl der Ebenen festlegen.

Der SPMD_USER_GUIDE bietet ein detailliertes Nutzerhandbuch, in dem die verschiedenen Umgebungen Variablen und Ein-/Aus-Schaltflächen des HF-Skripts. Beachten Sie, dass die LIBTPU_INIT_ARGS wird in PyTorch/XLA integriert und in zukünftigen Releases standardmäßig aktiviert.

Bereinigen

Löschen Sie Ihre TPU und die in die Warteschlange gestellten Ressourcen.

Benchmark-Ergebnisse

Der Durchsatz für alle drei Llama 2-Modellgrößen ist in der folgenden Tabelle aufgeführt.

	v5p-8	v5p-128	v5p-128
Modellgröße	7 Mrd.	13 B	70 Mrd.
Globale Batchgröße	96	1.024	128
Fragmentierung der Mesh-Form	(4, 1)	(64, 1)	(16, 4)
Modell-Flops-Auslastung (MFU)	56,67 %	55,80%	51,85 %

Support und Feedback

Wir freuen uns über jedes Feedback. Wenn Sie Feedback geben oder Support anfordern möchten, füllen Sie das Cloud TPU-Support- oder Feedbackformular aus.