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. Im Allgemeinen bietet v5p eine bis zu doppelte Leistung von v4 und bietet gleichzeitig doppelt so viele TPUs in einem Pod (6.000 größtes Slice im Vergleich zu 3.000 in v4), was zu einer bis zu vierfachen Leistung auf Pod-Ebene führt. 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 Cloud TPUs noch nicht kennen, lesen Sie den Einstiegsleitfaden für TPUs.
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.
Für jede Cloud TPU-Version sind bestimmte Beschleunigertypen für das Training oder die Inferenz erforderlich. Diese Beschleunigertypen werden in v5p-Konfigurationen beschrieben.
TPU-Ressourcen verwalten
Bei allen Befehlen in diesem Dokument wird davon ausgegangen, dass Sie eine TPU v5p-VM erstellen. Weitere Informationen zu den Befehlen zum Erstellen von TPU-VMs finden Sie unter TPUs verwalten oder im Nutzerhandbuch für anstehende Ressourcen. Damit die Befehle einfacher auszuführen sind, werden in den Codebeispielen in diesem Dokument die folgenden Umgebungsvariablen verwendet:
export PROJECT_ID=your-project export ACCELERATOR_TYPE=v5p-8 export ZONE=us-east5-a export RUNTIME_VERSION=v2-alpha-tpuv5 export TPU_NAME=your-tpu-name
Beschreibungen von Umgebungsvariablen
PROJECT_ID- Das Google Cloud-Projekt, in dem Sie die TPU erstellen.
ACCELERATOR_TYPE- Mit dem Beschleunigertyp geben Sie die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.
ZONE- Die Zone, in der Sie die Cloud TPU erstellen möchten.
RUNTIME_VERSION- Die Version der TPU-Software
TPU_NAME- Der benutzerdefinierte Name der verwendeten TPU.
Framework-Einrichtung
In diesem Abschnitt wird die allgemeine Einrichtung für das Modelltraining mit JAX oder PyTorch mit TPU v5p beschrieben.
Einrichtung für JAX
Wenn die Slice-Form mehr als 4 Chips hat, 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() gibt die Gesamtzahl der Chips im jeweiligen Ausschnitt 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 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 auf allen Multislice-Workern 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 Training mit v5p mit JAX zu beginnen.
Einrichtung für PyTorch
Die PJRT-Laufzeit ist die einzige unterstützte Laufzeit für v5p. In PyTorch 2.1 und höher wird PJRT als Standardlaufzeit für alle TPU-Versionen verwendet. 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 pip install numpy pip install torch torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu '
Verwenden Sie ein Python-Script mit PJRT, um die Installation zu validieren. Das Script zeigt die verfügbaren TPU-Geräte an. Die hier gezeigten Ausgaben wurden mit einem v5p-32-Speicherbereich 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 auf allen Multislice-Workern 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 die vorherigen Cloud TPU-Generationen. Weitere Informationen zur Profilerstellung finden Sie unter Profilerstellung Ihres Modells mit Cloud TPU-Tools und unter Cloud TPU-VMs überwachen.
Trainingsanleitungen
In diesem Abschnitt geht es hauptsächlich um Trainingsanleitungen für eine einzelne Ansicht.
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 in Multislice-Aufnahmen.
JAX-Anleitungen
Train Diffusion 2.1
In dieser Anleitung erfahren Sie, wie Sie das Stable Diffusion-Modell von HuggingFace mit dem Pokémon-Dataset auf Cloud TPU v5p trainieren.
Das Stable Diffusion-Modell ist ein latentes Text-zu-Bild-Modell, das fotorealistische Bilder aus beliebigen Texteingaben generiert. Weitere Informationen finden Sie in den folgenden Ressourcen:
Einrichten
Erstellen Sie Umgebungsvariablen:
export GCS_BUCKET_NAME=your-bucket export PROJECT_ID=your-project-ID export ACCELERATOR_TYPE=v5p-32 export ZONE=europe-west4-b export LOCATION=europe-west4 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 Informationen zu Ihrer TPU-Version finden Sie auf der Seite TPU-Versionen. ZONE Welche Zonen unterstützt werden, erfahren Sie im Dokument TPU-Regionen und ‑Zonen. ORT Die Google Cloud-Region, in der der Cloud Storage-Bucket erstellt werden soll. 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: tpu-service-account@myprojectID.iam.gserviceaccount.com TPU_NAME Die vom Nutzer zugewiesene Text-ID der TPU, die erstellt wird, wenn die anstehende Ressourcenanfrage zugewiesen wird. QUEUED_RESOURCE_ID Die vom Nutzer zugewiesene Text-ID der anstehenden Ressourcenanfrage. Informationen zu Ressourcen in der Warteschlange finden Sie im Dokument Ressourcen in der Warteschlange. QUOTA_TYPE Kann reservedoderspotsein. Wenn keines dieser Elemente angegeben ist, wird für QUOTA_TYPE standardmäßigon-demandverwendet. Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden, finden Sie unter Kontingente.VALID_UNTIL_DURATION Die Dauer, für die die Anfrage gültig ist. Informationen zu den verschiedenen gültigen Zeiträumen finden Sie unter Ressourcen in der Warteschlange. Richten Sie einen Speicher-Bucket für die Modellausgabe ein.
gcloud storage buckets create gs://$GCS_BUCKET_NAME \ --project=$PROJECT_ID \ --location=$LOCATION
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}
Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die in die Warteschlange gestellte Ressource im Status
ACTIVEbefindet. Prüfen Sie den Status der in der Warteschlange befindlichen Ressource mit dem folgenden Befehl:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} --zone ${ZONE}
Wenn sich die erwartete Ressource im Status
ACTIVEbefindet, sieht die Ausgabe in etwa so aus:state: ACTIVE
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. Wenn Sie eine in die Warteschlange gestellte Ressource löschen möchten, löschen Sie in zwei Schritten die Sliver und dann die Anfrage für die in die Warteschlange gestellte Ressource:
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
Benchmark-Ergebnisse
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/Sek.) |
12.10 |
18.08 |
19.10 |
MaxText
In dieser Anleitung erfahren Sie, wie Sie das MaxText-Modell mit einem synthetischen Dataset auf Cloud TPU trainieren.
MaxText ist ein leistungsstarker, beliebig skalierbarer, Open-Source-LLM, der in reiner Python/JAX geschrieben wurde und auf Cloud TPUs ausgerichtet ist. MaxText bietet Forschern und Entwicklern ein zugängliches und anpassbares Tool, mit dem sie die Grenzen der Forschung und Entwicklung im Bereich der Verarbeitung natürlicher Sprache (Natural Language Processing, NLP) erweitern können.
Bevor Sie diese Anleitung ausführen können, müssen Sie Ihre Cloud TPU-Umgebung einrichten.
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 Welche Zonen unterstützt werden, erfahren Sie im Dokument TPU-Regionen und ‑Zonen. RUNTIME_VERSION Verwenden Sie für v5p „v2-alpha-tpuv5“ als Laufzeitversion. RUN_NAME Vom Nutzer angegebener Name der Testausführung. Optionale Einrichtung für Multislice-Scans:
export NETWORK_NAME=your_network_name export FIREWALL_RULE_NAME=your_firewall_rule_name
Wenn Sie Multislice-Arbeitslasten ausführen und eine optimale Netzwerkleistung wünschen, sollten Sie ein eigenes Netzwerk mit einer MTU (Maximum Transmission Unit) von 8.896 Byte erstellen und entsprechende Firewallregeln konfigurieren. 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.createim Projekt erforderlich. In den folgenden Beispielen wird gezeigt, 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}
MaxText-Repository klonen
git clone https://github.com/google/maxtext.git
Modell trainieren
In den folgenden Abschnitten werden zwei Optionen zum Trainieren von MaxText beschrieben.
Option 1
Wenn Sie ein Script zum Verwalten des gesamten Workflows benötigen, von der Bereitstellung von Cloud TPUs und dem Installieren von Abhängigkeiten bis hin zum Ausführen Ihres Modells und zum Auflösen von Ressourcen, können Sie
multihost_job.pyverwenden.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 dem Starten des Scripts sollte im Protokoll eine Meldung ähnlich der folgenden angezeigt werden. Auf den Speicherort des Protokolls wird in der Ausgabenachricht verwiesen. Klicken Sie auf den ersten Link, um auf die Protokolle 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
Wenn Sie das Trainingsskript mehrmals auf einer bereitgestellten Cloud TPU ausführen möchten, verwenden Sie das multihost_runner.py-Script, um die Ressource zu verwenden.
Richten Sie Variablen ein, um eine TPU zu erstellen.
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
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}
Sie können eine SSH-Verbindung zu Ihren TPU-VMs herstellen, sobald Ihre
QueuedResourceden StatusACTIVEhat: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}
Wenn die in der Warteschlange befindliche Ressource den Status „AKTIV“ hat, sieht die Ausgabe in etwa so aus:
state: ACTIVE
Über SSH eine Verbindung zur TPU herstellen
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
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'
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=truehinzufü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
Löschen Sie Ihre TPU und die in die Warteschlange gestellten Ressourcen.
Benchmark-Ergebnisse
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. |
Auslastung der Modell-Flops (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 B |
v5p-512 |
3.15E+02 |
68,53% |
256 B |
v5p-1024 |
3.16E+02 |
68,82% |
512 B |
v5p-1024 |
2,94E+02 |
63,99% |
1.024 B |
v5p-2048 |
2,49E+02 |
64,05% |
1.024 B |
v5p-4096 |
2.97E+02 |
64,80% |
1.160 B |
v5p-7680 |
2,95E+02 |
64,27% |
1.160 B |
v5p-12288 |
3.04E+02 |
66,23% |
Das Modell mit 256 B Parametern wurde auf v5p-512 und v5p-1024 mit bf16- und int8-Gewichten getestet. 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 (v5p) trainieren
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 den Hilfeartikel Dataset vorbereiten.
Einrichten
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.18.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-8alsACCELERATOR_TYPE.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}
Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die in der Warteschlange befindliche Ressource im Status
ACTIVEbefindet. Verwenden Sie den folgenden Befehl, um den Status der in der Warteschlange befindlichen Ressource zu prüfen:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Über SSH eine Verbindung zur TPU herstellen
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Einige 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
Wechseln Sie in das Repository-Verzeichnis für Modelle und installieren Sie die erforderlichen Pakete.
cd ${MODELS_REPO} && git checkout r2.15.0 pip install -r official/requirements.txt
Modell trainieren
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 mit einem fiktiven Dataset auf v5p-16 oder höher trainieren. Wenn Sie ein anderes Dataset verwenden möchten, lesen Sie den Hilfeartikel Dataset vorbereiten.
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.18.0-pod-pjrt export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
ACCELERATOR_TYPEkann entwederv5p-16oder größer sein.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}
Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die in der Warteschlange befindliche Ressource im Status
ACTIVEbefindet.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}
Über SSH eine Verbindung zu Ihrer TPU (Worker 0) herstellen
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Einige 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
Wechseln Sie in das Repository-Verzeichnis für Modelle und installieren Sie die erforderlichen Pakete.
cd $MODELS_REPO && git checkout r2.15.0 pip install -r official/requirements.txt
Modell trainieren
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
Erstellen Sie Umgebungsvariablen.
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
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}
Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald Ihre
QueuedResourceden StatusACTIVEhat: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}
Wenn sich die in der Warteschlange befindliche Ressource im Status „AKTIV“ befindet, sieht die Ausgabe in etwa so aus:
state: ACTIVE
Installieren Sie Pytorch/XLA und die erforderlichen Abhängigkeiten.
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 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html '
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'
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"
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 ist eine ausführlichere Anleitung, in der alle Umgebungsvariablen und Ein-/Aus-Schaltflächen des HF-Scripts erläutert werden. Hinweis: LIBTPU_INIT_ARGS wird in zukünftigen Releases in PyTorch/XLA aufgenommen und 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 |
7B |
13 B |
70 Mrd. |
Globale Batchgröße |
96 |
1.024 |
128 |
Form des Sharding-Mesh |
(4, 1) |
(64, 1) |
(16, 4) |
Auslastung der Modell-Flops (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.