Cloud TPU v5p-Training

Cloud TPU v5p ist die Cloud TPU der fünften Generation von Google Cloud und der Nachfolger von v4 TPU. v5p ist für umfangreiches Training optimiert und eine führende Plattform für die Entwicklung grundlegender LLMs, Diffusionsmodelle und generativer KI. Allgemein bietet v5p eine bis zu doppelt so hohe Leistung wie v4. Gleichzeitig werden doppelt so viele TPUs in einen Pod gepackt (6.000 größtes Slice gegenüber 3.000 in v4), was zu einer bis zu 4-mal höheren Leistung auf Pod-Ebene führt. Es wird auch mit einer höheren Taktfrequenz betrieben (1,75 GHz gegenüber 1,05 GHz), SparseCore für große Einbettungen und dreifache HBM-Kapazität (High Bandwidth Memory) hinzugefügt.

Konzepte von Cloud TPU v5p

Wenn Sie mit Cloud TPUs noch nicht vertraut sind, lesen Sie die Informationen auf der Startseite der TPU-Dokumentation.

Cloud TPU-Konzepte (z. B. Segmente, 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 das Training oder die Inferenz. Diese Beschleunigertypen werden unter V5p-Konfigurationen beschrieben.

TPU-Ressourcen verwalten

Alle Befehle, mit denen Sie Ihre TPU-VMs verwalten können, werden unter TPUs verwalten oder im Nutzerhandbuch für Ressourcen in der Warteschlange zur Verwaltung von Ressourcen in der Warteschlange beschrieben.

Framework-Einrichtung

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

Einrichtung für JAX

Wenn Sie Slice-Formen mit mehr als vier Chips haben, haben Sie mehrere VMs in einem Slice. In diesem Fall müssen Sie das Flag --worker=all verwenden, um die Installation auf allen TPU-VMs mit einem einzigen Befehl 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 prüfen. Die hier gezeigten Ausgaben wurden mit einem v5p-32-Slice erstellt. Mit diesem Code wird geprüft, ob alles korrekt installiert ist. Dazu wird geprü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 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 im jeweiligen Segment an. jax.local_device_count() gibt die Anzahl der Chips an, auf die eine einzelne VM in diesem Slice zugreifen kann.

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

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

PyTorch einrichten

Die PJRT-Laufzeit ist die einzige unterstützte Laufzeit für v5p. PyTorch 2.1 und höher verwendet PJRT als Standardlaufzeit für alle TPU-Versionen. In diesem Abschnitt wird beschrieben, 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
'

Verwenden Sie ein Python-Skript mit PJRT, um die Installation zu validieren und die verfügbaren TPU-Geräte anzuzeigen. Die hier gezeigten Ausgaben wurden mit einem v5p-32-Slice erstellt.

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

Die PyTorch-Anleitungen in diesem Dokument unterstützen Sie beim Einstieg in das V5p-Training mit PyTorch.

Überwachen und Profil erstellen

Cloud TPU v5p unterstützt das Monitoring und die Profilerstellung mit denselben Methoden wie bei vorherigen Generationen von Cloud TPU. Weitere Informationen zur Profilerstellung finden Sie unter Modell mit Cloud TPU-Tools erstellen. Weitere Informationen zum Monitoring finden Sie unter Cloud TPU-VMs überwachen.

Trainingsanleitungen

Dieser Abschnitt konzentriert sich auf Trainingsanleitungen für ein einzelnes Slice. Wenn Sie diese Anleitungen für das Multislice-Training anpassen möchten, fügen Sie den SSH-Befehlen das Flag --node=all hinzu. Weitere Informationen und Best Practices finden Sie in der Einführung zu Multislice.

Hinweis: Bevor Sie eine dieser Anleitungen ausführen, müssen Sie Ihre Cloud TPU-Umgebung einrichten. Nachdem Sie Ihre TPU-Umgebung eingerichtet 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 „Stabile Diffusion“ aus HuggingFace mithilfe des Datasets Pokémon auf Cloud TPU v5p trainieren.

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

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

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

Einrichten

1. Umgebungsvariablen erstellen:

   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	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 v2-alpha-tpuv5 für die RUNTIME_VERSION.
   SERVICE_ACCOUNT	Dies ist die Adresse Ihres Dienstkontos, die Sie in der Google Cloud Console unter „IAM“ > „Dienstkonten“ finden. Beispiel: [e-mail]
   TPU_NAME	Die vom Nutzer zugewiesene Text-ID der TPU, die bei der Zuweisung der Ressourcenanfrage in der Warteschlange erstellt wird.
   QUEUED_RESOURCE_ID	Die vom Nutzer zugewiesene Text-ID der Ressourcenanfrage in der Warteschlange. Informationen zu Ressourcen in der Warteschlange finden Sie im Dokument Ressourcen in der Warteschlange.
   QUOTA_TYPE	Kann reserved oder best-effort sein. Wenn keine dieser Optionen angegeben ist, wird für QUOTA_TYPE standardmäßig on-demand verwendet. Informationen zu den verschiedenen Kontingenttypen, die von Cloud TPU unterstützt werden, finden Sie unter quotas.
   VALID_UNTIL_DURATION	Die Dauer, für die die Anfrage gültig ist. Informationen zur verschiedenen gültigen Dauer finden Sie unter Ressourcen in der Warteschlange.

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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved oder best-effort sein. Wenn Sie nichts angeben, wird der Standardwert on-demand verwendet. Informationen zu den verschiedenen Kontingenttypen, die von Cloud TPU unterstützt werden, finden Sie unter quotas.

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

   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 Ihre TPU- und Ressourcenanfrage in der Warteschlange am Ende der Sitzung oder entfernen Sie Ressourcenanfragen in der Warteschlange, die den Status „FAILED“ haben. Um eine in die Warteschlange gestellte Ressource zu löschen, löschen Sie die Segmente und dann die Anfrage für die in die Warteschlange gestellte Ressource 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 löschen Sie mit --force die Segmente und die Ressourcenanfrage in der Warteschlange in einem einzigen 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 wird beschrieben, wie Sie das MaxText-Modell mit einem synthetischen Dataset in Cloud TPU trainieren.

MaxText ist ein beliebig skalierbares, gut getestetes Open-Source-LLM, das in reinem Python/JAX geschrieben ist und auf Cloud TPUs abzielt. MaxText bietet Forschern und Entwicklern ein zugängliches und anpassungsfähiges Tool, mit dem sie die Grenzen der Forschung und Entwicklung bei Natural Language Processing (NLP) vorantreiben können.

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

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 Arbeitslasten mit mehreren Segmenten ausführen und eine optimale Netzwerkleistung erzielen möchten, sollten Sie ein dediziertes Netzwerk mit einer maximalen Übertragungseinheit (MTU) von 8.896 Byte erstellen und entsprechende Firewallregeln konfigurieren. Dieser Schritt ist zwar optional, kann jedoch die Leistung erheblich verbessern, insbesondere wenn die Anzahl der Segmente über das Rechenzentrumsnetzwerk (DCN) skaliert wird. Hinweis: Zum Erstellen eines Netzwerks ist die Berechtigung compute.networks.create im Projekt erforderlich. Die folgenden Beispiele zeigen, wie Sie ein dediziertes Netzwerk und eine Firewallregel erstellen.

   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 von Cloud TPUs und dem Installieren von Abhängigkeiten bis hin zur Ausführung des Modells und zum Entfernen von Ressourcen, 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 :Wenn Sie ein benutzerdefiniertes Netzwerk verwenden, fügen Sie --CQR_EXTRA_ARGS="--network=${NETWORK_NAME}" hinzu.

   Nach dem Starten des Skripts sollte im Log eine Meldung wie die folgende angezeigt werden. Der Speicherort des Logs wird in der Ausgabenachricht angegeben. Klicken Sie auf den ersten Link, um auf die Logs aller Worker zuzugreifen, sobald die TPU-Bereitstellung abgeschlossen ist.

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

Verwenden Sie zur Nutzung der Ressource das Skript multihost_runner.py, um das Trainingsskript auf einer bereitgestellten Cloud TPU mehrmals auszuführen.

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 beim Erstellen der TPU-Ressource im folgenden Schritt nicht gleichzeitig --node-id an. Wenn alle drei Flags angegeben sind, wird ein Fehler generiert.

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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved oder best-effort sein. Wenn Sie nichts angeben, wird der Standardwert on-demand verwendet. Informationen zu den verschiedenen Kontingenttypen, die von Cloud TPU unterstützt werden, finden Sie unter quotas.

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

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

    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, müssen Sie die Variablen ${TPU_NAME} und ${MAXTEXT_OUTPUT_PATH} festlegen, um sie in der TPU-VM zu verwenden:

   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.sh, 64b.sh aus. 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"
   

   Hinweis: Legen Sie für „Multislice“ --TPU_PREFIX=${NODE_PREFIX} anstelle von --TPU_PREFIX=${TPU_NAME} fest.

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 bf16- und int8-Gewichtungen auf v5p-512 und v5p-1024 getestet. In der folgenden Tabelle sind die Ergebnisse dieser Tests 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
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 mithilfe eines fiktiven Datasets 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-8
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.16.1-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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved oder best-effort sein. Wenn Sie nichts angeben, wird der Standardwert on-demand verwendet. Informationen zu den verschiedenen Kontingenttypen, die von Cloud TPU unterstützt werden, finden Sie unter quotas.

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

   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 mit einem fiktiven Dataset auf v5p-16 oder größer trainieren. Wenn Sie ein anderes Dataset verwenden möchten, finden Sie weitere Informationen unter Dataset vorbereiten.

Hinweis: Bevor Sie diese Anleitung ausführen, 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.16.1-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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved oder best-effort sein. Wenn Sie nichts angeben, wird der Standardwert on-demand verwendet. Informationen zu den verschiedenen Kontingenttypen, die von Cloud TPU unterstützt werden, finden Sie unter quotas.

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

   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 dieser Anleitung wird beschrieben, wie Sie das Llama 2 7B-Modell auf v5p mithilfe einer Abspaltung des HuggingFace-Repositorys auf PyTorch/XLA mit allgemeiner und skalierbarer Parallelisierung für ML-Computerrechner (GSPMD) trainieren.

Einrichtung

1. Erstellen Sie Variablen für die Projekt-ID, den Beschleunigertyp, die Zone, die Laufzeitversion und den 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 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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved oder best-effort sein. Wenn Sie nichts angeben, wird der Standardwert on-demand verwendet. Informationen zu den verschiedenen Kontingenttypen, die von Cloud TPU unterstützt werden, finden Sie unter quotas.

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

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

    state: ACTIVE
   

   Hinweis: Wenn Sie in einer Umgebung mit mehreren Segmenten arbeiten, verwenden Sie gcloud compute tpus queued-resources ssh --node=all ... anstelle von 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 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, müssen Sie das Flag --spmd_dcn_parallelism auf die Anzahl der Segmente festlegen.

SPMD_USER_GUIDE bietet ein ausführlicheres Nutzerhandbuch, in dem alle verschiedenen Umgebungsvariablen und Ein-/Aus-Schaltflächen des HF-Skripts erläutert werden. 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 der folgenden Tabelle enthalten.

	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 Cloud TPU-Support- oder Feedbackformular aus.