Cloud TPU v5p-Training

Cloud TPU v5p ist Cloud TPU von Google Cloud der fünften Generation und für die v4-TPU. v5p ist für umfangreiches Training optimiert und eine führende Plattform für die Entwicklung grundlegender LLMs, Diffusionsmodelle und Generative AI Allgemein gesprochen, v5p bietet eine 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. Er hat auch eine höhere Taktfrequenz (1,75 GHz vs. 1,05 GHz), fügt SparseCore hinzu für große Einbettungen und verdreifacht die Kapazität eines HBM (High Bandwidth Memory).

Konzepte von Cloud TPU v5p

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) sowie die Cloud TPU-Systemarchitektur für alle Cloud TPUs werden in den Cloud TPU-Systemarchitektur Seite.

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

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 Segmentformen mit mehr als 4 Chips haben, haben Sie mehrere VMs. in einer Scheibe. 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 "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Mit dem folgenden Befehl können Sie die Anzahl der Geräte (die die hier gezeigten Ausgaben mit einem v5p-32-Slice erstellt wurden. Dieser Code prüft, ob alles korrekt installiert ist, indem überprüft wird, 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 sieht in etwa so aus:

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 sieht in etwa so aus:

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

Probieren Sie die JAX-Anleitungen in diesem Dokument aus, um 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 mit einem Python-Skript mit PJRT durchführen 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\"))"
'

Probieren Sie die PyTorch-Anleitungen in diesem Dokument aus, um begannen wir mit dem V5p-Training mit PyTorch.

Überwachen und Profil erstellen

Cloud TPU v5p unterstützt Monitoring und Profilerstellung mithilfe der wie bei früheren Generationen von Cloud TPU. Sie können Weitere Informationen finden Sie unter Modell mit Cloud TPU-Tools erstellen. Weitere Informationen zur Profilerstellung und zum Monitoring von Cloud TPU-VMs um mehr über Monitoring zu erfahren.

Trainingsanleitungen

Dieser Abschnitt konzentriert sich auf Trainingsanleitungen für ein einzelnes Slice. Die Anpassung dieser Tutorials an das Multi-Slice-Training kann indem Sie SSH-Befehlen das Flag --node=all hinzufügen. Weitere Informationen und Best Practices finden Sie in der Einführung mit mehreren Segmenten

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

• Stabile Diffusion – Übersicht
• Codequelle für stabile Diffusion

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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=queued_resource_id
   export QUOTA_TYPE=quota_type
   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	Dies ist die Adresse Ihres Dienstkontos, die Sie finden in der Google Cloud Console -> IAM -> Dienstkonten. Beispiel: tpu-service-account@myprojectID.iam.gserviceaccount.com
   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 Ressourcenanfrage in der Warteschlange. Weitere Informationen finden Sie unter Das Dokument Ressourcen in der Warteschlange finden Sie Informationen zu Ressourcen in der Warteschlange.
   QUOTA_TYPE	Kann reserved sein oder best-effort. 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 <ph type="x-smartling-placeholder"></ph> Ressourcen in der Warteschlange finden Sie Informationen zur gültigen Dauer.

2. TPU-Ressource erstellen:

   gcloud alpha 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}
   

   Sie können eine SSH-Verbindung zur TPU-VM herstellen, sobald die Ressource in der Warteschlange hat den Status ACTIVE. 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
   

3. Installieren Sie JAX und die zugehörigen Abhängigkeiten.

   # compatible with v5p: only jax version 0.4.19 and later \
   # jax 0.4.19 requires py 3.10 \
   
   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} --zone=${ZONE} --worker=all \
   --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Laden Sie das HuggingFace-Repository herunter und Installationsvoraussetzungen.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
   

5. Modell trainieren

   Trainieren Sie das Modell mit einem vorab zugeordneten Zwischenspeicher mit 4 GB.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'
   

Bereinigen

Löschen Sie die TPU und die Ressourcenanfrage in der Warteschlange am Ende der Sitzung oder Ressourcenanfragen in der Warteschlange entfernen, die den Status „FAILED“ haben Bundesstaat. Zum Löschen einer Ressource in der Warteschlange löschen Sie die Segmente und anschließend die Ressource in der Warteschlange Anfrage in zwei 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

Oder verwenden Sie --force, um die Segmente und die Ressourcenanfrage in der Warteschlange zu löschen in nur einem Schritt:

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

Benchmarkergebnisse

Das Trainingsskript „Stable Diffusion“ wurde unter 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 leistungsstarkes, beliebig skalierbares Open-Source- getesteten LLM, das in reinem Python/JAX geschrieben ist und auf Cloud TPUs abzielt. 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	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 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 zwar optional, kann die Leistung deutlich verbessern, insbesondere bei der Skalierung der Anzahl der Segmente im Netzwerk des Rechenzentrums (DCN). Notiz wird erstellt Für ein Netzwerk ist die Berechtigung compute.networks.create im Projekt erforderlich. Die folgenden Beispiele zeigen, wie Sie ein dediziertes Netzwerk und eine eigene Firewall erstellen Regel.

   Dediziertes Netzwerk erstellen:

   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"
   

   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
   

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

2. Eine TPU-Ressource erstellen

   gcloud alpha 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}
   

   Sie können über SSH eine Verbindung zu Ihren TPU-VMs herstellen, sobald Sie QueuedResource befindet sich im Status ACTIVE:

   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

   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 Konfigurationsskripts wie 32b.sch, 64b.sch. Wenn Sie das Skript von einer TPU-VM 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"
   

Bereinigen

TPU und Ressourcen in der Warteschlange löschen

Benchmarkergebnisse

Das MaxText-Trainingsskript wurde von 32 B bis 1.160 B mit bf16-Genauigkeit 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)
32 Mrd.	v5p-128	3,28E+02	71,47%
64 Mrd.	v5p-128	3,23E+02	70,31%
128 Mrd.	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 Mrd.	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. In der folgenden Tabelle zeigt die Ergebnisse dieser Tests an.

	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
Modell-Flops-Auslastung (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 auf einer v5p-8-TPU trainieren mithilfe eines fiktiven Datasets. Wenn Sie ein anderes Dataset verwenden möchten, lesen Sie Dataset vorbereiten

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5p-8
   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 alpha 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}
   

   Sie können über SSH eine Verbindung zu Ihrer TPU-VM herstellen, Ressource in Warteschlange befindet sich in ACTIVE. 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 zum Verzeichnis des Modell-Repositorys und installieren Sie die Installationsvoraussetzungen.

   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

TPU und Ressourcen in der Warteschlange löschen

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, finden Sie weitere Informationen unter Dataset vorbereiten.

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. TPU-Ressource erstellen:

   gcloud alpha 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}
   

   Sie können über SSH eine Verbindung zu Ihrer TPU-VM herstellen, Ressource in Warteschlange befindet sich in ACTIVE.

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

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

3. Über SSH eine Verbindung zu Ihrer TPU (Worker Zero) 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 zum Verzeichnis des Modell-Repositorys und installieren Sie die Installationsvoraussetzungen.

   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

TPU und Ressourcen in der Warteschlange löschen

PyTorch/XLA

Llama 2

In diesem Tutorial erfahren Sie, wie Sie das Llama 2 7B-Modell mit einer Gabel auf V5p trainieren des HuggingFace-Repositorys auf PyTorch/XLA mit allgemeinen und skalierbaren Parallelisierung für ML-Berechnungsdiagramme (GSPMD).

Einrichtung

1. Variablen für Projekt-ID, Beschleunigertyp, Zone, Laufzeitversion und TPU-Name.

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

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

   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 Ressource in der Warteschlange im Status AKTIV befindet, wird die Ausgabe sieht in etwa so aus:

    state: ACTIVE
   

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

   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 Segmenten arbeiten, Legen Sie das Flag --spmd_dcn_parallelism auf die Anzahl der Segmente fest.

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

TPU und Ressourcen in der Warteschlange löschen

Benchmarkergebnisse

Der Durchsatz für alle drei Llama 2-Modellgrößen ist in .

	v5p-8	v5p-128	v5p-128
Modellgröße	7 Mrd.	13 Mrd.	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 Support- oder Feedbackformular für Cloud TPU aus.